Refactor Cosmos 2.5 model and tests for improved precision and functionality

- Added torchvision transforms import to the Cosmos 2.5 model for enhanced functionality.
- Updated max_size calculation in positional embedding classes to enforce strict behavior.
- Changed precision from bfloat16 to float32 in test cases to improve numerical stability.
- Tightened numerical difference assertions in tests to ensure higher accuracy in model outputs.

These changes enhance the model's robustness and ensure better alignment with reference implementations.

Author: KyleShao1016

fastvideo/models/dits/cosmos2_5.py

@@ -6,6 +6,7 @@
 import numpy as np
 import torch
 import torch.nn as nn
+from torchvision import transforms
 
 from fastvideo.attention import DistributedAttention, LocalAttention
 from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
@@ -18,6 +19,7 @@
 from fastvideo.platforms import AttentionBackendEnum
 
 
+
 class Cosmos25PatchEmbed(nn.Module):
     """
     COSMOS 2.5 patch embedding - converts video (B, C, T, H, W) to patches (B, T', H', W', D).
@@ -579,7 +581,7 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        self.max_size = [size // patch for size, patch in zip(max_size, patch_size, strict=False)]
+        self.max_size = [size // patch for size, patch in zip(max_size, patch_size, strict=True)]
         self.patch_size = patch_size
         self.base_fps = base_fps
         self.enable_fps_modulation = enable_fps_modulation
@@ -668,7 +670,7 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        self.max_size = [size // patch for size, patch in zip(max_size, patch_size, strict=False)]
+        self.max_size = [size // patch for size, patch in zip(max_size, patch_size, strict=True)]
         self.patch_size = patch_size
         self.eps = eps
 
@@ -907,7 +909,6 @@ def forward(
 
         # 2. Concatenate padding mask if needed
         if self.concat_padding_mask and padding_mask is not None:
-            from torchvision import transforms
             padding_mask = transforms.functional.resize(
                 padding_mask,
                 list(hidden_states.shape[-2:]),

fastvideo/tests/transformers/test_cosmos2_5.py

@@ -177,8 +177,7 @@ def load_reference_cosmos25_model(checkpoint_path: str, device, dtype):
 def test_cosmos25_transformer():
     """Test COSMOS 2.5 transformer against reference implementation."""
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16
-    precision_str = "bfloat16"
+    precision = torch.float32
 
     # Create COSMOS 2.5 specific config
     from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25ArchConfig
@@ -441,8 +440,8 @@ def test_cosmos25_transformer():
 
 
         # Allow for some numerical differences due to implementation details
-        assert max_diff < 1e-1, f"Maximum difference too large: {max_diff.item()}"
-        assert mean_diff < 1e-2, f"Mean difference too large: {mean_diff.item()}"
+        assert max_diff < 1e-4, f"Maximum difference too large: {max_diff.item()}"
+        assert mean_diff < 1e-5, f"Mean difference too large: {mean_diff.item()}"
 
         logger.info("✓ COSMOS 2.5 FastVideo implementation matches reference!")
     else:
@@ -454,8 +453,7 @@ def test_cosmos25_transformer():
 def test_cosmos25_transformer_video():
     """Test COSMOS 2.5 transformer with video input (multiple frames)."""
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16
-    precision_str = "bfloat16"
+    precision = torch.float32
 
     # Create COSMOS 2.5 specific config
     from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25ArchConfig