Push unet and scheduler outside of pipeline

rct_diffusion_pipeline.py

@@ -11,6 +11,12 @@ import pandas as pd
 from tqdm.auto import tqdm
 
 class RCTDiffusionPipeline(DiffusionPipeline):
+    def get_default_unet(hidden_dim):
+        return 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=hidden_dim,
+                            block_out_channels=(64, 128, 256), norm_num_groups=32)
+
     def __init__(self):
         super().__init__()
 
@@ -21,16 +27,14 @@ class RCTDiffusionPipeline(DiffusionPipeline):
         self.color3_dict = {}
         self.load_dictionaries_from_dataset()
 
-        self.scheduler = DDPMScheduler()
-
-        # the number of hidden features is dependant on the loaded dictionaries!
-        hidden_dim = self.get_class_labels_size()
-        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=(64, 128, 256), norm_num_groups=32)
-
-        self.unet.to(dtype=torch.float16)
+        self.scheduler = None
+        self.unet = None
+    
+    def set_unet(self, unet):
+        self.unet = unet
+    
+    def set_scheduler(self, scheduler):
+        self.scheduler = scheduler
     
     def load_dictionaries_from_dataset(self):
         dataset = load_dataset('frutiemax/rct_dataset')

train_model.py

@@ -10,6 +10,7 @@ import torch.nn.functional as F
 from diffusers.optimization import get_cosine_schedule_with_warmup
 from tqdm.auto import tqdm
 from accelerate import Accelerator
+from diffusers import DDPMScheduler, UNet2DConditionModel
 
 def save_and_test(pipeline, epoch):
     outputs = pipeline([[('aleppo pine tree', 1.0)]], [[('dark green', 1.0)]])
@@ -92,7 +93,8 @@ def train_model(batch_size=4, epochs=100, scheduler_num_timesteps=20, save_model
     del colors3
     del dataset
 
-    optimizer = torch.optim.Adam(model.unet.parameters(), lr=start_learning_rate)
+    unet = RCTDiffusionPipeline.get_default_unet(160)
+    optimizer = torch.optim.Adam(unet.parameters(), lr=start_learning_rate)
     lr_scheduler = get_cosine_schedule_with_warmup(
         optimizer=optimizer,
         num_warmup_steps=lr_warmup_steps,
@@ -108,10 +110,11 @@ def train_model(batch_size=4, epochs=100, scheduler_num_timesteps=20, save_model
         project_dir='logs',
     )
 
-    unet, scheduler, optimizer, lr_scheduler = accelerator.prepare(model.unet, model.scheduler, \
+    scheduler = DDPMScheduler(scheduler_num_timesteps)
+    unet, scheduler, optimizer, lr_scheduler = accelerator.prepare(unet, scheduler, \
       optimizer, lr_scheduler)
-
-    del model
+    
+    unet = unet.to(dtype=torch.float16)
     scheduler.set_timesteps(scheduler_num_timesteps)
     
     for epoch in range(epochs):
@@ -128,7 +131,7 @@ def train_model(batch_size=4, epochs=100, scheduler_num_timesteps=20, save_model
             noisy_images = scheduler.add_noise(clean_images, noise, timesteps).to(device='cuda', dtype=torch.float16)
 
             with accelerator.accumulate(unet):
-                noise_pred = unet(noisy_images, timesteps.to(device='cuda', dtype=torch.float16), class_labels[batch_index:batch_end].to(device='cuda',dtype=torch.float16), return_dict=False)[0]
+                noise_pred = unet(noisy_images, timesteps.to(device='cuda'), class_labels[batch_index:batch_end].to(device='cuda',dtype=torch.float16), return_dict=False)[0]
 
                 #noise_pred = noise_pred.to(device='cuda', dtype=torch.float16)
                 loss = F.mse_loss(noise_pred, noise.to('cuda', dtype=torch.float16))
@@ -143,6 +146,8 @@ def train_model(batch_size=4, epochs=100, scheduler_num_timesteps=20, save_model
             model.unet = accelerator.unwrap_model(unet)
             model.scheduler = scheduler
             save_and_test(model, epoch)
+            del model.unet
+            del model.scheduler
         progress_bar.update(1)