From 736904ef2f6fbb9ae89e78cf4bf08cfd3750fa72 Mon Sep 17 00:00:00 2001 From: Brett Kuprel Date: Thu, 7 Jul 2022 17:18:30 -0400 Subject: [PATCH] fix typing --- min_dalle/min_dalle.py | 10 +++++----- min_dalle/models/dalle_bart_decoder.py | 2 +- min_dalle/models/vqgan_detokenizer.py | 16 ++++++++-------- 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/min_dalle/min_dalle.py b/min_dalle/min_dalle.py index 802adcb..d8ba51c 100644 --- a/min_dalle/min_dalle.py +++ b/min_dalle/min_dalle.py @@ -2,7 +2,7 @@ import os from PIL import Image from matplotlib.pyplot import grid import numpy -from torch import LongTensor +from torch import LongTensor, FloatTensor from math import sqrt import torch import json @@ -148,7 +148,7 @@ class MinDalle: self, image_tokens: LongTensor, is_verbose: bool = False - ) -> LongTensor: + ) -> FloatTensor: if not self.is_reusable: del self.decoder if torch.cuda.is_available(): torch.cuda.empty_cache() if not self.is_reusable: self.init_detokenizer() @@ -158,7 +158,7 @@ class MinDalle: return images - def grid_from_images(self, images: LongTensor) -> Image.Image: + def grid_from_images(self, images: FloatTensor) -> Image.Image: grid_size = int(sqrt(images.shape[0])) images = images.reshape([grid_size] * 2 + list(images.shape[1:])) image = images.flatten(1, 2).transpose(0, 1).flatten(1, 2) @@ -175,7 +175,7 @@ class MinDalle: log2_k: int = 6, log2_supercondition_factor: int = 3, is_verbose: bool = False - ) -> Iterator[LongTensor]: + ) -> Iterator[FloatTensor]: assert(log2_mid_count in range(5)) if is_verbose: print("tokenizing text") tokens = self.tokenizer.tokenize(text, is_verbose=is_verbose) @@ -260,7 +260,7 @@ class MinDalle: log2_k: int = 6, log2_supercondition_factor: int = 3, is_verbose: bool = False - ) -> LongTensor: + ) -> FloatTensor: log2_mid_count = 0 images_stream = self.generate_images_stream( text, diff --git a/min_dalle/models/dalle_bart_decoder.py b/min_dalle/models/dalle_bart_decoder.py index d7fad35..5d34bac 100644 --- a/min_dalle/models/dalle_bart_decoder.py +++ b/min_dalle/models/dalle_bart_decoder.py @@ -1,6 +1,6 @@ from typing import Tuple, List import torch -from torch import LongTensor, nn, FloatTensor, BoolTensor +from torch import nn, LongTensor, FloatTensor, BoolTensor torch.set_grad_enabled(False) from .dalle_bart_encoder import GLU, AttentionBase diff --git a/min_dalle/models/vqgan_detokenizer.py b/min_dalle/models/vqgan_detokenizer.py index 3590f79..19881b5 100644 --- a/min_dalle/models/vqgan_detokenizer.py +++ b/min_dalle/models/vqgan_detokenizer.py @@ -1,5 +1,5 @@ import torch -from torch import Tensor +from torch import FloatTensor from torch.nn import Module, ModuleList, GroupNorm, Conv2d, Embedding torch.set_grad_enabled(False) @@ -16,7 +16,7 @@ class ResnetBlock(Module): if not self.is_middle: self.nin_shortcut = Conv2d(m, n, 1) - def forward(self, x: Tensor) -> Tensor: + def forward(self, x: FloatTensor) -> FloatTensor: h = x h = self.norm1.forward(h) h *= torch.sigmoid(h) @@ -39,7 +39,7 @@ class AttentionBlock(Module): self.v = Conv2d(n, n, 1) self.proj_out = Conv2d(n, n, 1) - def forward(self, x: Tensor) -> Tensor: + def forward(self, x: FloatTensor) -> FloatTensor: n, m = 2 ** 9, x.shape[0] h = x h = self.norm(h) @@ -67,7 +67,7 @@ class MiddleLayer(Module): self.attn_1 = AttentionBlock() self.block_2 = ResnetBlock(9, 9) - def forward(self, h: Tensor) -> Tensor: + def forward(self, h: FloatTensor) -> FloatTensor: h = self.block_1.forward(h) h = self.attn_1.forward(h) h = self.block_2.forward(h) @@ -81,7 +81,7 @@ class Upsample(Module): self.upsample = torch.nn.UpsamplingNearest2d(scale_factor=2) self.conv = Conv2d(n, n, 3, padding=1) - def forward(self, x: Tensor) -> Tensor: + def forward(self, x: FloatTensor) -> FloatTensor: x = self.upsample.forward(x.to(torch.float32)) x = self.conv.forward(x) return x @@ -116,7 +116,7 @@ class UpsampleBlock(Module): self.upsample = Upsample(log2_count_out) - def forward(self, h: Tensor) -> Tensor: + def forward(self, h: FloatTensor) -> FloatTensor: for j in range(3): h = self.block[j].forward(h) if self.has_attention: @@ -144,7 +144,7 @@ class Decoder(Module): self.norm_out = GroupNorm(2 ** 5, 2 ** 7) self.conv_out = Conv2d(2 ** 7, 3, 3, padding=1) - def forward(self, z: Tensor) -> Tensor: + def forward(self, z: FloatTensor) -> FloatTensor: z = self.conv_in.forward(z) z = self.mid.forward(z) @@ -165,7 +165,7 @@ class VQGanDetokenizer(Module): self.post_quant_conv = Conv2d(n, n, 1) self.decoder = Decoder() - def forward(self, z: Tensor) -> Tensor: + def forward(self, z: FloatTensor) -> FloatTensor: z = self.embedding.forward(z) z = z.view((z.shape[0], 2 ** 4, 2 ** 4, 2 ** 8)) z = z.permute(0, 3, 1, 2).contiguous()