Dont decode in the training phase

rct_diffusion_pipeline.py

@@ -278,10 +278,11 @@ class RCTDiffusionPipeline(DiffusionPipeline):
         images = torch.Tensor(size=(batch_size, 3, self.sample_size, self.sample_size)).to('cuda')
         images = noise_batches[:, :3]
 
-        #with torch.no_grad():
-            #image = noise_batches
-            #result = self.vae.decode(image).sample
-            #images = result
+        with torch.no_grad():
+            image = noise_batches
+            result = self.vae.decode(image).sample
+            images = result
+            images = images / self.vae.config.scaling_factor
 
         # convert those tensors to PIL images
         tensor_to_pil = T.ToPILImage()

train_model.py

@@ -124,6 +124,9 @@ def train_model(batch_size=4, total_images=-1, epochs=100, scheduler_num_timeste
     vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", use_safetensors=True)
     vae = vae.to(dtype=torch.float32, device='cuda')
 
+    vae.requires_grad_(False)
+    text_encoder.requires_grad_(False)
+
     optimizer = torch.optim.AdamW(unet.parameters(), lr=start_learning_rate)
     lr_scheduler = get_cosine_schedule_with_warmup(
         optimizer=optimizer,
@@ -149,26 +152,24 @@ def train_model(batch_size=4, total_images=-1, epochs=100, scheduler_num_timeste
             embeddings = create_embeddings(batch, model)
             clean_images = torch.reshape(clean_images, (batch['image'].size(0), SAMPLE_NUM_CHANNELS, SAMPLE_SIZE, SAMPLE_SIZE)).\
                 to(device='cuda')
+            
+            # use the vae to get the latent images
+            latent_images = vae.encode(clean_images).latent_dist.sample()
+            latent_images = latent_images * vae.config.scaling_factor
 
-            noise = torch.randn(clean_images.shape, dtype=torch.float32, device='cuda')
+            noise = torch.randn_like(latent_images)
             timesteps = torch.randint(0, scheduler.config.num_train_timesteps, (batch_size, )).to(device='cuda')
 
             #timesteps = timesteps.to(dtype=torch.int, device='cuda')
-            noisy_images = scheduler.add_noise(clean_images, noise, timesteps)
+            noisy_images = scheduler.add_noise(latent_images, noise, timesteps)
             
             batch_embeddings = embeddings
             batch_embeddings = batch_embeddings.to('cuda')
 
-            # use the vae to get the latent images
-            latent_images = vae.encode(noisy_images).latent_dist.sample()
-
             optimizer.zero_grad()
-            unet_results = unet(latent_images, timesteps, batch_embeddings).sample
-
-            # get back the upscale result
-            noise_pred = vae.decode(unet_results).sample
+            unet_results = unet(noisy_images, timesteps, batch_embeddings).sample
 
-            loss = loss_fn(noise_pred, noise)
+            loss = loss_fn(unet_results, noise)
             loss.backward()
             optimizer.step()
             lr_scheduler.step()