fixed issues

sd/ddpm.py

@@ -14,15 +14,15 @@ def __init__(self, generator: torch.Generator, num_training_steps: int = 1000, b
         self.generator = generator
         self.num_training_steps = num_training_steps
         # we should go from 1000 to 1 during denoising. so we reversed 
-        self.timesteps = torch.from_numpy(ndarray=np.arange(start=0, stop=num_training_steps))[::-1].copy()
+        self.timesteps = torch.from_numpy(np.arange(start=0, stop=num_training_steps)[::-1].copy())
 
     def set_inference_timesteps(self, num_inference_steps:int=50):
         self.num_inference_steps = num_inference_steps
         # 999, 998, 997 .... 0 = 1000 steps
         # 999, 999-20, 999-40, ... 0 = 50 steps
         step_ratio = self.num_training_steps // self.num_inference_steps
-        timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::,-1].copy().astype(np.int64)
-        self.timesteps = torch.from_numpy(ndarray=timesteps)
+        timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
+        self.timesteps = torch.from_numpy(timesteps)
 
     def _get_previous_timestep(self, timestep: int) -> int:
         prev_t = timestep - self.num_training_steps // self.num_inference_steps
@@ -99,7 +99,7 @@ def add_noise(self, original_samples: torch.Tensor, timesteps: torch.IntTensor)
         alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
         timesteps = timesteps.to(original_samples.device)
 
-        sqrt_alpha_prod = alphas_cumprod(timesteps) ** 0.5
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
         sqrt_alpha_prod = sqrt_alpha_prod.flatten()
         while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
             sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)

sd/diffusion.py

@@ -283,6 +283,25 @@ def __init__(self):
             SwitchSequential(UNET_ResidualBlock(in_channels=640, out_channels=320), UNET_AttentionBlock(n_head=8, n_embd=40)),
         ])
 
+    def forward(self, x, context, time):
+        # x: (Batch_Size, 4, Height / 8, Width / 8)
+        # context: (Batch_Size, Seq_Len, Dim) 
+        # time: (1, 1280)
+
+        skip_connections = []
+        for layers in self.encoders:
+            x = layers(x, context, time)
+            skip_connections.append(x)
+
+        x = self.bottleneck(x, context, time)
+
+        for layers in self.decoders:
+            # Since we always concat with the skip connection of the encoder, the number of features increases before being sent to the decoder's layer
+            x = torch.cat((x, skip_connections.pop()), dim=1) 
+            x = layers(x, context, time)
+
+        return x
+
 class UNET_OutputLayer(nn.Module):
     def __init__(self, in_channels, out_channels):
         super().__init__()

sd/model_loader.py

@@ -10,15 +10,19 @@ def preload_models_from_standard_weights(ckpt_path: str, device: str):
     device = 'cpu'
     encoder = VAE_Encoder().to(device)
     encoder.load_state_dict(state_dict['encoder'], strict=True)
+    encoder.half()
 
     decoder = VAE_Decoder().to(device)
     decoder.load_state_dict(state_dict['decoder'], strict=True)
+    decoder.half()
 
     diffusion = Diffusion().to(device)
     diffusion.load_state_dict(state_dict['diffusion'], strict=True)
+    diffusion.half()
 
     clip = CLIP().to(device)
     clip.load_state_dict(state_dict['clip'], strict=True)
