tanner/min-dalle-test
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Use all logical cores in Torch mode

Browse Source
This commit is contained in:
Adam Novak 2022-06-28 22:22:54 -04:00
parent 97a55f169c
commit 28c812c832
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
min_dalle/min_dalle_torch.py
View File

@ -1,8 +1,10 @@
import numpy import numpy
import os
from typing import Dict from typing import Dict
from torch import LongTensor, FloatTensor from torch import LongTensor, FloatTensor
import torch import torch
torch.set_grad_enabled(False) torch.set_grad_enabled(False)
torch.set_num_threads(os.cpu_count())
from .models.vqgan_detokenizer import VQGanDetokenizer from .models.vqgan_detokenizer import VQGanDetokenizer
from .models.dalle_bart_encoder_torch import DalleBartEncoderTorch from .models.dalle_bart_encoder_torch import DalleBartEncoderTorch
@ -114,4 +116,3 @@ def detokenize_torch(image_tokens: LongTensor, is_torch: bool) -> numpy.ndarray:
image = detokenizer.forward(image_tokens).to(torch.uint8) image = detokenizer.forward(image_tokens).to(torch.uint8)
del detokenizer, params del detokenizer, params
return image.to('cpu').detach().numpy() return image.to('cpu').detach().numpy()