From 3c0ebd7d70eaedcee519e7fc83cfddb14741c0b9 Mon Sep 17 00:00:00 2001
From: mayukhdeb <mayukhmainak2000@gmail.com>
Date: Fri, 28 Jun 2024 04:38:43 -0700
Subject: [PATCH] apply mask to loss + some hardcoding for banny debugging

---
 main_sd3.py | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/main_sd3.py b/main_sd3.py
index e461a81..f6a19f4 100644
--- a/main_sd3.py
+++ b/main_sd3.py
@@ -26,6 +26,7 @@ todos:
     - [x] [later] clip grad norms after loss backward
     - [x] wandb log loss
     - [x] loss go down
+    - [x] apply mask to denoising loss
 [x] - Save checkpoint after training
 [] - save checkpoint during training
 [x] - inference + visualize examples
@@ -538,9 +539,22 @@ def main(config: TrainingConfig, wandb_log = False):
             else:
                 captions, vae_latent, mask = train_dataset.get_aspect_ratio_bucketed_batch()
 
-            print(f"Step: {global_step}\n TI lr: {optimizer_ti.param_groups[0]['lr']} Example prompt: {captions[0]}")
             model_input = vae_latent
             prompts = captions
+            """
+            some hardcoding on the captions just to see whether it works
+            """
+            prompts = [
+                x.replace("<s0><s1>, ", "") for x in prompts
+            ]
+            prompts = [
+                x.replace("bananaman", "<s0><s1>") for x in prompts
+            ]
+            """
+            done with hardcoding
+            """
+            print(f"Step: {global_step}\n Example prompt: {prompts[0]}")
+
             prompt_embeds, pooled_prompt_embeds = compute_text_embeddings(
                 prompt = prompts, 
                 text_encoders = text_encoders, 
@@ -608,8 +622,18 @@ def main(config: TrainingConfig, wandb_log = False):
             target = model_input
 
             # Compute regular loss.
+            loss_term = (
+                    weighting.float().to(model_pred.device) * (model_pred.float() - target.float().to(model_pred.device)) ** 2
+                )
+            """
+            apply mask
+            """
+            assert loss_term.shape == mask.shape
+
+            loss_term = loss_term * mask.to(loss_term.device)
+
             loss = torch.mean(
-                (weighting.float().to(model_pred.device) * (model_pred.float() - target.float().to(model_pred.device)) ** 2).reshape(target.shape[0], -1),
+                loss_term.reshape(target.shape[0], -1),
                 1,
             )
             loss = loss.mean()