fix dit inference bug. Add wanbd to inference code.

Signed-off-by: sajadn <snorouzi@nvidia.com>

Author: sajadn

dfm/src/megatron/model/dit/edm/edm_pipeline.py

@@ -179,21 +179,17 @@ def training_step(
         # import pdb; pdb.set_trace()
         # Get the input data to noise and denoise~(image, video) and the corresponding conditioner.
         self.net = model
-        x0_from_data_batch, x0, condition = self.get_data_and_condition(data_batch)
+        x0, condition = self.get_data_and_condition(data_batch)
 
         # Sample pertubation noise levels and N(0, 1) noises
         sigma, epsilon = self.draw_training_sigma_and_epsilon(x0.size(), condition)
 
         if parallel_state.is_pipeline_last_stage():
-            output_batch, pred_mse, edm_loss = self.compute_loss_with_epsilon_and_sigma(
-                data_batch, x0_from_data_batch, x0, condition, epsilon, sigma
-            )
+            output_batch, pred_mse, edm_loss = self.compute_loss_with_epsilon_and_sigma(x0, condition, epsilon, sigma)
 
             return output_batch, edm_loss
         else:
-            net_output = self.compute_loss_with_epsilon_and_sigma(
-                data_batch, x0_from_data_batch, x0, condition, epsilon, sigma
-            )
+            net_output = self.compute_loss_with_epsilon_and_sigma(x0, condition, epsilon, sigma)
             return net_output
 
     def denoise(self, xt: torch.Tensor, sigma: torch.Tensor, condition: dict[str, torch.Tensor]):
@@ -232,8 +228,6 @@ def denoise(self, xt: torch.Tensor, sigma: torch.Tensor, condition: dict[str, to
 
     def compute_loss_with_epsilon_and_sigma(
         self,
-        data_batch: dict[str, torch.Tensor],
-        x0_from_data_batch: torch.Tensor,
         x0: torch.Tensor,
         condition: dict[str, torch.Tensor],
         epsilon: torch.Tensor,
@@ -294,14 +288,14 @@ def get_per_sigma_loss_weights(self, sigma: torch.Tensor):
 
     def get_condition_uncondition(self, data_batch: Dict):
         """Returns conditioning and unconditioning for classifier-free guidance."""
-        _, _, condition = self.get_data_and_condition(data_batch, dropout_rate=0.0)
+        _, condition = self.get_data_and_condition(data_batch, dropout_rate=0.0)
 
         if "neg_context_embeddings" in data_batch:
             data_batch["context_embeddings"] = data_batch["neg_context_embeddings"]
             data_batch["context_mask"] = data_batch["context_mask"]
-            _, _, uncondition = self.get_data_and_condition(data_batch, dropout_rate=1.0)
+            _, uncondition = self.get_data_and_condition(data_batch, dropout_rate=1.0)
         else:
-            _, _, uncondition = self.get_data_and_condition(data_batch, dropout_rate=1.0)
+            _, uncondition = self.get_data_and_condition(data_batch, dropout_rate=1.0)
 
         return condition, uncondition
 
@@ -419,13 +413,14 @@ def get_data_and_condition(self, data_batch: dict[str, Tensor], dropout_rate=0.2
             Raw data, latent data, and conditioning information.
         """
         # Latent state
-        raw_state = data_batch["video"] * self.sigma_data
-        # assume data is already encoded
-        latent_state = raw_state
-
-        # Condition
-        data_batch["crossattn_emb"] = self.random_dropout_input(
+        latent_state = data_batch["video"] * self.sigma_data
+        condition = {}  # Create a new dictionary for condition
+        # Copy all keys from data_batch except 'video'
+        for key, value in data_batch.items():
+            if key not in ["video", "context_embeddings"]:
+                condition[key] = value
+
+        condition["crossattn_emb"] = self.random_dropout_input(
             data_batch["context_embeddings"], dropout_rate=dropout_rate
         )
-
-        return raw_state, latent_state, data_batch
+        return latent_state, condition

examples/megatron/recipes/dit/inference_dit_model.py

@@ -24,6 +24,7 @@
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from transformers import T5EncoderModel, T5TokenizerFast
 
+import wandb
 from dfm.src.common.tokenizers.cosmos.cosmos1.causal_video_tokenizer import CausalVideoTokenizer
 from dfm.src.common.utils.save_video import save_video
 from dfm.src.megatron.model.dit.edm.edm_pipeline import EDMPipeline
@@ -143,12 +144,7 @@ def get_pos_id_3d(self, *, t, h, w):
         return self.grid[:t, :h, :w]
 
 
-def prepare_data_batch(args, t5_embeding_max_length=512):
-    tokenizer = T5TokenizerFast.from_pretrained("google-t5/t5-11b", cache_dir=args.t5_cache_dir)
-    text_encoder = T5EncoderModel.from_pretrained("google-t5/t5-11b", cache_dir=args.t5_cache_dir)
-    text_encoder.to("cuda")
-    text_encoder.eval()
-
+def prepare_data_batch(args, tokenizer, text_encoder, t5_embeding_max_length=512):
     print("[args.prompt]: ", args.prompt)
     # Encode text to T5 embedding
     out = encode_for_batch(tokenizer, text_encoder, [args.prompt])
@@ -254,8 +250,7 @@ def load_model_from_checkpoint(args):
     if isinstance(model, list):
         model = model[0]
 
-    model = model.cuda().to(torch.bfloat16)
-    model.eval()
+    model = model.cuda().to(torch.bfloat16).eval()
 
     print_rank_0(f"✅ Model loaded successfully from {checkpoint_path}")
 
@@ -311,12 +306,25 @@ def set_seed(seed):
 
     set_seed(42)
 
+    rank = torch.distributed.get_rank()
+    if rank == 0:
+        gather_list = [None for _ in range(ps.get_data_parallel_world_size())]
+        wandb.init(project="dit-inference-video", name="inference_generation")
+    else:
+        gather_list = None
+
     # Load model from checkpoint or initialize from scratch
     print_rank_0("Loading model from checkpoint...")
     model, diffusion_pipeline, model_config = load_model_from_checkpoint(args)
 
+    tokenizer = T5TokenizerFast.from_pretrained("google-t5/t5-11b", cache_dir=args.t5_cache_dir, dtype=torch.bfloat16)
+    text_encoder = T5EncoderModel.from_pretrained(
+        "google-t5/t5-11b", cache_dir=args.t5_cache_dir, dtype=torch.bfloat16
+    )
+    text_encoder.to("cuda").eval()
+
     print_rank_0("preparing data batch...")
-    data_batch, state_shape = prepare_data_batch(args)
+    data_batch, state_shape = prepare_data_batch(args, tokenizer, text_encoder)
     vae = CausalVideoTokenizer.from_pretrained(args.tokenizer_model, cache_dir=args.tokenizer_cache_dir)
     vae.to("cuda").eval()
 
@@ -356,6 +364,22 @@ def set_seed(seed):
     )
     print_rank_0(f"saved video to rank={rank}_{args.video_save_path}")
 
+    torch.distributed.gather_object(
+        obj=(decoded_video[0], args.prompt),
+        object_gather_list=gather_list,
+        dst=0,
+        group=ps.get_data_parallel_group(),
+    )
+
+    if rank == 0 and wandb.run is not None:
+        videos = []
+        for video, caption in gather_list:
+            video_data_transposed = video.transpose(0, 3, 1, 2)
+            videos.append(wandb.Video(video_data_transposed, fps=args.fps, format="mp4", caption=caption))
+        wandb.log({"generated_videos": videos})
+        wandb.finish()
+        print_rank_0("✅ All videos gathered and logged to wandb with captions")
+
 
 if __name__ == "__main__":
     args = parse_args()