+    clip.half()
 
     return {
         'clip': clip,

sd/pipeline.py

@@ -20,7 +20,7 @@ def rescale(x: torch.Tensor, old_range: tuple[int, int], new_range: tuple[int, i
 
 def get_time_embedding(timestep):
     # shape: (160,)
-    freqs = torch.pow(input=10000, exponent=-torch.arange(start=0, end=160, dtype=torch.float32)/160)
+    freqs = torch.pow(10000, -torch.arange(start=0, end=160, dtype=torch.float32)/160)
     # shape(1, 160)
     x = torch.tensor(data=[timestep], dtype=torch.float32)[:, None] * freqs[None]
 
@@ -55,14 +55,14 @@ def generate(prompt: str, uncond_prompt: str, input_image: None | torch.Tensor,
 
         if do_cfg:
             # convert the prompt into tokens
-            cond_tokens = tokenizer.batch_encode_plus([prompt], padding='max_length', max_length=77).input_ids
+            cond_tokens = tokenizer([prompt], padding='max_length', max_length=77).input_ids
             # convert input ids into tensor (batch_size, seq len)
             cond_tokens = torch.tensor(cond_tokens, dtype=torch.long, device=device)
 
             # (batch size, seq len) -> (batch size, seq len, dim)
             cond_context = clip(cond_tokens)
 
-            uncond_tokens = tokenizer.batch_encode_plus([uncond_prompt], padding='max_length', max_length=77).input_ids
+            uncond_tokens = tokenizer([uncond_prompt], padding='max_length', max_length=77).input_ids
             uncond_tokens = torch.tensor(uncond_tokens, dtype=torch.long, device=device)
             # (batch size, seq len) -> (batch size, seq len, dim)
             uncond_context = clip(uncond_tokens)
@@ -72,7 +72,7 @@ def generate(prompt: str, uncond_prompt: str, input_image: None | torch.Tensor,
 
         else:
             # convert it into a list of tokens
-            tokens = tokenizer.batch_encode_plus([prompt], padding='max_length', max_length=77).input_ids
+            tokens = tokenizer([prompt], padding='max_length', max_length=77).input_ids
             tokens = torch.tensor(tokens, dtype=torch.long, device=device)
             # (1, 77, 768)
             context = clip(tokens)
@@ -81,7 +81,7 @@ def generate(prompt: str, uncond_prompt: str, input_image: None | torch.Tensor,
 
         if sampler_name == 'ddpm':
             sampler = DDPMSampler(generator)
-            sampler.set_inference_steps(n_inference_steps)
+            sampler.set_inference_timesteps(n_inference_steps)
         else:
             raise ValueError(f'unknown sampler: {sampler_name}')
         latents_shape = (1, 4, LATENTS_HEIGHT, LATENTS_WIDTH)
@@ -94,7 +94,7 @@ def generate(prompt: str, uncond_prompt: str, input_image: None | torch.Tensor,
             # height, width, channels
             input_image_tensor = np.array(input_image_tensor)
             # convert to tensor
-            input_image_tensor = torch.tensor(input_image_tensor, dtype=torch.float32, device=device)
+            input_image_tensor = torch.tensor(input_image_tensor, dtype=torch.float16, device=device)
             # unet accepts image in the range [-1, 1]
             input_image_tensor = rescale(x=input_image_tensor, old_range=(0, 255), new_range=(-1, 1))
             # height, width, channels -> batch_size, height, width, channels
@@ -103,7 +103,7 @@ def generate(prompt: str, uncond_prompt: str, input_image: None | torch.Tensor,
             input_image_tensor = input_image_tensor.permute(0, 3, 1, 2)
 
             # batch_size, 4, latent height, latent width
-            encoder_noise = torch.randn(size=latents_shape, generator=generator, device=device)
+            encoder_noise = torch.randn(size=latents_shape, generator=generator, device=device, dtype=torch.float16)
             # batch_size, 4, latent_height, latent_width
             latents = encoder(input_image_tensor, encoder_noise)
 
@@ -116,15 +116,15 @@ def generate(prompt: str, uncond_prompt: str, input_image: None | torch.Tensor,
         else:
             # if we are running text-to-image, start with random noise N(0, I)
             # (Batch_Size, 4, Latents_Height, Latents_Width)
-            latents = torch.randn(size=latents_shape, generator=generator, device=device)
+            latents = torch.randn(size=latents_shape, generator=generator, device=device, dtype=torch.float16)
 
         diffusion = models['diffusion']
         diffusion.to(device)
 
         timesteps = tqdm(sampler.timesteps)
         for i, timestep in enumerate(timesteps):
             # (1, 320)
-            time_embedding = get_time_embedding(timestep=timestep).to(device)
+            time_embedding = get_time_embedding(timestep=timestep).to(device, dtype=torch.float16)
 
             # (batch_szie, 4, latent_height, latent_width)
             model_input = latents