pre-commit

Author: kevin314

fastvideo/configs/models/dits/__init__.py

@@ -1,7 +1,9 @@
+from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.stepvideo import StepVideoConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
-print("WOW")
-from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
 
-__all__ = ["HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig", "CosmosVideoConfig"]
+__all__ = [
+    "HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig",
+    "CosmosVideoConfig"
+]

fastvideo/configs/models/dits/cosmos.py

@@ -90,7 +90,6 @@ class CosmosArchConfig(DiTArchConfig):
     qk_norm: str = "rms_norm"
     eps: float = 1e-6
     exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
-    
 
     def __post_init__(self):
         super().__post_init__()

fastvideo/configs/models/encoders/t5.py

@@ -94,6 +94,7 @@ class T5Config(TextEncoderConfig):
 @dataclass
 class T5LargeConfig(TextEncoderConfig):
     """T5 Large configuration for your specific model."""
-    arch_config: TextEncoderArchConfig = field(default_factory=T5LargeArchConfig)
+    arch_config: TextEncoderArchConfig = field(
+        default_factory=T5LargeArchConfig)
 
     prefix: str = "t5"

fastvideo/configs/models/vaes/__init__.py

@@ -1,7 +1,7 @@
+from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
 from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
 from fastvideo.configs.models.vaes.stepvideovae import StepVideoVAEConfig
 from fastvideo.configs.models.vaes.wanvae import WanVAEConfig
