Fix training

rct_diffusion_pipeline.py

@@ -28,7 +28,7 @@ class RCTDiffusionPipeline(DiffusionPipeline):
         self.unet = UNet2DConditionModel(sample_size=256, in_channels=12, out_channels=12, \
                         down_block_types=('CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'DownBlock2D'),\
                               up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D'), cross_attention_dim=160,
-                            block_out_channels=(12, 24, 30), norm_num_groups=6)
+                            block_out_channels=(64, 128, 256), norm_num_groups=32)
 
         self.unet.to(device='cuda', dtype=torch.float16)
     
@@ -167,7 +167,7 @@ class RCTDiffusionPipeline(DiffusionPipeline):
             colors3.append(c3)
 
         # now put those weights into a tensor
-        class_labels = self.pack_labels_to_tensor(batch_size, object_descriptions, colors1, colors2, colors3)
+        class_labels = self.pack_labels_to_tensor(batch_size, object_descriptions, colors1, colors2, colors3).to(device='cuda',dtype=torch.float16)
 
         # we need those class labels for the 12 channels
         #new_class_labels = torch.Tensor(size=(batch_size, 12, self.get_class_labels_size()))

train_model.py

@@ -108,11 +108,14 @@ def train_model(batch_size=4, epochs=100, save_model_interval=10, start_learning
             clean_images = targets[batch_index:batch_end].to(device='cuda', dtype=torch.float16)
             clean_images = torch.reshape(clean_images, (batch_size, 12, 256, 256))
 
-            noise = torch.randn(clean_images.shape).to('cuda')
-            timesteps = torch.randint(0, model.scheduler.config.num_train_timesteps, (batch_size, )).to('cuda')
+            noise = torch.randn(clean_images.shape).to(device='cuda', dtype=torch.float16)
+            timesteps = torch.randint(0, model.scheduler.config.num_train_timesteps, (batch_size, ))
+            timesteps = timesteps.to(dtype=torch.int, device='cuda')
             noisy_images = model.scheduler.add_noise(clean_images, noise, timesteps).to(device='cuda', dtype=torch.float16)
             noise_pred = model.unet(noisy_images, timesteps, class_labels[batch_index:batch_end].to(device='cuda',dtype=torch.float16), return_dict=False)[0]
-            loss = F.mse_loss(noise_pred, noise)
+
+            noise_pred = noise_pred.to(device='cuda', dtype=torch.float16)
+            loss = F.mse_loss(noise_pred, noise).to(device='cuda', dtype=torch.float16)
             loss.backward()
 
             optimizer.step()