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

use cuda if available

Browse Source
This commit is contained in:
Brett Kuprel 2022-06-28 12:47:11 -04:00
parent 8544f59576
commit 5aa6fe49bf
3 changed files with 13 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
min_dalle/load_params.py
View File

@ -2,8 +2,9 @@ import os
import numpy import numpy
from copy import deepcopy from copy import deepcopy
from typing import Dict from typing import Dict
import torch
from flax import traverse_util, serialization from flax import traverse_util, serialization
import torch
torch.no_grad()
def load_vqgan_torch_params(path: str) -> Dict[str, torch.Tensor]: def load_vqgan_torch_params(path: str) -> Dict[str, torch.Tensor]:
@ -29,6 +30,7 @@ def load_vqgan_torch_params(path: str) -> Dict[str, torch.Tensor]:
for i in P: for i in P:
P[i] = torch.tensor(P[i]) P[i] = torch.tensor(P[i])
if torch.cuda.is_available(): P[i] = P[i].cuda()
P['embedding.weight'] = P.pop('quantize.embedding.embedding') P['embedding.weight'] = P.pop('quantize.embedding.embedding')
@ -85,6 +87,7 @@ def convert_dalle_bart_torch_from_flax_params(
for i in P: for i in P:
P[i] = torch.tensor(P[i]) P[i] = torch.tensor(P[i])
if torch.cuda.is_available(): P[i] = P[i].cuda()
for i in list(P): for i in list(P):
if 'kernel' in i: if 'kernel' in i:

11
min_dalle/min_dalle_torch.py
View File

@ -73,7 +73,6 @@ def decode_torch(
print("sampling image tokens") print("sampling image tokens")
torch.manual_seed(seed) torch.manual_seed(seed)
text_tokens = torch.tensor(text_tokens).to(torch.long)
image_tokens = decoder.forward(text_tokens, encoder_state) image_tokens = decoder.forward(text_tokens, encoder_state)
return image_tokens return image_tokens
@ -84,10 +83,9 @@ def generate_image_tokens_torch(
config: dict, config: dict,
params: dict, params: dict,
image_token_count: int image_token_count: int
) -> numpy.ndarray: ) -> LongTensor:
text_tokens = torch.tensor(text_tokens).to(torch.long) text_tokens = torch.tensor(text_tokens).to(torch.long)
if torch.cuda.is_available(): if torch.cuda.is_available(): text_tokens = text_tokens.cuda()
text_tokens = text_tokens.cuda()
encoder_state = encode_torch( encoder_state = encode_torch(
text_tokens, text_tokens,
config, config,
@ -101,16 +99,15 @@ def generate_image_tokens_torch(
params, params,
image_token_count image_token_count
) )
return image_tokens.detach().numpy() return image_tokens
def detokenize_torch(image_tokens: numpy.ndarray) -> numpy.ndarray: def detokenize_torch(image_tokens: LongTensor) -> numpy.ndarray:
print("detokenizing image") print("detokenizing image")
model_path = './pretrained/vqgan' model_path = './pretrained/vqgan'
params = load_vqgan_torch_params(model_path) params = load_vqgan_torch_params(model_path)
detokenizer = VQGanDetokenizer() detokenizer = VQGanDetokenizer()
detokenizer.load_state_dict(params) detokenizer.load_state_dict(params)
image_tokens = torch.tensor(image_tokens).to(torch.long)
image = detokenizer.forward(image_tokens).to(torch.uint8) image = detokenizer.forward(image_tokens).to(torch.uint8)
return image.detach().numpy() return image.detach().numpy()

5
min_dalle/models/dalle_bart_decoder_torch.py
View File

@ -127,6 +127,10 @@ class DalleBartDecoderTorch(nn.Module):
self.start_token = torch.tensor([start_token]).to(torch.long) self.start_token = torch.tensor([start_token]).to(torch.long)
self.pad_token = torch.tensor([1]).to(torch.long) self.pad_token = torch.tensor([1]).to(torch.long)
self.condition_factor = torch.tensor([10]).to(torch.float) self.condition_factor = torch.tensor([10]).to(torch.float)
if torch.cuda.is_available():
self.start_token = self.start_token.cuda()
self.pad_token = self.pad_token.cuda()
self.condition_factor = self.condition_factor.cuda()
self.image_token_count = image_token_count self.image_token_count = image_token_count
self.embed_tokens = nn.Embedding(image_vocab_size + 1, embed_count) self.embed_tokens = nn.Embedding(image_vocab_size + 1, embed_count)
self.embed_positions = nn.Embedding(image_token_count, embed_count) self.embed_positions = nn.Embedding(image_token_count, embed_count)
@ -200,6 +204,7 @@ class DalleBartDecoderTorch(nn.Module):
for i in range(self.sample_token_count): for i in range(self.sample_token_count):
token_index = torch.tensor([i]).to(torch.long) token_index = torch.tensor([i]).to(torch.long)
if torch.cuda.is_available(): token_index = token_index.cuda()
probs, keys_values_state = self.decode_step( probs, keys_values_state = self.decode_step(
text_tokens = text_tokens, text_tokens = text_tokens,
encoder_state = encoder_state, encoder_state = encoder_state,