-from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
 
 __all__ = [
     "HunyuanVAEConfig",

fastvideo/configs/pipelines/__init__.py

@@ -1,14 +1,13 @@
 from fastvideo.configs.pipelines.base import (PipelineConfig,
                                               SlidingTileAttnConfig)
+from fastvideo.configs.pipelines.cosmos import CosmosConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.registry import (
     get_pipeline_config_cls_from_name)
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
 from fastvideo.configs.pipelines.wan import (WanI2V480PConfig, WanI2V720PConfig,
                                              WanT2V480PConfig, WanT2V720PConfig)
 
-from fastvideo.configs.pipelines.cosmos import CosmosConfig
-
 __all__ = [
     "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig",
     "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig",

fastvideo/configs/pipelines/cosmos.py

@@ -5,10 +5,8 @@
 import torch
 
 from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
-
 from fastvideo.configs.models.dits import CosmosVideoConfig
-from fastvideo.configs.models.encoders import (BaseEncoderOutput,
-                                                  T5LargeConfig)
+from fastvideo.configs.models.encoders import BaseEncoderOutput, T5LargeConfig
 from fastvideo.configs.models.vaes import CosmosVAEConfig
 from fastvideo.configs.pipelines.base import PipelineConfig
 
@@ -19,39 +17,35 @@ def t5_large_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
     Return raw last_hidden_state without truncation/padding.
     """
     hidden_state = outputs.last_hidden_state
-    
+
     if hidden_state is None:
         raise ValueError("T5 Large outputs missing last_hidden_state")
-    
+
     nan_count = torch.isnan(hidden_state).sum()
     if nan_count > 0:
         hidden_state = hidden_state.masked_fill(torch.isnan(hidden_state), 0.0)
-    
+
     return hidden_state
 
 
 @dataclass
 class CosmosConfig(PipelineConfig):
     """Configuration for Cosmos2 Video2World pipeline matching diffusers."""
 
-
     dit_config: DiTConfig = field(default_factory=CosmosVideoConfig)
-    
 
     vae_config: VAEConfig = field(default_factory=CosmosVAEConfig)
-    
 
     text_encoder_configs: tuple[EncoderConfig, ...] = field(
         default_factory=lambda: (T5LargeConfig(), ))
     postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
                                   ...] = field(default_factory=lambda:
                                                (t5_large_postprocess_text, ))
 
-
     dit_precision: str = "bf16"
     vae_precision: str = "fp16"
     text_encoder_precisions: tuple[str, ...] = field(
-        default_factory=lambda: ("bf16",))
+        default_factory=lambda: ("bf16", ))
 
     conditioning_strategy: str = "frame_replace"
     min_num_conditional_frames: int = 1
@@ -61,13 +55,12 @@ class CosmosConfig(PipelineConfig):
     state_ch: int = 16
     state_t: int = 24
     text_encoder_class: str = "T5"
-    
 
     embedded_cfg_scale: int = 6
-    flow_shift: float = 1.0 
+    flow_shift: float = 1.0
 
     def __post_init__(self):
         self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
-        
-        self._vae_latent_dim = 16
+
+        self._vae_latent_dim = 16

fastvideo/configs/pipelines/registry.py

@@ -5,13 +5,13 @@
 from collections.abc import Callable
 
 from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.pipelines.cosmos import CosmosConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
 from fastvideo.configs.pipelines.wan import (FastWan2_1_T2V_480P_Config,
                                              FastWan2_2_TI2V_5B_Config,
                                              WanI2V480PConfig, WanI2V720PConfig,
                                              WanT2V480PConfig, WanT2V720PConfig)
-from fastvideo.configs.pipelines.cosmos import CosmosConfig
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
                              verify_model_config_and_directory)

fastvideo/image_processor.py

@@ -4,8 +4,6 @@
 This module provides lightweight image preprocessing without external dependencies beyond PyTorch/NumPy/PIL.
 """
 
-from typing import Optional, Union
-
 import numpy as np
 import PIL.Image
 import torch
@@ -29,9 +27,9 @@ def __init__(self, vae_scale_factor: int = 8) -> None:
 
     def preprocess(
         self,
-        image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
-        height: Optional[int] = None,
-        width: Optional[int] = None,
+        image: PIL.Image.Image | np.ndarray | torch.Tensor,
+        height: int | None = None,
+        width: int | None = None,
     ) -> torch.Tensor:
         """
         Preprocess an image to a normalized torch tensor.
@@ -55,14 +53,13 @@ def preprocess(
         else:
             raise ValueError(
                 f"Unsupported image type: {type(image)}. "
-                "Supported types: PIL.Image.Image, np.ndarray, torch.Tensor"
-            )
+                "Supported types: PIL.Image.Image, np.ndarray, torch.Tensor")
 
     def _preprocess_pil(
         self,
         image: PIL.Image.Image,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
+        height: int | None = None,
+        width: int | None = None,
     ) -> torch.Tensor:
         """Preprocess a PIL image."""
         if height is None:
@@ -73,7 +70,8 @@ def _preprocess_pil(
         height = height - (height % self.vae_scale_factor)
         width = width - (width % self.vae_scale_factor)
 
-        image = image.resize((width, height), resample=PIL.Image.Resampling.LANCZOS)
+        image = image.resize((width, height),
+                             resample=PIL.Image.Resampling.LANCZOS)
 
         image_np = np.array(image, dtype=np.float32) / 255.0
 
@@ -85,8 +83,8 @@ def _preprocess_pil(
     def _preprocess_numpy(
         self,
         image: np.ndarray,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
+        height: int | None = None,
+        width: int | None = None,
     ) -> torch.Tensor:
         """Preprocess a numpy array."""
         # Determine target dimensions if not provided
@@ -115,7 +113,8 @@ def _preprocess_numpy(
                 image_uint8 = image.astype(np.uint8)
             pil_image = PIL.Image.fromarray(image_uint8)
 
-        pil_image = pil_image.resize((width, height), resample=PIL.Image.Resampling.LANCZOS)
+        pil_image = pil_image.resize((width, height),
+                                     resample=PIL.Image.Resampling.LANCZOS)
         image_np = np.array(pil_image, dtype=np.float32) / 255.0
 
         # Ensure 3D shape
@@ -127,8 +126,8 @@ def _preprocess_numpy(
     def _preprocess_tensor(
         self,
         image: torch.Tensor,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
+        height: int | None = None,
+        width: int | None = None,
     ) -> torch.Tensor:
         """Preprocess a torch tensor."""
         # Determine target dimensions
@@ -158,9 +157,10 @@ def _preprocess_tensor(
             else:  # (H, W, C) - need to rearrange
                 image = image.permute(2, 0, 1).unsqueeze(0)  # (1, C, H, W)
 
-        image = torch.nn.functional.interpolate(
-            image, size=(height, width), mode="bilinear", align_corners=False
-        )
+        image = torch.nn.functional.interpolate(image,
+                                                size=(height, width),
+                                                mode="bilinear",
+                                                align_corners=False)
 
         if image.max() > 1.0:  # Assume [0, 255] range
             image = image / 255.0
@@ -181,9 +181,11 @@ def _normalize_to_tensor(self, image_np: np.ndarray) -> torch.Tensor:
         """
         # Convert to tensor
         if image_np.ndim == 2:  # (H, W) - grayscale
-            tensor = torch.from_numpy(image_np).unsqueeze(0).unsqueeze(0)  # (1, 1, H, W)
+            tensor = torch.from_numpy(image_np).unsqueeze(0).unsqueeze(
+                0)  # (1, 1, H, W)
         elif image_np.ndim == 3:  # (H, W, C)
-            tensor = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(0)  # (1, C, H, W)
+            tensor = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(
+                0)  # (1, C, H, W)
         else:
             raise ValueError(f"Expected 2D or 3D array, got {image_np.ndim}D")
 

fastvideo/layers/layernorm.py

@@ -43,7 +43,8 @@ def forward_diffusers(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """Forward method that matches Diffusers RMSNorm implementation exactly."""
         input_dtype = hidden_states.dtype
         variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        hidden_states = hidden_states * torch.rsqrt(variance +
+                                                    self.variance_epsilon)
 
         if self.has_weight and self.weight is not None:
             if self.weight.dtype in [torch.float16, torch.bfloat16]:

fastvideo/layers/rotary_embedding.py

@@ -46,10 +46,10 @@ def _rotate_gptj(x: torch.Tensor) -> torch.Tensor:
 
 def apply_rotary_emb(
     x: torch.Tensor,
-    freqs_cis: torch.Tensor | tuple[torch.Tensor],
+    freqs_cis: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
     use_real: bool = True,
     use_real_unbind_dim: int = -1,
-) -> tuple[torch.Tensor, torch.Tensor]:
+) -> torch.Tensor:
     """
     Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
     to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are