use cuda if available
This commit is contained in:
parent
8544f59576
commit
5aa6fe49bf
|
@ -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:
|
||||||
|
|
|
@ -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()
|
||||||
|
|
|
@ -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,
|
||||||
|
|
Loading…
Reference in New Issue
Block a user