Merge pull request #141 from aryamancodes/aryamang/fix-mv-inference

fix Text2world-MultiView inference by adding missing config options

Author: lynetcha-nv

README.md

@@ -3,7 +3,7 @@
 </p>
 
 <h1 align="center">
-    
+
 > 🚨 **Update Notice**  
 >
 > The latest version of our Cosmos-Predict is now live!

cosmos_predict1/diffusion/config/inference/cosmos-1-diffusion-text2world-multiview.py

@@ -37,6 +37,10 @@
                 88,
                 160,
             ],
+            net=dict(
+                extra_per_block_abs_pos_emb=True,
+                extra_per_block_abs_pos_emb_type="sincos",
+            ),
             tokenizer=dict(
                 video_vae=dict(
                     pixel_chunk_duration=57,
@@ -59,10 +63,12 @@
             name="Cosmos_Predict1_Text2World_7B_Multiview_post_trained",
         ),
         model=dict(
-             net=dict(
+            net=dict(
                 n_views=5,
                 view_condition_dim=3,
-                add_repeat_frame_embedding=False, 
+                add_repeat_frame_embedding=False,
+                extra_per_block_abs_pos_emb=True,
+                extra_per_block_abs_pos_emb_type="sincos",
             ),
             latent_shape=[
                 16,

cosmos_predict1/diffusion/config/inference/cosmos-1-diffusion-video2world-multiview.py

@@ -84,4 +84,3 @@
     Cosmos_Predict1_Video2World_7B_Multiview_post_trained,
 ]:
     cs.store(group="experiment", package="_global_", name=_item["job"]["name"], node=_item)
-

cosmos_predict1/diffusion/inference/world_generation_pipeline.py

@@ -1012,9 +1012,15 @@ def _run_tokenizer_decoding(self, sample: torch.Tensor) -> np.ndarray:
         video = (1.0 + self.model.decode(sample)).clamp(0, 2) / 2  # [B, 3, T, H, W]
         video_segments = einops.rearrange(video, "b c (v t) h w -> b c v t h w", v=self.n_views)
         video_arrangement = [1, 0, 2, 4, 3, 5]
-	    # Fill one blank view for 5view
+        # Fill one blank view for 5view
         if self.n_views == 5:
-            ones_tensor = torch.zeros_like(video_segments[:, :, 0,],).unsqueeze(2)
+            ones_tensor = torch.zeros_like(
+                video_segments[
+                    :,
+                    :,
+                    0,
+                ],
+            ).unsqueeze(2)
             video_segments = torch.cat((video_segments, ones_tensor), dim=2)
             video_arrangement = [1, 0, 2, 3, 5, 4]
         grid_video = torch.stack(

scripts/download_diffusion_checkpoints.py

@@ -101,7 +101,7 @@ def convert_pixtral_checkpoint(checkpoint_dir: str, checkpoint_name: str, vit_ty
         allow_patterns=["params.json", "consolidated.safetensors"],
         local_dir=pixtral_ckpt_dir,
         local_dir_use_symlinks=False,
-        revision="db3e3ed01201248694fcb170c7bd292ecfcad22b"
+        revision="db3e3ed01201248694fcb170c7bd292ecfcad22b",
     )
     orig_dtype = torch.get_default_dtype()
     dtype = torch.bfloat16