[Do not merge] V1 encoders and model loading (#261)

Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>

assets/prompt.txt

@@ -1,5 +1,5 @@
-A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature.
 Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.
+A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature.
 A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere.
 A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest. The playful yet serene atmosphere is complemented by soft natural light filtering through the petals. Mid-shot, warm and cheerful tones.
 A superintelligent humanoid robot waking up. The robot has a sleek metallic body with futuristic design features. Its glowing red eyes are the focal point, emanating a sharp, intense light as it powers on. The scene is set in a dimly lit, high-tech laboratory filled with glowing control panels, robotic arms, and holographic screens. The setting emphasizes advanced technology and an atmosphere of mystery. The ambiance is eerie and dramatic, highlighting the moment of awakening and the robots immense intelligence. Photorealistic style with a cinematic, dark sci-fi aesthetic. Aspect ratio: 16:9 --v 6.1

fastvideo/attention/__init__.py

@@ -1 +1,3 @@
-from .distributed_attn import DistributedAttention
+from .flash_attn import (DistributedAttention, TextEncoderFlashAttention, LocalAttention)
+
+__all__ = ["DistributedAttention", "TextEncoderFlashAttention", "LocalAttention"]

fastvideo/attention/flash_attn.py

@@ -1,9 +1,11 @@
+from itertools import accumulate
+from typing import List, Optional
+
 import torch
 import torch.nn as nn
-from typing import Optional
 from fastvideo.distributed.communication_op import sequence_model_parallel_all_to_all_4D, sequence_model_parallel_all_gather
 from fastvideo.distributed.parallel_state import get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size
-from flash_attn import flash_attn_func
+from flash_attn import flash_attn_func, flash_attn_varlen_func
 
 
 class DistributedAttention(nn.Module):
@@ -33,6 +35,7 @@ class DistributedAttention(nn.Module):
         self.dropout_rate = dropout_rate
         self.causal = causal
         self.softmax_scale = softmax_scale
+
     def forward(
         self,
         q: torch.Tensor,
@@ -132,4 +135,99 @@ class LocalAttention(nn.Module):
             causal=self.causal
         )
         
-        return output
+        return output
+
+class TextEncoderFlashAttention:
+    """Simplified Flash Attention implementation for text encoders in diffusion models.
+    
+    This implementation removes KV caching and other generation-specific features,
+    focusing only on the encoder attention pattern needed for text encoders in
+    diffusion models like DiT.
+    """
+
+    def __init__(
+        self,
+        softmax_scale: Optional[float] = None,
+    ) -> None:
+        """Initialize the text encoder flash attention implementation.
+        
+        Args:
+            num_heads: Number of attention heads
+            head_size: Size of each attention head
+            scale: Attention scale factor. If None, defaults to 1/sqrt(head_size)
+        """
+        self.softmax_scale = softmax_scale
+        
+        # Set defaults for optional parameters
+        self.alibi_slopes = None
+        self.sliding_window = (-1, -1)  # No sliding window by default
+        self.logits_soft_cap = 0  # No soft cap by default
+        
+        
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        seq_lens: Optional[List[int]] = None,
+        output: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Forward pass with FlashAttention for text encoder.
+        
+        Args:
+            query: shape = [num_tokens, num_heads, head_size]
+            key: shape = [num_tokens, num_kv_heads, head_size]
+            value: shape = [num_tokens, num_kv_heads, head_size]
+            seq_lens: List of sequence lengths for batched input
+            output: shape = [num_tokens, num_heads, head_size]
+                   Pre-allocated output tensor (optional)
+                   
+        Returns:
+            torch.Tensor: Output tensor with shape [num_tokens, num_heads, head_size]
+        """
+        device = query.device
+        batch_size = len(seq_lens) if seq_lens is not None else 1
+        
+        if output is None:
+            output = torch.empty_like(query)
+            
+        if seq_lens is None:
+            # Single sequence case - use standard flash attention
+            return flash_attn_varlen_func(
+                q=query,
+                k=key,
+                v=value,
+                softmax_scale=self.softmax_scale,
+                causal=False,  # Non-causal for encoder
+                # window_size=self.sliding_window,
+                # alibi_slopes=self.alibi_slopes,
+                # softcap=self.logits_soft_cap,
+                out=output,
+                # fa_version=self.vllm_flash_attn_version,
+            )
+        else:
+            assert False, "Not implemented"
+            # Multiple sequences case - use varlen flash attention
+            max_seq_len = max(seq_lens)
+            seq_start_loc = list(accumulate(seq_lens, initial=0))
+            seq_start_loc_tensor = torch.tensor(seq_start_loc, 
+                                               dtype=torch.int32, 
+                                               device=device)
+            
+            return flash_attn_varlen_func(
+                q=query,
+                k=key,
+                v=value,
+                cu_seqlens_q=seq_start_loc_tensor,
+                cu_seqlens_k=seq_start_loc_tensor,
+                max_seqlen_q=max_seq_len,
+                max_seqlen_k=max_seq_len,
+                softmax_scale=self.scale,
+                causal=False,  # Non-causal for encoder
+                window_size=self.sliding_window,
+                alibi_slopes=self.alibi_slopes,
+                softcap=self.logits_soft_cap,
+                out=output,
+                fa_version=self.vllm_flash_attn_version,
+            )

fastvideo/attention/layer.py

@@ -0,0 +1,135 @@
+import torch
+import torch.nn as nn
+from typing import Optional, Dict, Any, Type
+
+from fastvideo.attention.flash_attn import LocalAttention, DistributedAttention
+
+
+class AttentionBackend:
+    """Base class for attention backends that provides implementation classes."""
+    
+    @staticmethod
+    def get_impl_cls() -> Type:
+        """Returns the implementation class for the attention backend."""
+        raise NotImplementedError("Subclasses must implement get_impl_cls")
+
+
+class LocalAttentionBackend(AttentionBackend):
+    """Backend for local attention implementation."""
+    
+    @staticmethod
+    def get_impl_cls() -> Type:
+        """Returns the LocalAttention implementation class."""
+        return LocalAttention
+
+
+class DistributedAttentionBackend(AttentionBackend):
+    """Backend for distributed attention implementation."""
+    
+    @staticmethod
+    def get_impl_cls() -> Type:
+        """Returns the DistributedAttention implementation class."""
+        return DistributedAttention
+
+
+class AttentionLayer(nn.Module):
+    """Base attention layer that uses an implementation class from a backend."""
+    
+    def __init__(
+        self,
+        attn_backend: AttentionBackend,
+        dropout_rate: float = 0.0,
+        causal: bool = False,
+        softmax_scale: Optional[float] = None,
+        **extra_impl_args
+    ):
+        super().__init__()
+        impl_cls = attn_backend.get_impl_cls()
+        self.impl = impl_cls(
+            dropout_rate=dropout_rate,
+            causal=causal,
+            softmax_scale=softmax_scale,
+            **extra_impl_args
+        )
+    
+    def forward(self, *args, **kwargs):
+        """Forward pass delegated to the implementation."""
+        return self.impl(*args, **kwargs)
+
+
+class LocalAttentionLayer(AttentionLayer):
+    """Wrapper for local attention implementation."""
+    
+    def __init__(
+        self,
+        dropout_rate: float = 0.0,
+        causal: bool = False,
+        softmax_scale: Optional[float] = None,
+        **extra_impl_args
+    ):
+        super().__init__(
+            LocalAttentionBackend(),
+            dropout_rate=dropout_rate,
+            causal=causal,
+            softmax_scale=softmax_scale,
+            **extra_impl_args
+        )
+    
+    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+        """Apply local attention.
+        
+        Args:
+            q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim]
+            k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim]
+            v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim]
+            
+        Returns:
+            torch.Tensor: Output tensor after attention
+        """
+        return self.impl(q, k, v)
+
+
+class DistributedAttentionLayer(AttentionLayer):
+    """Wrapper for distributed attention implementation."""
+    
+    def __init__(
+        self,
+        dropout_rate: float = 0.0,
+        causal: bool = False,
+        softmax_scale: Optional[float] = None,
+        **extra_impl_args
+    ):
+        super().__init__(
+            DistributedAttentionBackend(),
+            dropout_rate=dropout_rate,
+            causal=causal,
+            softmax_scale=softmax_scale,
+            **extra_impl_args
+        )
+    
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        replicated_q: Optional[torch.Tensor] = None,
+        replicated_k: Optional[torch.Tensor] = None,
+        replicated_v: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Forward pass for distributed attention.
+        
+        Args:
+            q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim]
+            k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim]
+            v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim]
+            replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens
+            replicated_k (Optional[torch.Tensor]): Replicated key tensor
+            replicated_v (Optional[torch.Tensor]): Replicated value tensor
+            
+        Returns:
+            Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing:
+                - o (torch.Tensor): Output tensor after attention for the main sequence
+                - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided
+        """
+        return self.impl(q, k, v, replicated_q, replicated_k, replicated_v)
+

fastvideo/distributed/parallel_state.py

@@ -850,7 +850,6 @@ def get_sp_group() -> GroupCoordinator:
 
 def initialize_model_parallel(
     tensor_model_parallel_size: int = 1,
-    pipeline_model_parallel_size: int = 1,
     sequence_model_parallel_size: int = 1,
     backend: Optional[str] = None,
 ) -> None:
@@ -885,15 +884,15 @@ def initialize_model_parallel(
         get_world_group().device_group)
 
     # Ensure the world size is compatible with the parallelism configuration
-    assert world_size % (tensor_model_parallel_size * pipeline_model_parallel_size * sequence_model_parallel_size) == 0, \
-        f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size}) * " \
-        f"pipeline_model_parallel_size ({pipeline_model_parallel_size}) * " \
-        f"sequence_model_parallel_size ({sequence_model_parallel_size})"
+    # assert world_size % (tensor_model_parallel_size * pipeline_model_parallel_size * sequence_model_parallel_size) == 0, \
+    #     f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size}) * " \
+    #     f"pipeline_model_parallel_size ({pipeline_model_parallel_size}) * " \
+    #     f"sequence_model_parallel_size ({sequence_model_parallel_size})"
 
     # Check for incompatible parallelism configurations
-    if sequence_model_parallel_size > 1:
-        assert tensor_model_parallel_size == 1, "Sequence parallelism (SP) is incompatible with tensor parallelism (TP). Please use SP=1 or TP=1."
-        assert pipeline_model_parallel_size == 1, "Sequence parallelism (SP) is incompatible with pipeline parallelism (PP). Please use SP=1 or PP=1."
+    # if sequence_model_parallel_size > 1:
+    #     assert tensor_model_parallel_size == 1, "Sequence parallelism (SP) is incompatible with tensor parallelism (TP). Please use SP=1 or TP=1."
+    #     assert pipeline_model_parallel_size == 1, "Sequence parallelism (SP) is incompatible with pipeline parallelism (PP). Please use SP=1 or PP=1."
 
     # Build the tensor model-parallel groups.
     num_tensor_model_parallel_groups: int = (world_size //
@@ -933,21 +932,6 @@ def initialize_model_parallel(
                                    backend,
                                    group_name="sp")
 
-    # Build the pipeline model-parallel groups.
-    num_pipeline_model_parallel_groups: int = (world_size //
-                                               pipeline_model_parallel_size)
-    global _PP
-    assert _PP is None, (
-        "pipeline model parallel group is already initialized")
-    group_ranks = []
-    for i in range(num_pipeline_model_parallel_groups):
-        ranks = list(range(i, world_size, num_pipeline_model_parallel_groups))
-        group_ranks.append(ranks)
-    _PP = init_model_parallel_group(group_ranks,
-                                    get_world_group().local_rank,
-                                    backend,
-                                    group_name="pp")
-
 
 def get_sequence_model_parallel_world_size():
     """Return world size for the sequence model parallel group."""
@@ -1172,3 +1156,267 @@ def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],
             aggregated_data += rank_data
 
     return [x == 1 for x in aggregated_data.tolist()]
+
+
+def initialize_tensor_parallel_group(
+    tensor_model_parallel_size: int = 1,
+    backend: Optional[str] = None,
+    group_name_suffix: str = ""
+) -> GroupCoordinator:
+    """Initialize a tensor parallel group for a specific model.
+    
+    This function creates a tensor parallel group that can be used with the
+    patch_tensor_parallel_group context manager. It allows different models
+    to use different tensor parallelism configurations.
+    
+    Arguments:
+        tensor_model_parallel_size: number of GPUs used for tensor model parallelism.
+        backend: communication backend to use.
+        group_name_suffix: optional suffix to make the group name unique.
+        
+    Returns:
+        A GroupCoordinator for tensor parallelism that can be used with
+        the patch_tensor_parallel_group context manager.
+        
+    Example usage:
+        ```python
+        # Initialize tensor parallel group for model1
+        tp_group_model1 = initialize_tensor_parallel_group(
+            tensor_model_parallel_size=4,
+            group_name_suffix="model1"
+        )
+        
+        # Use tensor parallelism for model1
+        with patch_tensor_parallel_group(tp_group_model1):
+            # Run model1 with tensor parallelism
+            output1 = model1(input1)
+        ```
+    """
+    # Get world size and rank. Ensure some consistencies.
+    assert torch.distributed.is_initialized()
+    world_size: int = torch.distributed.get_world_size()
+    backend = backend or torch.distributed.get_backend(
+        get_world_group().device_group)
+
+    # Ensure the world size is compatible with the parallelism configuration
+    assert world_size % tensor_model_parallel_size == 0, \
+        f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size})"
+
+    # Build the tensor model-parallel groups.
+    num_tensor_model_parallel_groups: int = (world_size // tensor_model_parallel_size)
+    tp_group_ranks = []
+    for i in range(num_tensor_model_parallel_groups):
+        ranks = list(
+            range(i * tensor_model_parallel_size,
+                  (i + 1) * tensor_model_parallel_size))
+        tp_group_ranks.append(ranks)
+
+    # Create TP group coordinator with a unique name
+    group_name = f"tp_{group_name_suffix}" if group_name_suffix else "tp"
+    tp_group = init_model_parallel_group(tp_group_ranks,
+                                        get_world_group().local_rank,
+                                        backend,
+                                        use_message_queue_broadcaster=True,
+                                        group_name=group_name)
+
+    return tp_group
+
+
+def initialize_sequence_parallel_group(
+    sequence_model_parallel_size: int = 1,
+    backend: Optional[str] = None,
+    group_name_suffix: str = ""
+) -> GroupCoordinator:
+    """Initialize a sequence parallel group for a specific model.
+    
+    This function creates a sequence parallel group that can be used with the
+    patch_sequence_parallel_group context manager. It allows different models
+    to use different sequence parallelism configurations.
+    
+    Arguments:
+        sequence_model_parallel_size: number of GPUs used for sequence model parallelism.
+        backend: communication backend to use.
+        group_name_suffix: optional suffix to make the group name unique.
+        
+    Returns:
+        A GroupCoordinator for sequence parallelism that can be used with
+        the patch_sequence_parallel_group context manager.
+        
+    Example usage:
+        ```python
+        # Initialize sequence parallel group for model2
+        sp_group_model2 = initialize_sequence_parallel_group(
+            sequence_model_parallel_size=2,
+            group_name_suffix="model2"
+        )
+        
+        # Use sequence parallelism for model2
+        with patch_sequence_parallel_group(sp_group_model2):
+            # Run model2 with sequence parallelism
+            output2 = model2(input2)
+        ```
+    """
+    # Get world size and rank. Ensure some consistencies.
+    assert torch.distributed.is_initialized()
+    world_size: int = torch.distributed.get_world_size()
+    backend = backend or torch.distributed.get_backend(
+        get_world_group().device_group)
+
+    # Ensure the world size is compatible with the parallelism configuration
+    assert world_size % sequence_model_parallel_size == 0, \
+        f"World size ({world_size}) must be divisible by sequence_model_parallel_size ({sequence_model_parallel_size})"
+
+    # Build the sequence model-parallel groups.
+    num_sequence_model_parallel_groups: int = (world_size // sequence_model_parallel_size)
+    sp_group_ranks = []
+    
+    for i in range(num_sequence_model_parallel_groups):
+        # Create groups of consecutive ranks
+        ranks = list(range(i * sequence_model_parallel_size, 
+                           (i + 1) * sequence_model_parallel_size))
+        sp_group_ranks.append(ranks)
+
+    # Create SP group coordinator with a unique name
+    group_name = f"sp_{group_name_suffix}" if group_name_suffix else "sp"
+    sp_group = init_model_parallel_group(sp_group_ranks,
+                                       get_world_group().local_rank,
+                                       backend,
+                                       group_name=group_name)
+
+    return sp_group
+
+
+_SP_STATE_PATCHED = False
+
+@contextmanager
+def patch_sequence_parallel_group(sp_group: GroupCoordinator):
+    """Patch the sp group temporarily until this function ends.
+    
+    This method allows running a model with SP while another model uses TP.
+    
+    Args:
+        sp_group (GroupCoordinator): the sp group coordinator
+    """
+    global _SP_STATE_PATCHED
+    assert not _SP_STATE_PATCHED, "Should not call when it's already patched"
+
+    _SP_STATE_PATCHED = True
+    old_sp_group = get_sp_group()
+    global _SP
+    _SP = sp_group
+    try:
+        yield
+    finally:
+        # restore the original state
+        _SP_STATE_PATCHED = False
+        _SP = old_sp_group
+
+
+_PP_STATE_PATCHED = False
+
+@contextmanager
+def patch_pipeline_parallel_group(pp_group: GroupCoordinator):
+    """Patch the pp group temporarily until this function ends.
+    
+    This method allows running a model with a different PP configuration.
+    
+    Args:
+        pp_group (GroupCoordinator): the pp group coordinator
+    """
+    global _PP_STATE_PATCHED
+    assert not _PP_STATE_PATCHED, "Should not call when it's already patched"
+
+    _PP_STATE_PATCHED = True
+    old_pp_group = get_pp_group()
+    global _PP
+    _PP = pp_group
+    try:
+        yield
+    finally:
+        # restore the original state
+        _PP_STATE_PATCHED = False
+        _PP = old_pp_group
+
+
+# Example of how to use the independent parallelism functions
+"""
+Here's a complete example of how to use the independent parallelism functions
+for different models:
+
+```python
+import torch
+from fastvideo.distributed.parallel_state import (
+    init_distributed_environment,
+    initialize_tensor_parallel_group,
+    initialize_sequence_parallel_group,
+    patch_tensor_parallel_group,
+    patch_sequence_parallel_group,
+    destroy_model_parallel,
+    destroy_distributed_environment
+)
+
+# Initialize the distributed environment
+init_distributed_environment(
+    world_size=8,
+    rank=torch.distributed.get_rank(),
+    distributed_init_method="tcp://localhost:12345",
+    local_rank=torch.distributed.get_rank() % torch.cuda.device_count()
+)
+
+try:
+    # Create a tensor parallel group for model1 with TP=4
+    tp_group_model1 = initialize_tensor_parallel_group(
+        tensor_model_parallel_size=4,
+        group_name_suffix="model1"
+    )
+    
+    # Create a sequence parallel group for model2 with SP=2
+    sp_group_model2 = initialize_sequence_parallel_group(
+        sequence_model_parallel_size=2,
+        group_name_suffix="model2"
+    )
+    
+    # Create another tensor parallel group for model3 with TP=2
+    tp_group_model3 = initialize_tensor_parallel_group(
+        tensor_model_parallel_size=2,
+        group_name_suffix="model3"
+    )
+    
+    # Use model1 with tensor parallelism
+    with patch_tensor_parallel_group(tp_group_model1):
+        # Inside this context, get_tp_group() returns tp_group_model1
+        # Run model1 with tensor parallelism
+        output1 = model1(input1)
+    
+    # Use model2 with sequence parallelism
+    with patch_sequence_parallel_group(sp_group_model2):
+        # Inside this context, get_sp_group() returns sp_group_model2
+        # Run model2 with sequence parallelism
+        output2 = model2(input2)
+    
+    # Use model3 with a different tensor parallelism configuration
+    with patch_tensor_parallel_group(tp_group_model3):
+        # Inside this context, get_tp_group() returns tp_group_model3
+        # Run model3 with tensor parallelism
+        output3 = model3(input3)
+    
+    # You can switch between models as needed
+    with patch_tensor_parallel_group(tp_group_model1):
+        # Back to using model1
+        more_output1 = model1(more_input1)
+    
+finally:
+    # Clean up
+    destroy_model_parallel()
+    destroy_distributed_environment()
+```
+
+This approach allows you to:
+1. Create separate parallel groups for each model
+2. Use different parallelism strategies for different models
+3. Switch between models as needed
+4. Use unique group names to avoid conflicts
+
+Note that each model can use its own optimal parallelism strategy without
+interfering with other models.
+"""

fastvideo/inference.py

@@ -1,250 +0,0 @@
-import os
-import random
-import time
-from pathlib import Path
-from typing import Dict, List, Optional, Union, Any
-
-import torch
-from safetensors.torch import load_file as safetensors_load_file
-import json
-
-from fastvideo.models.hunyuan.constants import NEGATIVE_PROMPT, PRECISION_TO_TYPE, PROMPT_TEMPLATE
-# from fastvideo.models.hunyuan.diffusion.pipelines import HunyuanVideoPipeline
-from fastvideo.pipelines.pipeline_base import DiffusionPipelineBase
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler
-# from fastvideo.models.hunyuan.modules import load_model
-# from fastvideo.models.hunyuan.text_encoder import TextEncoder
-from fastvideo.models.hunyuan.utils.data_utils import align_to
-# from fastvideo.models.hunyuan.vae import load_vae
-# from fastvideo.utils.parallel_states import nccl_info
-from fastvideo.inference_args import InferenceArgs
-from fastvideo.platforms import current_platform
-from fastvideo.logger import init_logger
-from fastvideo.pipelines.loader import PipelineLoader, get_pipeline_loader
-from fastvideo.distributed.parallel_state import get_sp_group
-
-logger = init_logger(__name__)
-
-
-class DiffusionInference:
-    """
-    Unified inference class that works with any diffusion pipeline.
-    This combines model loading, pipeline creation, and inference in a flexible way.
-    """
-    
-    def __init__(self, 
-                args: InferenceArgs,
-                pipeline: DiffusionPipelineBase,
-                default_negative_prompt: str = NEGATIVE_PROMPT):
-        """
-        Initialize the inference class with a pipeline and args.
-        
-        Args:
-            args: The inference arguments
-            pipeline: The diffusion pipeline to use
-            default_negative_prompt: The default negative prompt to use
-        """
-        self.args = args
-        self.pipeline = pipeline
-        self.default_negative_prompt = default_negative_prompt
-
-    
-        
-    @classmethod
-    def load_pipeline(cls, inference_args: InferenceArgs):
-        """
-        Create an inference instance from pretrained model components.
-        
-        Args:
-            inference_args: The inference arguments containing model path and other settings
-            
-        Returns:
-            A DiffusionInference instance ready for inference
-        """
-        logger.debug("start loading pipeline")
-        models_root_path = Path(inference_args.model_path)
-        if not models_root_path.exists():
-            raise ValueError(f"`models_root` not exists: {models_root_path}")
-
-        # Create save folder to save the samples
-        save_path = inference_args.output_path
-        os.makedirs(os.path.dirname(save_path), exist_ok=True)
-
-        # Log which platform we're using
-        print(f"Using platform: {current_platform.__class__.__name__}")
-        assert current_platform.is_cuda(), "CUDA is not available"
-
-        # init distributed
-        # _initialize_distributed()
-
-        
-        # get loader for the pipeline
-
-        # Create pipeline using the components
-        pipeline_loader = get_pipeline_loader(inference_args)
-        pipeline = pipeline_loader.load_pipeline(
-            inference_args=inference_args,
-        )
-        
-        return cls(inference_args, pipeline, default_negative_prompt=NEGATIVE_PROMPT)
-    
-    def predict(self, prompt: str, inference_args: InferenceArgs) -> Dict[str, Any]:
-        """
-        Run inference with the pipeline.
-            
-        Returns:
-            Dictionary with generated samples and metadata
-        """
-
-        out_dict = dict()
-
-        batch_size = inference_args.batch_size
-        num_videos_per_prompt = inference_args.num_videos
-        seed = inference_args.seed
-        height = inference_args.height
-        width = inference_args.width
-        video_length = inference_args.num_frames
-        # prompt = inference_args.prompt
-        negative_prompt = inference_args.neg_prompt
-        infer_steps = inference_args.num_inference_steps
-        guidance_scale = inference_args.guidance_scale
-        flow_shift = inference_args.flow_shift
-        # mask_strategy = inference_args.mask_strategy
-        embedded_guidance_scale = inference_args.embedded_cfg_scale
-        
-
-        # ========================================================================
-        # Arguments: seed
-        # ========================================================================
-        if isinstance(seed, torch.Tensor):
-            seed = seed.tolist()
-        if seed is None:
-            seeds = [random.randint(0, 1_000_000) for _ in range(batch_size * num_videos_per_prompt)]
-        elif isinstance(seed, int):
-            seeds = [seed + i for _ in range(batch_size) for i in range(num_videos_per_prompt)]
-        elif isinstance(seed, (list, tuple)):
-            if len(seed) == batch_size:
-                seeds = [int(seed[i]) + j for i in range(batch_size) for j in range(num_videos_per_prompt)]
-            elif len(seed) == batch_size * num_videos_per_prompt:
-                seeds = [int(s) for s in seed]
-            else:
-                raise ValueError(
-                    f"Length of seed must be equal to number of prompt(batch_size) or "
-                    f"batch_size * num_videos_per_prompt ({batch_size} * {num_videos_per_prompt}), got {seed}.")
-        else:
-            raise ValueError(f"Seed must be an integer, a list of integers, or None, got {seed}.")
-        # Peiyuan: using GPU seed will cause A100 and H100 to generate different results...
-        generator = [torch.Generator("cpu").manual_seed(seed) for seed in seeds]
-        out_dict["seeds"] = seeds
-
-        # ========================================================================
-        # Arguments: target_width, target_height, target_video_length
-        # ========================================================================
-        if width <= 0 or height <= 0 or video_length <= 0:
-            raise ValueError(
-                f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={video_length}"
-            )
-        if (video_length - 1) % 4 != 0:
-            raise ValueError(f"`video_length-1` must be a multiple of 4, got {video_length}")
-
-        logger.info(f"Input (height, width, video_length) = ({height}, {width}, {video_length})")
-
-        target_height = align_to(height, 16)
-        target_width = align_to(width, 16)
-        target_video_length = video_length
-
-        out_dict["size"] = (target_height, target_width, target_video_length)
-
-        # ========================================================================
-        # Arguments: prompt, new_prompt, negative_prompt
-        # ========================================================================
-        if not isinstance(prompt, str):
-            raise TypeError(f"`prompt` must be a string, but got {type(prompt)}")
-        prompt = [prompt.strip()]
-
-        # negative prompt
-        if negative_prompt is None or negative_prompt == "":
-            negative_prompt = self.default_negative_prompt
-        if not isinstance(negative_prompt, str):
-            raise TypeError(f"`negative_prompt` must be a string, but got {type(negative_prompt)}")
-        negative_prompt = [negative_prompt.strip()]
-
-        # ========================================================================
-        # Scheduler
-        # ========================================================================
-        scheduler = FlowMatchDiscreteScheduler(
-            shift=flow_shift,
-            reverse=self.args.flow_reverse,
-            solver=self.args.flow_solver,
-        )
-        self.pipeline.scheduler = scheduler
-
-        if "884" in self.args.vae:
-            latents_size = [(video_length - 1) // 4 + 1, height // 8, width // 8]
-        elif "888" in self.args.vae:
-            latents_size = [(video_length - 1) // 8 + 1, height // 8, width // 8]
-        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
-
-        # ========================================================================
-        # Print infer args
-        # ========================================================================
-        debug_str = f"""
-                        height: {target_height}
-                         width: {target_width}
-                  video_length: {target_video_length}
-                        prompt: {prompt}
-                    neg_prompt: {negative_prompt}
-                          seed: {seed}
-                   infer_steps: {infer_steps}
-         num_videos_per_prompt: {num_videos_per_prompt}
-                guidance_scale: {guidance_scale}
-                      n_tokens: {n_tokens}
-                    flow_shift: {flow_shift}
-       embedded_guidance_scale: {embedded_guidance_scale}"""
-        logger.info(debug_str)
-        # return
-        # sp_group = get_sp_group()
-        # local_rank = sp_group.rank
-        import os
-        local_rank = os.environ.get("LOCAL_RANK", 0)
-        device = torch.device(f"cuda:{local_rank}")
-        batch = ForwardBatch(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_videos_per_prompt=num_videos_per_prompt,
-            height=inference_args.height,
-            width=inference_args.width,
-            num_frames=inference_args.num_frames,
-            num_inference_steps=inference_args.num_inference_steps,
-            guidance_scale=inference_args.guidance_scale,
-            generator=generator,
-            eta=0.0,
-            n_tokens=n_tokens,
-            data_type="video" if inference_args.num_frames > 1 else "image",
-            device=device,
-            extra={},  # Any additional parameters
-        )
-
-        print('===============================================')
-        print(batch)
-        print('===============================================')
-        print('===============================================')
-        print('===============================================')
-        print(inference_args)
-
-        # ========================================================================
-        # Pipeline inference
-        # ========================================================================
-        start_time = time.time()
-        samples = self.pipeline.forward(
-            batch=batch,
-            inference_args=inference_args,
-        )[0]
-        out_dict["samples"] = samples
-        out_dict["prompts"] = prompt
-
-        gen_time = time.time() - start_time
-        logger.info(f"Success, time: {gen_time}")
-
-        return out_dict

fastvideo/inference_args.py

@@ -32,6 +32,10 @@ class InferenceArgs:
     dit_weight: Optional[str] = None
     model_dir: Optional[str] = None
     
+    # HuggingFace specific parameters
+    trust_remote_code: bool = False
+    revision: Optional[str] = None
+    
     # Parallelism
     tp_size: int = 1
     sp_size: int = 1
@@ -126,6 +130,8 @@ class InferenceArgs:
     output_path: str = "outputs/"
     seed: int = 1024
     seeds: Optional[List[int]] = None
+    device_str: Optional[str] = None
+    device = None
 
     def __post_init__(self):
         pass
@@ -155,6 +161,20 @@ class InferenceArgs:
             type=str,
             help="Directory containing StepVideo model",
         )
+        
+        # HuggingFace specific parameters
+        parser.add_argument(
+            "--trust-remote-code",
+            action="store_true",
+            default=InferenceArgs.trust_remote_code,
+            help="Trust remote code when loading HuggingFace models",
+        )
+        parser.add_argument(
+            "--revision",
+            type=str,
+            default=InferenceArgs.revision,
+            help="The specific model version to use (can be a branch name, tag name, or commit id)",
+        )
 
         # Parallelism
         parser.add_argument(

fastvideo/inference_engine.py

@@ -72,6 +72,11 @@ class InferenceEngine:
         """
         try:
             logger.info(f"Building pipeline...")
+            # TODO(will): probably a better place to set device_str?
+            local_rank = os.environ.get("LOCAL_RANK", 0)
+            device_str = f"cuda:{local_rank}"
+            inference_args.device_str = device_str
+            inference_args.device = torch.device(device_str)
             # TODO(will): I don't really like this api.
             # it should be something closer to pipeline_cls.from_pretrained(...)
             # this way for training we can just do pipeline_cls.from_pretrained(
@@ -190,9 +195,7 @@ class InferenceEngine:
         # return
         # sp_group = get_sp_group()
         # local_rank = sp_group.rank
-        import os
-        local_rank = os.environ.get("LOCAL_RANK", 0)
-        device = torch.device(f"cuda:{local_rank}")
+        device = torch.device(inference_args.device_str)
         batch = ForwardBatch(
             prompt=prompt,
             negative_prompt=negative_prompt,

fastvideo/layers/pooler.py

@@ -1,322 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from enum import IntEnum
-from typing import List, Optional, Union
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from transformers import PretrainedConfig
-from typing_extensions import assert_never
-
-from vllm.config import PoolerConfig
-from vllm.model_executor.pooling_metadata import (PoolingMetadata,
-                                                  PoolingTensors)
-from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput
-from vllm.transformers_utils.config import (
-    get_cross_encoder_activation_function)
-
-
-class PoolingType(IntEnum):
-    """Enumeration for different types of pooling methods."""
-    LAST = 0
-    ALL = 1
-    CLS = 2
-    STEP = 3
-    MEAN = 4
-
-
-class SimplePooler(nn.Module):
-    """A layer that pools specific information from hidden states.
-
-    This layer does the following:
-    1. Extracts specific tokens or aggregates data based on pooling method.
-    2. Normalizes output if specified.
-    3. Returns structured results as `PoolerOutput`.
-
-    Attributes:
-        pooling_type: The type of pooling to use.
-        normalize: Whether to normalize the pooled data.
-    """
-
-    @staticmethod
-    def from_pooling_type(
-        pooling_type: PoolingType,
-        *,
-        normalize: bool,
-        softmax: bool,
-        step_tag_id: Optional[int] = None,
-        returned_token_ids: Optional[List[int]] = None,
-    ) -> "SimplePooler":
-        if pooling_type == PoolingType.LAST:
-            assert step_tag_id is None and returned_token_ids is None
-            return LastPool(normalize=normalize, softmax=softmax)
-        if pooling_type == PoolingType.ALL:
-            assert step_tag_id is None and returned_token_ids is None
-            return AllPool(normalize=normalize, softmax=softmax)
-        if pooling_type == PoolingType.CLS:
-            assert step_tag_id is None and returned_token_ids is None
-            return CLSPool(normalize=normalize, softmax=softmax)
-        if pooling_type == PoolingType.MEAN:
-            assert step_tag_id is None and returned_token_ids is None
-            return MeanPool(normalize=normalize, softmax=softmax)
-        if pooling_type == PoolingType.STEP:
-            return StepPool(normalize=normalize,
-                            softmax=softmax,
-                            step_tag_id=step_tag_id,
-                            returned_token_ids=returned_token_ids)
-
-        assert_never(pooling_type)
-
-    def __init__(self, *, normalize: bool, softmax: bool) -> None:
-        super().__init__()
-
-        self.head = PoolerHead(normalize=normalize, softmax=softmax)
-
-    def get_prompt_lens(
-        self,
-        hidden_states: torch.Tensor,
-        pooling_metadata: PoolingMetadata,
-    ) -> torch.Tensor:
-        return PoolingTensors.from_pooling_metadata(
-            pooling_metadata, hidden_states.device).prompt_lens
-
-    def extract_states(
-        self,
-        hidden_states: torch.Tensor,
-        pooling_metadata: PoolingMetadata,
-    ) -> Union[list[torch.Tensor], torch.Tensor]:
-        raise NotImplementedError
-
-    def build_output(self, data: torch.Tensor) -> PoolingSequenceGroupOutput:
-        return PoolingSequenceGroupOutput(data)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        pooling_metadata: PoolingMetadata,
-    ) -> PoolerOutput:
-        pooled_data = self.extract_states(hidden_states, pooling_metadata)
-        pooled_data = self.head(pooled_data)
-        pooled_outputs = [self.build_output(data) for data in pooled_data]
-        return PoolerOutput(outputs=pooled_outputs)
-
-
-class CLSPool(SimplePooler):
-
-    def extract_states(
-        self,
-        hidden_states: torch.Tensor,
-        pooling_metadata: PoolingMetadata,
-    ) -> Union[list[torch.Tensor], torch.Tensor]:
-        prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
-
-        first_token_flat_indices = torch.zeros_like(prompt_lens)
-        first_token_flat_indices[1:] += torch.cumsum(prompt_lens, dim=0)[:-1]
-        return hidden_states[first_token_flat_indices]
-
-
-class LastPool(SimplePooler):
-
-    def extract_states(
-        self,
-        hidden_states: torch.Tensor,
-        pooling_metadata: PoolingMetadata,
-    ) -> Union[list[torch.Tensor], torch.Tensor]:
-        prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
-
-        last_token_flat_indices = torch.cumsum(prompt_lens, dim=0) - 1
-        return hidden_states[last_token_flat_indices]
-
-
-class AllPool(SimplePooler):
-
-    def extract_states(
-        self,
-        hidden_states: torch.Tensor,
-        pooling_metadata: PoolingMetadata,
-    ) -> Union[list[torch.Tensor], torch.Tensor]:
-        prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
-
-        offset = 0
-        pooled_data = list[torch.Tensor]()
-        for prompt_len in prompt_lens:
-            pooled_data.append(hidden_states[offset:offset + prompt_len])
-            offset += prompt_len
-
-        return pooled_data
-
-
-class MeanPool(SimplePooler):
-
-    def extract_states(
-        self,
-        hidden_states: torch.Tensor,
-        pooling_metadata: PoolingMetadata,
-    ) -> Union[list[torch.Tensor], torch.Tensor]:
-        prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
-
-        cumsum = torch.cumsum(hidden_states, dim=0)
-        start_indices = torch.cat([
-            torch.tensor([0], device=hidden_states.device),
-            torch.cumsum(prompt_lens[:-1], dim=0)
-        ])
-        end_indices = torch.cumsum(prompt_lens, dim=0)
-        return (cumsum[end_indices - 1] - cumsum[start_indices] +
-                hidden_states[start_indices]) / prompt_lens.unsqueeze(1)
-
-
-class StepPool(SimplePooler):
-
-    def __init__(
-        self,
-        *,
-        normalize: bool,
-        softmax: bool,
-        step_tag_id: Optional[int] = None,
-        returned_token_ids: Optional[List[int]] = None,
-    ):
-        super().__init__(normalize=normalize, softmax=softmax)
-
-        self.step_tag_id = step_tag_id
-        self.returned_token_ids = returned_token_ids
-
-    def extract_states(
-        self,
-        hidden_states: torch.Tensor,
-        pooling_metadata: PoolingMetadata,
-    ) -> Union[list[torch.Tensor], torch.Tensor]:
-        prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
-
-        returned_token_ids = self.returned_token_ids
-        if returned_token_ids is not None and len(returned_token_ids) > 0:
-            hidden_states = hidden_states[:, returned_token_ids]
-
-        step_tag_id = self.step_tag_id
-
-        offset = 0
-        pooled_data = list[torch.Tensor]()
-        for prompt_len, seq_data_i in zip(prompt_lens,
-                                          pooling_metadata.seq_data.values()):
-            pooled_data_i = hidden_states[offset:offset + prompt_len]
-            if step_tag_id is not None:
-                token_ids = torch.tensor(seq_data_i.prompt_token_ids)
-                pooled_data_i = pooled_data_i[token_ids == step_tag_id]
-
-            offset += prompt_len
-            pooled_data.append(pooled_data_i)
-
-        return pooled_data
-
-
-class PoolerHead(nn.Module):
-
-    def __init__(self, *, normalize: bool, softmax: bool) -> None:
-        super().__init__()
-
-        self.normalize = normalize
-        self.softmax = softmax
-
-    def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor]):
-        if self.normalize:
-            if isinstance(pooled_data, list):
-                pooled_data = [
-                    F.normalize(data, p=2, dim=1) for data in pooled_data
-                ]
-            else:
-                pooled_data = F.normalize(pooled_data, p=2, dim=1)
-
-        if self.softmax:
-            if isinstance(pooled_data, list):
-                pooled_data = [F.softmax(data, dim=-1) for data in pooled_data]
-            else:
-                pooled_data = F.softmax(pooled_data, dim=-1)
-
-        return pooled_data
-
-
-class Pooler(nn.Module):
-
-    @classmethod
-    def from_config_with_defaults(
-        cls,
-        pooler_config: PoolerConfig,
-        pooling_type: PoolingType,
-        normalize: bool,
-        softmax: bool,
-        step_tag_id: Optional[int] = None,
-        returned_token_ids: Optional[List[int]] = None,
-    ) -> SimplePooler:
-        return SimplePooler.from_pooling_type(
-            pooling_type=PoolingType[pooler_config.pooling_type]
-            if pooler_config.pooling_type is not None else pooling_type,
-            normalize=pooler_config.normalize
-            if pooler_config.normalize is not None else normalize,
-            softmax=pooler_config.softmax
-            if pooler_config.softmax is not None else softmax,
-            step_tag_id=pooler_config.step_tag_id
-            if pooler_config.step_tag_id is not None else step_tag_id,
-            returned_token_ids=pooler_config.returned_token_ids
-            if pooler_config.returned_token_ids is not None else
-            returned_token_ids,
-        )
-
-
-class CrossEncodingPooler(nn.Module):
-    """A layer that pools specific information from hidden states.
-
-    This layer does the following:
-    1. Extracts specific tokens or aggregates data based on pooling method.
-    2. Normalizes output if specified.
-    3. Returns structured results as `PoolerOutput`.
-
-    Attributes:
-        pooling_type: The type of pooling to use.
-        normalize: Whether to normalize the pooled data.
-    """
-
-    def __init__(
-        self,
-        config: PretrainedConfig,
-        classifier: nn.Module,
-        pooler: Optional[nn.Module] = None,
-    ):
-        super().__init__()
-        self.classifier = classifier
-        self.pooler = pooler
-        self.default_activation_function = \
-            get_cross_encoder_activation_function(config)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        pooling_metadata: PoolingMetadata,
-    ) -> PoolerOutput:
-        """Pools sentence pair scores from the hidden_states."""
-
-        prompt_lens = PoolingTensors.from_pooling_metadata(
-            pooling_metadata, hidden_states.device).prompt_lens
-
-        offset = 0
-        pooled_data_lst = []
-        for prompt_len in prompt_lens:
-            pooled_data_i = hidden_states[offset:offset + prompt_len]
-
-            if self.pooler is not None:
-                final_shape_tensor = self.pooler(pooled_data_i)
-            else:
-                final_shape_tensor = self.classifier(pooled_data_i)
-
-            pooled_data_lst.append(final_shape_tensor)
-            offset += prompt_len
-
-        pooled_output = torch.stack(pooled_data_lst)
-
-        if self.pooler is not None:
-            # apply classifier once on the full batch if possible
-            pooled_output = self.classifier(pooled_output)
-
-        scores = self.default_activation_function(pooled_output).squeeze(-1)
-
-        pooled_outputs = [PoolingSequenceGroupOutput(data) for data in scores]
-        return PoolerOutput(outputs=pooled_outputs)

fastvideo/layers/resampler.py

@@ -1,269 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-# Adapted from
-# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
-# https://huggingface.co/Qwen/Qwen-7B/blob/main/modeling_qwen.py
-# https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20
-#
-# Copyright 2023 The Qwen team.
-# Copyright 2023 The vLLM team.
-# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
-#
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
-# and OPT implementations in this library. It has been modified from its
-# original forms to accommodate minor architectural differences compared
-# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Shared resampler perceiver network used in multimodal models and
-related helpers for sincos positional embeddings.
-
-Example models: Qwen (Qwen-VL), MiniCPM-V 2.0
-"""
-import math
-from functools import partial
-from typing import Callable, Optional, Tuple, Union
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-from vllm.model_executor.layers.linear import ReplicatedLinear
-from vllm.model_executor.layers.quantization import QuantizationConfig
-
-DEFAULT_LN = partial(nn.LayerNorm, eps=1e-6)
-
-
-def get_abs_pos(abs_pos: torch.Tensor, tgt_size: Union[torch.Tensor,
-                                                       int]) -> torch.Tensor:
-    # abs_pos: L, C
-    # tgt_size: (H, W)
-    # return: M, C
-    src_size = int(math.sqrt(abs_pos.size(0)))
-    dtype = abs_pos.dtype
-    if isinstance(tgt_size, int):
-        tgt_size = (tgt_size, tgt_size)
-    if (src_size == tgt_size[0] and src_size == tgt_size[1]):
-        return abs_pos
-    return (F.interpolate(
-        abs_pos.float().reshape(1, src_size, src_size, -1).permute(0, 3, 1, 2),
-        size=(tgt_size[0], tgt_size[1]),
-        mode="bicubic",
-        align_corners=False,
-    ).permute(0, 2, 3, 1).flatten(0, 2).to(dtype=dtype))
-
-
-# sin/cos positional embedding helpers are adapted from:
-# https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20
-def get_1d_sincos_pos_embed_from_grid(
-    embed_dim: int, pos: np.ndarray,
-    version: Tuple[int, int] = (2, 0)) -> torch.Tensor:
-    """
-    embed_dim: output dimension for each position
-    pos: a list of positions to be encoded: size (M,) / (H, W)
-    out: (M, D) / (H, W, D)
-    """
-    assert embed_dim % 2 == 0
-    omega = np.arange(embed_dim // 2, dtype=np.float32)
-    omega /= embed_dim / 2.0
-    omega = 1.0 / 10000**omega  # (D/2,)
-
-    if version == (2, 0):
-        pos = pos.reshape(-1)  # (M,)
-        out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
-        emb_sin = np.sin(out)  # (M, D/2)
-        emb_cos = np.cos(out)  # (M, D/2)
-        emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
-    else:
-        out = np.einsum("hw,d->hwd", pos, omega)  # (H, W, D/2), outer product
-        emb_sin = np.sin(out)  # (H, W, D/2)
-        emb_cos = np.cos(out)  # (H, W, D/2)
-        emb = np.concatenate([emb_sin, emb_cos], axis=-1)  # (H, W, D)
-    return emb
-
-
-def get_2d_sincos_pos_embed_from_grid(
-    embed_dim: int, grid: np.ndarray,
-    version: Tuple[int, int] = (2, 0)) -> torch.Tensor:
-    assert embed_dim % 2 == 0
-
-    # use half of dimensions to encode grid_h
-    emb_h = get_1d_sincos_pos_embed_from_grid(
-        embed_dim // 2, grid[0], version)  # (H*W, D/2) or (H, W, D/2)
-    emb_w = get_1d_sincos_pos_embed_from_grid(
-        embed_dim // 2, grid[1], version)  # (H*W, D/2) or (H, W, D/2)
-
-    if version == (2, 0):
-        emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
-    else:
-        emb = np.concatenate([emb_h, emb_w], axis=-1)  # (H, W, D)
-    return emb
-
-
-def get_2d_sincos_pos_embed(
-        embed_dim: int,
-        grid_size: Union[int, Tuple[int, int]],
-        cls_token: bool = False,
-        version: Tuple[int, int] = (2, 0),
-) -> torch.Tensor:
-    """
-    grid_size: int of the grid height and width
-    return:
-    pos_embed: [grid_size*grid_size, embed_dim] or
-                [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
-    """
-    if isinstance(grid_size, int):
-        grid_h_size, grid_w_size = grid_size, grid_size
-    else:
-        grid_h_size, grid_w_size = grid_size[0], grid_size[1]
-
-    grid_h = np.arange(grid_h_size, dtype=np.float32)
-    grid_w = np.arange(grid_w_size, dtype=np.float32)
-    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
-    grid = np.stack(grid, axis=0)
-    assert isinstance(grid, np.ndarray) and \
-        grid.shape == (2, grid_h_size, grid_w_size)
-
-    if version == (2, 0):
-        grid = grid.reshape([2, 1, grid_h_size, grid_w_size])
-        pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid, version)
-        if cls_token:
-            pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed],
-                                       axis=0)
-    else:
-        pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid, version)
-    return pos_embed
-
-
-class BaseResampler(nn.Module):
-    """
-    A 2D perceiver-resampler network with one cross attention layers by
-        (grid_size**2) learnable queries and 2d sincos pos_emb.
-    Outputs:
-        A tensor with the shape of (grid_size**2, embed_dim)
-    """
-
-    def __init__(self,
-                 num_queries: int,
-                 embed_dim: int,
-                 num_heads: int,
-                 kv_dim: Optional[int] = None,
-                 norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN,
-                 do_post_projection: bool = True,
-                 quant_config: Optional[QuantizationConfig] = None,
-                 prefix: str = "") -> None:
-        super().__init__()
-
-        self.num_queries = num_queries
-        self.embed_dim = embed_dim
-        self.num_heads = num_heads
-
-        self.query = nn.Parameter(torch.empty(self.num_queries, embed_dim))
-
-        if kv_dim is not None and kv_dim != embed_dim:
-            self.kv_proj = ReplicatedLinear(kv_dim,
-                                            embed_dim,
-                                            bias=False,
-                                            quant_config=quant_config,
-                                            prefix=f"{prefix}.kv_proj")
-        else:
-            # Maintain the same return value with ReplicatedLinear.forward
-            self.kv_proj = lambda *args, **kwargs: (  # type: ignore # noqa 
-                nn.Identity()(*args, **kwargs),
-                None,
-            )
-        self.attn = nn.MultiheadAttention(embed_dim, num_heads)
-        self.ln_q = norm_layer(embed_dim)
-        self.ln_kv = norm_layer(embed_dim)
-        self.do_post_projection = do_post_projection
-        self.ln_post = norm_layer(embed_dim) if do_post_projection else None
-        self.proj = nn.Parameter(
-            (embed_dim**-0.5) *
-            torch.empty(embed_dim, embed_dim)) if do_post_projection else None
-
-    def _repeat(self, query, N: int):
-        return query.unsqueeze(1).repeat(1, N, 1)
-
-
-class Resampler2(BaseResampler):
-    """Resampler-perceiver network to be used for a variety of model types,
-    e.g., Qwen-vl / Minicpmv 2.0. The main difference is the addition of the
-    do_post_projection arg, which indicates whether or not there should be
-    a post layer normalization and projector after the attention. This is
-    present in minicpmv2.0, but not qwen-vl.
-    """
-
-    def __init__(self,
-                 grid_size: int,
-                 embed_dim: int,
-                 num_heads: int,
-                 kv_dim: Optional[int] = None,
-                 norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN,
-                 adaptive: bool = False,
-                 do_post_projection: bool = True,
-                 quant_config: Optional[QuantizationConfig] = None,
-                 prefix: str = "") -> None:
-        super().__init__(grid_size**2,
-                         embed_dim,
-                         num_heads,
-                         kv_dim,
-                         norm_layer,
-                         do_post_projection=do_post_projection,
-                         quant_config=quant_config,
-                         prefix=prefix)
-
-        self.adaptive = adaptive
-        pos_embed_arr = get_2d_sincos_pos_embed(embed_dim,
-                                                grid_size,
-                                                version=(2, 0))
-
-        self.pos_embed = nn.Parameter(
-            torch.from_numpy(pos_embed_arr).requires_grad_(False))
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        tgt_sizes: Optional[torch.Tensor] = None,
-        attn_mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        if tgt_sizes is None:
-            tgt_sizes = int(math.sqrt(x.size(1)))
-        if self.adaptive:
-            pos_embed_arr = get_2d_sincos_pos_embed(self.embed_dim,
-                                                    tgt_sizes,
-                                                    version=(2, 0))
-            pos_embed = torch.from_numpy(pos_embed_arr).to(device=x.device,
-                                                           dtype=x.dtype)
-        else:
-            pos_embed = get_abs_pos(self.pos_embed,
-                                    tgt_sizes).to(device=x.device,
-                                                  dtype=x.dtype)
-
-        x, _ = self.kv_proj(x)
-        x = self.ln_kv(x).permute(1, 0, 2)
-
-        N = x.shape[1]
-        q = self.ln_q(self.query)
-        out = self.attn(
-            self._repeat(q, N) + self.pos_embed.unsqueeze(1),
-            x + pos_embed.unsqueeze(1),
-            x,
-            attn_mask=attn_mask,
-        )[0]
-        x = out.permute(1, 0, 2)
-        if self.do_post_projection:
-            x = self.ln_post(x)
-            x = x @ self.proj
-        return x

fastvideo/layers/rotary_embedding.py

@@ -500,4 +500,42 @@ def get_rotary_pos_embed(
         sp_rank=sp_rank,
         sp_world_size=sp_world_size
     )
-    return freqs_cos, freqs_sin
+    return freqs_cos, freqs_sin
+
+_ROPE_DICT: Dict[Tuple, RotaryEmbedding] = {}
+
+def get_rope(
+    head_size: int,
+    rotary_dim: int,
+    max_position: int,
+    base: int,
+    is_neox_style: bool = True,
+    rope_scaling: Optional[Dict[str, Any]] = None,
+    dtype: Optional[torch.dtype] = None,
+    partial_rotary_factor: float = 1.0,
+) -> RotaryEmbedding:
+    if dtype is None:
+        dtype = torch.get_default_dtype()
+    if rope_scaling is not None:
+        # Transforms every value that is a list into a tuple for caching calls
+        rope_scaling_tuple = {
+            k: tuple(v) if isinstance(v, list) else v
+            for k, v in rope_scaling.items()
+        }
+        rope_scaling_args = tuple(rope_scaling_tuple.items())
+    else:
+        rope_scaling_args = None
+    if partial_rotary_factor < 1.0:
+        rotary_dim = int(rotary_dim * partial_rotary_factor)
+    key = (head_size, rotary_dim, max_position, base, is_neox_style,
+           rope_scaling_args, dtype)
+    if key in _ROPE_DICT:
+        return _ROPE_DICT[key]
+
+    if rope_scaling is None:
+        rotary_emb = RotaryEmbedding(head_size, rotary_dim, max_position, base,
+                                     is_neox_style, dtype)
+    else:
+        raise ValueError(f"Unknown RoPE scaling {rope_scaling}")
+    _ROPE_DICT[key] = rotary_emb
+    return rotary_emb

fastvideo/layers/vocab_parallel_embedding.py

@@ -7,14 +7,14 @@ import torch
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter, UninitializedParameter
 
-from vllm.distributed import (divide, get_tensor_model_parallel_rank,
+from fastvideo.distributed import (divide, get_tensor_model_parallel_rank,
                               get_tensor_model_parallel_world_size,
                               tensor_model_parallel_all_reduce)
 from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding)
-from vllm.model_executor.parameter import BasevLLMParameter
-from vllm.model_executor.utils import set_weight_attrs
-from vllm.platforms import current_platform
+from fastvideo.models.parameter import BasevLLMParameter
+from fastvideo.models.utils import set_weight_attrs
+from fastvideo.platforms import current_platform
 
 DEFAULT_VOCAB_PADDING_SIZE = 64
 

fastvideo/loader/fsdp_load.py

@@ -3,7 +3,7 @@
 # Copyright 2025 The FastVideo Authors.
 
 from collections import defaultdict
-from typing import Any, Callable, Dict, Generator, List, Optional, Tuple
+from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Type
 from itertools import chain
 import torch
 from torch import nn
@@ -13,7 +13,7 @@ from torch.distributed._composable.fsdp import CPUOffloadPolicy, fully_shard
 from torch.distributed._tensor import distribute_tensor
 from torch.nn.modules.module import _IncompatibleKeys
 from vllm.model_executor.model_loader.weight_utils import safetensors_weights_iterator
-from fastvideo.models.dits.registry import ModelRegistry
+
 import contextlib
 import re
 
@@ -79,14 +79,13 @@ def get_param_names_mapping(mapping_dict: Dict[str, str]) -> Callable[[str], str
 
 # TODO(PY): add compile option
 def load_fsdp_model(
-    model_name: str,
+    model_cls: Type[nn.Module],
     init_params: Dict[str, Any],
     weight_dir_list: List[str],
     device: torch.device,
     cpu_offload: bool = False,
     default_dtype: Optional[torch.dtype] = torch.bfloat16,
 ) -> torch.nn.Module:
-    model_cls, _ = ModelRegistry.resolve_model_cls(model_name)
     with set_default_dtype(default_dtype), torch.device("meta"):
         model = model_cls(**init_params)
     device_mesh = init_device_mesh(

fastvideo/loader/loader.py

@@ -21,48 +21,22 @@ from contextlib import contextmanager
 from typing import (Any, Callable, Dict, Generator, Iterable, List, Optional,
                     Tuple, cast)
 
-import gguf
 import huggingface_hub
-import numpy as np
 import torch
-from huggingface_hub import HfApi
 from torch import nn
-from transformers import AutoModelForCausalLM
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
 
-from vllm.attention import Attention
-from vllm.config import (LoadConfig, LoadFormat, ModelConfig, ParallelConfig,
-                         VllmConfig, set_current_vllm_config)
-from vllm.distributed import (get_tensor_model_parallel_rank,
-                              get_tensor_model_parallel_world_size)
-from vllm.envs import VLLM_USE_MODELSCOPE
-from vllm.logger import init_logger
-from vllm.model_executor.layers.linear import (LinearBase,
-                                               MergedColumnParallelLinear,
-                                               QKVParallelLinear,
-                                               ReplicatedLinear,
-                                               RowParallelLinear)
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizeMethodBase)
-from vllm.model_executor.model_loader.tensorizer import (
-    TensorizerConfig, is_vllm_tensorized, load_with_tensorizer,
-    serialize_vllm_model, tensorizer_weights_iterator)
-from vllm.model_executor.model_loader.utils import (ParamMapping,
-                                                    configure_quant_config,
-                                                    get_model_architecture,
-                                                    set_default_torch_dtype)
-from vllm.model_executor.model_loader.weight_utils import (
+from fastvideo.logger import init_logger
+from fastvideo.loader.utils import set_default_torch_dtype
+from fastvideo.loader.weight_utils import (
     download_safetensors_index_file_from_hf, download_weights_from_hf,
     filter_duplicate_safetensors_files, filter_files_not_needed_for_inference,
-    get_gguf_extra_tensor_names, get_lock, gguf_quant_weights_iterator,
-    initialize_dummy_weights, np_cache_weights_iterator, pt_weights_iterator,
-    runai_safetensors_weights_iterator, safetensors_weights_iterator)
-from vllm.model_executor.utils import set_weight_attrs
-from vllm.platforms import current_platform
-from vllm.transformers_utils.s3_utils import glob as s3_glob
-from vllm.transformers_utils.utils import is_s3
-from vllm.utils import is_pin_memory_available
+    pt_weights_iterator,
+    safetensors_weights_iterator)
+from fastvideo.platforms import current_platform
+from fastvideo.models.registry import ModelRegistry
 
+from fastvideo.inference_args import InferenceArgs
 
 @contextmanager
 def device_loading_context(module: torch.nn.Module,
@@ -86,7 +60,7 @@ def device_loading_context(module: torch.nn.Module,
 
     finally:
         # Restore parameters to their original devices, ignoring new parameters
-        pin_memory = is_pin_memory_available()
+        pin_memory = True
         for name, p in module.named_parameters():
             if name in original_device_states:
                 original_device: torch.device = original_device_states[name]
@@ -111,90 +85,78 @@ logger = init_logger(__name__)
 
 
 def _initialize_model(
-    vllm_config: VllmConfig,
-    *,
-    prefix: str = "",
+    model_path: str,
+    model_config,
+    inference_args: InferenceArgs,
 ) -> nn.Module:
-    """Initialize a model with the given configurations."""
-    model_config = vllm_config.model_config
-    model_class, _ = get_model_architecture(model_config)
+    """Initialize a model with the given configurations.
+    
+    Args:
+        model_name: The name of the model architecture
+        model_path: Path to the model directory
+        inference_args: Inference arguments
+        component_name: Optional component name (e.g., "text_encoder", "vae")
+            If provided, will load the config from the component subdirectory
+    
+    Returns:
+        The initialized model
+    """
+    architectures = getattr(model_config, "architectures", [])
+    if len(architectures) == 0:
+        raise ValueError("Model config does not contain a valid model architecture")
+        logger.info("Trying to load model from diffusers format")
+        class_name = model_config.pop("_class_name")
+        architectures = [class_name]
+    if architectures is None:
+        raise ValueError("Model config does not contain a valid model architecture")
+    model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
+    logger.info(f"Loading model_cls: {model_cls}")
 
-    if vllm_config.quant_config is not None:
-        configure_quant_config(vllm_config.quant_config, model_class)
-
-    signatures = inspect.signature(model_class.__init__)
-    all_params = [param.name for param in signatures.parameters.values()]
-    if "vllm_config" in all_params and "prefix" in all_params:
-        # new-style model class
-        with set_current_vllm_config(vllm_config, check_compile=True):
-            return model_class(vllm_config=vllm_config, prefix=prefix)
-
-    msg = ("vLLM model class should accept `vllm_config` and `prefix` as "
-           "input arguments. Possibly you have an old-style model class"
-           " registered from out of tree and it is used for new vLLM version. "
-           "Check https://docs.vllm.ai/en/latest/design/arch_overview.html "
-           "for the design and update the model class accordingly.")
-    warnings.warn(msg, DeprecationWarning, stacklevel=2)
-
-    logger.warning(
-        "Trying to guess the arguments for old-style model class %s",
-        model_class,
-    )
-    # try to be compatible with old-style model class
-    kwargs = {}
-    if "prefix" in all_params:
-        kwargs["prefix"] = prefix
-    if "config" in all_params:
-        kwargs["config"] = model_config.hf_config
-    if "cache_config" in all_params:
-        kwargs["cache_config"] = vllm_config.cache_config
-    if "quant_config" in all_params:
-        kwargs["quant_config"] = vllm_config.quant_config
-    if "lora_config" in all_params:
-        kwargs["lora_config"] = vllm_config.lora_config
-    if "scheduler_config" in all_params:
-        kwargs["scheduler_config"] = vllm_config.scheduler_config
-    with set_current_vllm_config(vllm_config, check_compile=True):
-        return model_class(**kwargs)
-
-
-def _process_weights_after_loading(model: nn.Module, model_config: ModelConfig,
-                                   target_device: torch.device) -> None:
-    for _, module in model.named_modules():
-        quant_method = getattr(module, "quant_method", None)
-        if isinstance(quant_method, QuantizeMethodBase):
-            # When quant methods need to process weights after loading
-            # (for repacking, quantizing, etc), they expect parameters
-            # to be on the global target device. This scope is for the
-            # case where cpu offloading is used, where we will move the
-            # parameters onto device for processing and back off after.
-            with device_loading_context(module, target_device):
-                quant_method.process_weights_after_loading(module)
-
-    # Currently only used by MLA.
-    # NOTE: This intentionally happens after other modules so we can easily
-    # decompress the weights for MLA.
-    for _, module in model.named_modules():
-        if isinstance(module, Attention) and \
-            hasattr(module, "process_weights_after_loading"):
-            # TODO(lucas): see if there is a way to unify the signatures
-            # of process_weights_after_loading
-            module.process_weights_after_loading(model_config.dtype)
+    return model_cls(model_config)
 
+    
+    # Store the config for later use
+    if not hasattr(inference_args, 'hf_configs'):
+        inference_args.hf_configs = {}
+    
+    if component_name:
+        inference_args.hf_configs[component_name] = hf_config
+    else:
+        inference_args.hf_config = hf_config
+    
+    # Initialize the model based on the model_name
+    if model_name == "CLIPVisionModel":
+        from fastvideo.models.encoders.clip import CLIPVisionModel
+        from transformers import CLIPVisionConfig
+        
+        # If the config is not a CLIPVisionConfig, convert it
+        if not isinstance(hf_config, CLIPVisionConfig):
+            # This might need adjustment based on your specific requirements
+            clip_config = CLIPVisionConfig.from_dict(hf_config.to_dict())
+        else:
+            clip_config = hf_config
+        
+        # Initialize the model with the config
+        model = CLIPVisionModel(config=clip_config)
+    else:
+        # Add more model types as needed
+        raise ValueError(f"Unsupported model type: {model_name}")
+    
+    return model
 
 class BaseModelLoader(ABC):
     """Base class for model loaders."""
 
-    def __init__(self, load_config: LoadConfig):
-        self.load_config = load_config
+    def __init__(self, inference_args: InferenceArgs):
+        self.inference_args = inference_args
 
     @abstractmethod
-    def download_model(self, model_config: ModelConfig) -> None:
+    def download_model(self, inference_args: InferenceArgs) -> None:
         """Download a model so that it can be immediately loaded."""
         raise NotImplementedError
 
     @abstractmethod
-    def load_model(self, *, vllm_config: VllmConfig) -> nn.Module:
+    def load_model(self, *, inference_args: InferenceArgs) -> nn.Module:
         """Load a model with the given configurations."""
         raise NotImplementedError
 
@@ -224,40 +186,8 @@ class DefaultModelLoader(BaseModelLoader):
     counter_before_loading_weights: float = 0.0
     counter_after_loading_weights: float = 0.0
 
-    def __init__(self, load_config: LoadConfig):
-        super().__init__(load_config)
-        if load_config.model_loader_extra_config:
-            raise ValueError(f"Model loader extra config is not supported for "
-                             f"load format {load_config.load_format}")
-
-    def _maybe_download_from_modelscope(
-            self, model: str, revision: Optional[str]) -> Optional[str]:
-        """Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True.
-
-        Returns the path to the downloaded model, or None if the model is not
-        downloaded from ModelScope."""
-        if VLLM_USE_MODELSCOPE:
-            # download model from ModelScope hub,
-            # lazy import so that modelscope is not required for normal use.
-            # pylint: disable=C.
-            from modelscope.hub.snapshot_download import snapshot_download
-
-            if not os.path.exists(model):
-                # Use file lock to prevent multiple processes from
-                # downloading the same model weights at the same time.
-                with get_lock(model, self.load_config.download_dir):
-                    model_path = snapshot_download(
-                        model_id=model,
-                        cache_dir=self.load_config.download_dir,
-                        local_files_only=huggingface_hub.constants.
-                        HF_HUB_OFFLINE,
-                        revision=revision,
-                        ignore_file_pattern=self.load_config.ignore_patterns,
-                    )
-            else:
-                model_path = model
-            return model_path
-        return None
+    def __init__(self, inference_args: InferenceArgs):
+        super().__init__(inference_args)
 
     def _prepare_weights(
         self,
@@ -269,29 +199,32 @@ class DefaultModelLoader(BaseModelLoader):
         """Prepare weights for the model.
 
         If the model is not local, it will be downloaded."""
-        model_name_or_path = (self._maybe_download_from_modelscope(
-            model_name_or_path, revision) or model_name_or_path)
+        # model_name_or_path = (self._maybe_download_from_modelscope(
+        #     model_name_or_path, revision) or model_name_or_path)
 
         is_local = os.path.isdir(model_name_or_path)
-        load_format = self.load_config.load_format
+        assert is_local, "Model path must be a local directory"
+        # load_format = self.load_config.load_format
+        # load_format = LoadFormat.AUTO
         use_safetensors = False
         index_file = SAFE_WEIGHTS_INDEX_NAME
+        allow_patterns = ["*.safetensors", "*.bin"]
         # Some quantized models use .pt files for storing the weights.
-        if load_format == LoadFormat.AUTO:
-            allow_patterns = ["*.safetensors", "*.bin"]
-        elif load_format == LoadFormat.SAFETENSORS:
-            use_safetensors = True
-            allow_patterns = ["*.safetensors"]
-        elif load_format == LoadFormat.MISTRAL:
-            use_safetensors = True
-            allow_patterns = ["consolidated*.safetensors"]
-            index_file = "consolidated.safetensors.index.json"
-        elif load_format == LoadFormat.PT:
-            allow_patterns = ["*.pt"]
-        elif load_format == LoadFormat.NPCACHE:
-            allow_patterns = ["*.bin"]
-        else:
-            raise ValueError(f"Unknown load_format: {load_format}")
+        # if load_format == LoadFormat.AUTO:
+        #     allow_patterns = ["*.safetensors", "*.bin"]
+        # elif load_format == LoadFormat.SAFETENSORS:
+        #     use_safetensors = True
+        #     allow_patterns = ["*.safetensors"]
+        # elif load_format == LoadFormat.MISTRAL:
+        #     use_safetensors = True
+        #     allow_patterns = ["consolidated*.safetensors"]
+        #     index_file = "consolidated.safetensors.index.json"
+        # elif load_format == LoadFormat.PT:
+        #     allow_patterns = ["*.pt"]
+        # elif load_format == LoadFormat.NPCACHE:
+        #     allow_patterns = ["*.bin"]
+        # else:
+        #     raise ValueError(f"Unknown load_format: {load_format}")
 
         if fall_back_to_pt:
             allow_patterns += ["*.pt"]
@@ -350,31 +283,22 @@ class DefaultModelLoader(BaseModelLoader):
         hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
             source.model_or_path, source.revision, source.fall_back_to_pt,
             source.allow_patterns_overrides)
-        if self.load_config.load_format == LoadFormat.NPCACHE:
-            # Currently np_cache only support *.bin checkpoints
-            assert use_safetensors is False
-            weights_iterator = np_cache_weights_iterator(
-                source.model_or_path,
-                self.load_config.download_dir,
-                hf_folder,
-                hf_weights_files,
-            )
-        elif use_safetensors:
+        if use_safetensors:
             weights_iterator = safetensors_weights_iterator(hf_weights_files)
         else:
             weights_iterator = pt_weights_iterator(hf_weights_files)
 
-        if current_platform.is_tpu():
-            # In PyTorch XLA, we should call `xm.mark_step` frequently so that
-            # not too many ops are accumulated in the XLA program.
-            import torch_xla.core.xla_model as xm
+        # if current_platform.is_tpu():
+        #     # In PyTorch XLA, we should call `xm.mark_step` frequently so that
+        #     # not too many ops are accumulated in the XLA program.
+        #     import torch_xla.core.xla_model as xm
 
-            def _xla_weights_iterator(iterator: Generator):
-                for weights in iterator:
-                    yield weights
-                    xm.mark_step()
+        #     def _xla_weights_iterator(iterator: Generator):
+        #         for weights in iterator:
+        #             yield weights
+        #             xm.mark_step()
 
-            weights_iterator = _xla_weights_iterator(weights_iterator)
+        #     weights_iterator = _xla_weights_iterator(weights_iterator)
 
         if self.counter_before_loading_weights == 0.0:
             self.counter_before_loading_weights = time.perf_counter()
@@ -384,12 +308,12 @@ class DefaultModelLoader(BaseModelLoader):
 
     def _get_all_weights(
         self,
-        model_config: ModelConfig,
+        model_config: Dict[str, Any],
         model: nn.Module,
     ) -> Generator[Tuple[str, torch.Tensor], None, None]:
         primary_weights = DefaultModelLoader.Source(
             model_config.model,
-            model_config.revision,
+            "",
             prefix="",
             fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load",
                                     True),
@@ -405,21 +329,25 @@ class DefaultModelLoader(BaseModelLoader):
         for source in secondary_weights:
             yield from self._get_weights_iterator(source)
 
-    def download_model(self, model_config: ModelConfig) -> None:
-        self._prepare_weights(model_config.model,
-                              model_config.revision,
-                              fall_back_to_pt=True,
-                              allow_patterns_overrides=None)
+    def download_model(self, inference_args: InferenceArgs) -> None:
+        raise NotImplementedError
+        # self._prepare_weights(inference_args.model,
+        #                       inference_args.revision,
+        #                       fall_back_to_pt=True,
+        #                       allow_patterns_overrides=None)
 
-    def load_model(self, vllm_config: VllmConfig) -> nn.Module:
-        device_config = vllm_config.device_config
-        model_config = vllm_config.model_config
-        target_device = torch.device(device_config.device)
-        with set_default_torch_dtype(model_config.dtype):
+    def load_model(self, model_path: str,
+                   model_config: Dict[str, Any],
+                   inference_args: InferenceArgs) -> nn.Module:
+        logger.info(f"Loading model on device: {inference_args.device_str}")
+        target_device = torch.device(inference_args.device_str)
+        # TODO(will): add support for other dtypes
+        with set_default_torch_dtype(torch.bfloat16):
             with target_device:
-                model = _initialize_model(vllm_config=vllm_config)
-
+                model = _initialize_model(model_path, model_config, inference_args)
+                
             weights_to_load = {name for name, _ in model.named_parameters()}
+            model_config.model = model_path
             loaded_weights = model.load_weights(
                 self._get_all_weights(model_config, model))
             self.counter_after_loading_weights = time.perf_counter()
@@ -429,24 +357,23 @@ class DefaultModelLoader(BaseModelLoader):
                 self.counter_before_loading_weights)
             # We only enable strict check for non-quantized models
             # that have loaded weights tracking currently.
-            if model_config.quantization is None and loaded_weights is not None:
+            if loaded_weights is not None:
                 weights_not_loaded = weights_to_load - loaded_weights
                 if weights_not_loaded:
                     raise ValueError(
                         "Following weights were not initialized from "
                         f"checkpoint: {weights_not_loaded}")
 
-            _process_weights_after_loading(model, model_config, target_device)
-
+        # TODO(will): add support for training/finetune
         return model.eval()
 
 
 
 
 
-def get_model_loader(load_config: LoadConfig) -> BaseModelLoader:
+def get_model_loader(inference_args: InferenceArgs) -> BaseModelLoader:
     """Get a model loader based on the load format."""
-    if isinstance(load_config.load_format, type):
-        return load_config.load_format(load_config)
+    # if isinstance(load_config.load_format, type):
+    #     return load_config.load_format(load_config)
 
-    return DefaultModelLoader(load_config)
+    return DefaultModelLoader(inference_args)

fastvideo/loader/utils.py

@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Utilities for selecting and loading models."""
+import contextlib
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple, Type
+
+import torch
+import transformers
+from torch import nn
+from transformers.dynamic_module_utils import get_class_from_dynamic_module
+
+from vllm.config import ModelConfig, ModelImpl
+from fastvideo.logger import init_logger
+from vllm.model_executor.models import ModelRegistry
+from vllm.model_executor.models.adapters import (as_classification_model,
+                                                 as_embedding_model,
+                                                 as_reward_model)
+
+logger = init_logger(__name__)
+
+
+@contextlib.contextmanager
+def set_default_torch_dtype(dtype: torch.dtype):
+    """Sets the default torch dtype to the given dtype."""
+    old_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    yield
+    torch.set_default_dtype(old_dtype)
+
+
+def load_hf_config_from_subdir(model_path: str, component_name: str, trust_remote_code: bool = False, revision: Optional[str] = None):
+    """
+    Load a HuggingFace config from a component subdirectory.
+    
+    Args:
+        model_path: Path to the model directory
+        component_name: Name of the component subdirectory (e.g., "text_encoder", "vae")
+        trust_remote_code: Whether to trust remote code when loading the config
+        revision: Optional revision to use when loading from HuggingFace Hub
+        
+    Returns:
+        The loaded HuggingFace config
+    """
+    import os
+    from transformers import AutoConfig
+    
+    component_path = os.path.join(model_path, component_name)
+    
+    if not os.path.exists(component_path):
+        raise ValueError(f"Component directory {component_path} does not exist")
+    
+    config_path = os.path.join(component_path, "config.json")
+    if not os.path.exists(config_path):
+        raise ValueError(f"Config file {config_path} does not exist")
+    
+    logger.info(f"Loading config from {config_path}")
+    config = AutoConfig.from_pretrained(
+        component_path,
+        trust_remote_code=trust_remote_code,
+        revision=revision
+    )
+    
+    return config
+
+
+def is_transformers_impl_compatible(
+        arch: str,
+        module: Optional[transformers.PreTrainedModel] = None) -> bool:
+    mod = module or getattr(transformers, arch, None)
+    if mod is None:
+        return False
+    if hasattr(mod, "supports_backend"):
+        return mod.is_backend_compatible()
+    else:
+        return mod._supports_flex_attn
+
+
+def resolve_transformers_fallback(model_config: ModelConfig,
+                                  architectures: list[str]):
+    for i, arch in enumerate(architectures):
+        if arch == "TransformersModel":
+            continue
+        auto_map: dict[str, str] = getattr(model_config.hf_config, "auto_map",
+                                           None) or dict()
+        # Make sure that config class is always initialized before model class,
+        # otherwise the model class won't be able to access the config class,
+        # the expected auto_map should have correct order like:
+        # "auto_map": {
+        #     "AutoConfig": "<your-repo-name>--<config-name>",
+        #     "AutoModel": "<your-repo-name>--<config-name>",
+        #     "AutoModelFor<Task>": "<your-repo-name>--<config-name>",
+        # },
+        auto_modules = {
+            name: get_class_from_dynamic_module(module, model_config.model)
+            for name, module in sorted(auto_map.items(), key=lambda x: x[0])
+        }
+        custom_model_module = auto_modules.get("AutoModel")
+        # TODO(Isotr0py): Further clean up these raises.
+        # perhaps handled them in _ModelRegistry._raise_for_unsupported?
+        if model_config.model_impl == ModelImpl.TRANSFORMERS:
+            if not is_transformers_impl_compatible(arch, custom_model_module):
+                raise ValueError(
+                    f"The Transformers implementation of {arch} is not "
+                    "compatible with vLLM.")
+            architectures[i] = "TransformersModel"
+        if model_config.model_impl == ModelImpl.AUTO:
+            if not is_transformers_impl_compatible(arch, custom_model_module):
+                raise ValueError(
+                    f"{arch} has no vLLM implementation and the Transformers "
+                    "implementation is not compatible with vLLM.")
+            logger.warning(
+                "%s has no vLLM implementation, falling back to Transformers "
+                "implementation. Some features may not be supported and "
+                "performance may not be optimal.", arch)
+            architectures[i] = "TransformersModel"
+    return architectures
+
+
+def get_model_architecture(
+        model_config: ModelConfig) -> Tuple[Type[nn.Module], str]:
+    architectures = getattr(model_config.hf_config, "architectures", [])
+
+    # Special handling for quantized Mixtral.
+    # FIXME(woosuk): This is a temporary hack.
+    mixtral_supported = [
+        "fp8", "compressed-tensors", "gptq_marlin", "awq_marlin"
+    ]
+
+    if (model_config.quantization is not None
+            and model_config.quantization not in mixtral_supported
+            and "MixtralForCausalLM" in architectures):
+        architectures = ["QuantMixtralForCausalLM"]
+
+    vllm_supported_archs = ModelRegistry.get_supported_archs()
+    is_vllm_supported = any(arch in vllm_supported_archs
+                            for arch in architectures)
+    if (not is_vllm_supported
+            or model_config.model_impl == ModelImpl.TRANSFORMERS):
+        architectures = resolve_transformers_fallback(model_config,
+                                                      architectures)
+
+    model_cls, arch = ModelRegistry.resolve_model_cls(architectures)
+    if model_config.task == "embed":
+        model_cls = as_embedding_model(model_cls)
+    elif model_config.task == "classify":
+        model_cls = as_classification_model(model_cls)
+    elif model_config.task == "reward":
+        model_cls = as_reward_model(model_cls)
+
+    return model_cls, arch
+
+
+def get_architecture_class_name(model_config: ModelConfig) -> str:
+    return get_model_architecture(model_config)[1]
+
+
+@dataclass
+class ParamMapping:
+    """
+    A class to handle parameter mapping for model weight loading.
+    It creates a bidirectional mapping between packed parameters and their 
+    constituent parts.
+    """
+    packed_mapping: Dict[str, List[str]]
+    inverse_packed_mapping: Dict[str, Tuple[str,
+                                            int]] = field(default_factory=dict)
+
+    def __post_init__(self):
+        for packed_name, sub_params in self.packed_mapping.items():
+            # Skip self-contained cases (e.g., {"W_pack": ["W_pack"]})
+            if len(sub_params) == 1 and sub_params[0] == packed_name:
+                continue
+            for index, param_name in enumerate(sub_params):
+                self.inverse_packed_mapping[param_name] = (
+                    packed_name,
+                    index,
+                )
+
+    def get_sub_modules(self,
+                        module_name: str) -> Optional[Tuple[str, List[str]]]:
+        for key, value in self.packed_mapping.items():
+            if module_name.endswith(key):
+                return key, value
+        return None

fastvideo/loader/weight_utils.py

@@ -25,23 +25,23 @@ from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download
 from safetensors.torch import load_file, safe_open, save_file
 from tqdm.auto import tqdm
 
-from vllm.config import LoadConfig, ModelConfig
-from vllm.distributed import get_tensor_model_parallel_rank
-from vllm.logger import init_logger
-from vllm.model_executor.layers.quantization import (QuantizationConfig,
-                                                     get_quantization_config)
-from vllm.platforms import current_platform
-from vllm.utils import PlaceholderModule
+# from vllm.config import LoadConfig, ModelConfig
+from fastvideo.distributed import get_tensor_model_parallel_rank
+from fastvideo.logger import init_logger
+# from vllm.model_executor.layers.quantization import (QuantizationConfig,
+#                                                      get_quantization_config)
+from fastvideo.platforms import current_platform
+# from vllm.utils import PlaceholderModule
 
-try:
-    from runai_model_streamer import SafetensorsStreamer
-except (ImportError, OSError):
-    # see https://github.com/run-ai/runai-model-streamer/issues/26
-    # OSError will be raised on arm64 platform
-    runai_model_streamer = PlaceholderModule(
-        "runai_model_streamer")  # type: ignore[assignment]
-    SafetensorsStreamer = runai_model_streamer.placeholder_attr(
-        "SafetensorsStreamer")
+# try:
+#     from runai_model_streamer import SafetensorsStreamer
+# except (ImportError, OSError):
+#     # see https://github.com/run-ai/runai-model-streamer/issues/26
+#     # OSError will be raised on arm64 platform
+#     runai_model_streamer = PlaceholderModule(
+#         "runai_model_streamer")  # type: ignore[assignment]
+#     SafetensorsStreamer = runai_model_streamer.placeholder_attr(
+#         "SafetensorsStreamer")
 
 logger = init_logger(__name__)
 
@@ -137,89 +137,89 @@ def convert_bin_to_safetensor_file(
 
 
 # TODO(woosuk): Move this to other place.
-def get_quant_config(model_config: ModelConfig,
-                     load_config: LoadConfig) -> QuantizationConfig:
+# def get_quant_config(model_config: ModelConfig,
+#                      load_config: LoadConfig) -> QuantizationConfig:
 
-    quant_cls = get_quantization_config(model_config.quantization)
+#     quant_cls = get_quantization_config(model_config.quantization)
 
-    # GGUF doesn't have config file
-    if model_config.quantization == "gguf":
-        return quant_cls.from_config({})
+#     # GGUF doesn't have config file
+#     if model_config.quantization == "gguf":
+#         return quant_cls.from_config({})
 
-    # Read the quantization config from the HF model config, if available.
-    hf_quant_config = getattr(model_config.hf_config, "quantization_config",
-                              None)
-    # some vision model may keep quantization_config in their text_config
-    hf_text_config = getattr(model_config.hf_config, "text_config", None)
-    if hf_quant_config is None and hf_text_config is not None:
-        hf_quant_config = getattr(hf_text_config, "quantization_config", None)
-    if hf_quant_config is None:
-        # compressed-tensors uses a compressions_config
-        hf_quant_config = getattr(model_config.hf_config, "compression_config",
-                                  None)
-    if hf_quant_config is not None:
-        return quant_cls.from_config(hf_quant_config)
-    # In case of bitsandbytes/QLoRA, get quant config from the adapter model.
-    if model_config.quantization == "bitsandbytes":
-        if (not load_config.model_loader_extra_config
-                or "qlora_adapter_name_or_path"
-                not in load_config.model_loader_extra_config):
-            return quant_cls.from_config({"adapter_name_or_path": ""})
-        model_name_or_path = load_config.model_loader_extra_config[
-            "qlora_adapter_name_or_path"]
+#     # Read the quantization config from the HF model config, if available.
+#     hf_quant_config = getattr(model_config.hf_config, "quantization_config",
+#                               None)
+#     # some vision model may keep quantization_config in their text_config
+#     hf_text_config = getattr(model_config.hf_config, "text_config", None)
+#     if hf_quant_config is None and hf_text_config is not None:
+#         hf_quant_config = getattr(hf_text_config, "quantization_config", None)
+#     if hf_quant_config is None:
+#         # compressed-tensors uses a compressions_config
+#         hf_quant_config = getattr(model_config.hf_config, "compression_config",
+#                                   None)
+#     if hf_quant_config is not None:
+#         return quant_cls.from_config(hf_quant_config)
+#     # In case of bitsandbytes/QLoRA, get quant config from the adapter model.
+#     if model_config.quantization == "bitsandbytes":
+#         if (not load_config.model_loader_extra_config
+#                 or "qlora_adapter_name_or_path"
+#                 not in load_config.model_loader_extra_config):
+#             return quant_cls.from_config({"adapter_name_or_path": ""})
+#         model_name_or_path = load_config.model_loader_extra_config[
+#             "qlora_adapter_name_or_path"]
 
-    else:
-        model_name_or_path = model_config.model
-    is_local = os.path.isdir(model_name_or_path)
-    if not is_local:
-        # Download the config files.
-        with get_lock(model_name_or_path, load_config.download_dir):
-            hf_folder = snapshot_download(
-                model_name_or_path,
-                revision=model_config.revision,
-                allow_patterns="*.json",
-                cache_dir=load_config.download_dir,
-                local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
-                tqdm_class=DisabledTqdm,
-            )
-    else:
-        hf_folder = model_name_or_path
+#     else:
+#         model_name_or_path = model_config.model
+#     is_local = os.path.isdir(model_name_or_path)
+#     if not is_local:
+#         # Download the config files.
+#         with get_lock(model_name_or_path, load_config.download_dir):
+#             hf_folder = snapshot_download(
+#                 model_name_or_path,
+#                 revision=model_config.revision,
+#                 allow_patterns="*.json",
+#                 cache_dir=load_config.download_dir,
+#                 local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+#                 tqdm_class=DisabledTqdm,
+#             )
+#     else:
+#         hf_folder = model_name_or_path
 
-    possible_config_filenames = quant_cls.get_config_filenames()
+#     possible_config_filenames = quant_cls.get_config_filenames()
 
-    # If the quantization config is not found, use the default config.
-    if not possible_config_filenames:
-        return quant_cls()
+#     # If the quantization config is not found, use the default config.
+#     if not possible_config_filenames:
+#         return quant_cls()
 
-    config_files = glob.glob(os.path.join(hf_folder, "*.json"))
+#     config_files = glob.glob(os.path.join(hf_folder, "*.json"))
 
-    quant_config_files = [
-        f for f in config_files if any(
-            f.endswith(x) for x in possible_config_filenames)
-    ]
-    if len(quant_config_files) == 0:
-        raise ValueError(
-            f"Cannot find the config file for {model_config.quantization}")
-    if len(quant_config_files) > 1:
-        raise ValueError(
-            f"Found multiple config files for {model_config.quantization}: "
-            f"{quant_config_files}")
+#     quant_config_files = [
+#         f for f in config_files if any(
+#             f.endswith(x) for x in possible_config_filenames)
+#     ]
+#     if len(quant_config_files) == 0:
+#         raise ValueError(
+#             f"Cannot find the config file for {model_config.quantization}")
+#     if len(quant_config_files) > 1:
+#         raise ValueError(
+#             f"Found multiple config files for {model_config.quantization}: "
+#             f"{quant_config_files}")
 
-    quant_config_file = quant_config_files[0]
-    with open(quant_config_file) as f:
-        config = json.load(f)
+#     quant_config_file = quant_config_files[0]
+#     with open(quant_config_file) as f:
+#         config = json.load(f)
 
-        if model_config.quantization == "bitsandbytes":
-            config["adapter_name_or_path"] = model_name_or_path
-        elif model_config.quantization == "modelopt":
-            if config["producer"]["name"] == "modelopt":
-                return quant_cls.from_config(config)
-            else:
-                raise ValueError(
-                    f"Unsupported quantization config"
-                    f" found for {model_config.quantization} in {f}.")
+#         if model_config.quantization == "bitsandbytes":
+#             config["adapter_name_or_path"] = model_name_or_path
+#         elif model_config.quantization == "modelopt":
+#             if config["producer"]["name"] == "modelopt":
+#                 return quant_cls.from_config(config)
+#             else:
+#                 raise ValueError(
+#                     f"Unsupported quantization config"
+#                     f" found for {model_config.quantization} in {f}.")
 
-    return quant_cls.from_config(config)
+#     return quant_cls.from_config(config)
 
 
 def download_weights_from_hf(
@@ -436,21 +436,21 @@ def safetensors_weights_iterator(
                 yield name, param
 
 
-def runai_safetensors_weights_iterator(
-    hf_weights_files: List[str]
-) -> Generator[Tuple[str, torch.Tensor], None, None]:
-    """Iterate over the weights in the model safetensor files."""
-    enable_tqdm = not torch.distributed.is_initialized(
-    ) or torch.distributed.get_rank() == 0
-    with SafetensorsStreamer() as streamer:
-        for st_file in tqdm(
-                hf_weights_files,
-                desc="Loading safetensors using Runai Model Streamer",
-                disable=not enable_tqdm,
-                bar_format=_BAR_FORMAT,
-        ):
-            streamer.stream_file(st_file)
-            yield from streamer.get_tensors()
+# def runai_safetensors_weights_iterator(
+#     hf_weights_files: List[str]
+# ) -> Generator[Tuple[str, torch.Tensor], None, None]:
+#     """Iterate over the weights in the model safetensor files."""
+#     enable_tqdm = not torch.distributed.is_initialized(
+#     ) or torch.distributed.get_rank() == 0
+#     with SafetensorsStreamer() as streamer:
+#         for st_file in tqdm(
+#                 hf_weights_files,
+#                 desc="Loading safetensors using Runai Model Streamer",
+#                 disable=not enable_tqdm,
+#                 bar_format=_BAR_FORMAT,
+#         ):
+#             streamer.stream_file(st_file)
+#             yield from streamer.get_tensors()
 
 
 def pt_weights_iterator(

fastvideo/model_loader.py

@@ -1,273 +0,0 @@
-"""
-Model loader for diffusion models.
-
-This module provides a base class and implementations for loading model components
-needed by diffusion pipelines.
-"""
-
-from abc import ABC, abstractmethod
-from typing import Any, Dict, Optional, Type, Union
-import os
-import torch
-import traceback
-
-from fastvideo.inference_args import InferenceArgs
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-class ModelLoader(ABC):
-    """
-    Abstract base class for model loaders.
-    
-    A model loader is responsible for loading the components needed by a diffusion
-    pipeline, such as the UNet, VAE, text encoder, etc.
-    """
-    
-    @abstractmethod
-    def load_components(
-        self,
-        inference_args: InferenceArgs,
-    ) -> Dict[str, Any]:
-        """
-        Load the components needed by a diffusion pipeline.
-        
-        Args:
-            inference_args: The inference arguments.
-            
-        Returns:
-            A dictionary of component names to component objects.
-        """
-        pass
-
-
-class HunYuanModelLoader(ModelLoader):
-    """
-    Model loader for HunYuan models.
-    
-    This loader loads the components needed by HunYuan diffusion pipelines.
-    """
-    
-    def load_components(
-        self,
-        inference_args: InferenceArgs,
-    ) -> Dict[str, Any]:
-        """
-        Load the components needed by a HunYuan diffusion pipeline.
-        
-        Args:
-            inference_args: The inference arguments.
-            
-        Returns:
-            A dictionary of component names to component objects.
-            
-        Raises:
-            FileNotFoundError: If a required model file is not found.
-            RuntimeError: If there is an error loading a model component.
-        """
-        try:
-            from fastvideo.models.unet import UNet3DConditionModel
-            from fastvideo.models.vae import AutoencoderKLCausal3D
-            from fastvideo.models.text_encoder import TextEncoder
-            from fastvideo.models.tokenizer import Tokenizer
-            from fastvideo.models.scheduler import DDIMScheduler
-            
-            logger.info(f"Loading model components from {inference_args.model_path}")
-            
-            # Check if model path exists
-            if not os.path.exists(inference_args.model_path):
-                raise FileNotFoundError(f"Model path not found: {inference_args.model_path}")
-            
-            # Check if DiT weight exists
-            if inference_args.dit_weight and not os.path.exists(inference_args.dit_weight):
-                raise FileNotFoundError(f"DiT weight not found: {inference_args.dit_weight}")
-            
-            # Load UNet
-            try:
-                logger.info("Loading UNet...")
-                unet = UNet3DConditionModel.from_pretrained(
-                    inference_args.dit_weight,
-                    subfolder="",
-                    load_key=inference_args.load_key,
-                )
-                logger.info("UNet loaded successfully")
-            except Exception as e:
-                logger.error(f"Error loading UNet: {e}")
-                logger.error(traceback.format_exc())
-                raise RuntimeError(f"Failed to load UNet: {e}")
-            
-            # Load VAE
-            try:
-                logger.info("Loading VAE...")
-                vae = AutoencoderKLCausal3D.from_pretrained(
-                    inference_args.model_path,
-                    subfolder="vae",
-                    use_safetensors=True,
-                )
-                logger.info("VAE loaded successfully")
-            except Exception as e:
-                logger.error(f"Error loading VAE: {e}")
-                logger.error(traceback.format_exc())
-                raise RuntimeError(f"Failed to load VAE: {e}")
-            
-            # Configure VAE for sequence parallelism if needed
-            if inference_args.vae_sp:
-                try:
-                    logger.info("Configuring VAE for sequence parallelism")
-                    vae.enable_sequence_parallel()
-                    logger.info("VAE configured for sequence parallelism")
-                except Exception as e:
-                    logger.error(f"Error configuring VAE for sequence parallelism: {e}")
-                    logger.error(traceback.format_exc())
-                    logger.warning("Continuing without sequence parallelism for VAE")
-            
-            # Load text encoders
-            try:
-                logger.info("Loading primary text encoder...")
-                text_encoder = TextEncoder.from_pretrained(
-                    inference_args.model_path,
-                    subfolder="text_encoder",
-                    use_safetensors=True,
-                )
-                logger.info("Primary text encoder loaded successfully")
-            except Exception as e:
-                logger.error(f"Error loading primary text encoder: {e}")
-                logger.error(traceback.format_exc())
-                raise RuntimeError(f"Failed to load primary text encoder: {e}")
-            
-            try:
-                logger.info("Loading secondary text encoder...")
-                text_encoder_2 = TextEncoder.from_pretrained(
-                    inference_args.model_path,
-                    subfolder="text_encoder_2",
-                    use_safetensors=True,
-                )
-                logger.info("Secondary text encoder loaded successfully")
-            except Exception as e:
-                logger.error(f"Error loading secondary text encoder: {e}")
-                logger.error(traceback.format_exc())
-                raise RuntimeError(f"Failed to load secondary text encoder: {e}")
-            
-            # Load tokenizers
-            try:
-                logger.info("Loading primary tokenizer...")
-                tokenizer = Tokenizer.from_pretrained(
-                    inference_args.model_path,
-                    subfolder="tokenizer",
-                )
-                logger.info("Primary tokenizer loaded successfully")
-            except Exception as e:
-                logger.error(f"Error loading primary tokenizer: {e}")
-                logger.error(traceback.format_exc())
-                raise RuntimeError(f"Failed to load primary tokenizer: {e}")
-            
-            try:
-                logger.info("Loading secondary tokenizer...")
-                tokenizer_2 = Tokenizer.from_pretrained(
-                    inference_args.model_path,
-                    subfolder="tokenizer_2",
-                )
-                logger.info("Secondary tokenizer loaded successfully")
-            except Exception as e:
-                logger.error(f"Error loading secondary tokenizer: {e}")
-                logger.error(traceback.format_exc())
-                raise RuntimeError(f"Failed to load secondary tokenizer: {e}")
-            
-            # Load scheduler
-            try:
-                logger.info("Loading scheduler...")
-                scheduler = DDIMScheduler.from_pretrained(
-                    inference_args.model_path,
-                    subfolder="scheduler",
-                )
-                logger.info("Scheduler loaded successfully")
-            except Exception as e:
-                logger.error(f"Error loading scheduler: {e}")
-                logger.error(traceback.format_exc())
-                raise RuntimeError(f"Failed to load scheduler: {e}")
-            
-            # Move models to the appropriate device
-            device = torch.device(f"cuda:{os.environ.get('LOCAL_RANK', '0')}")
-            logger.info(f"Moving models to device: {device}")
-            
-            try:
-                unet.to(device)
-                vae.to(device)
-                text_encoder.to(device)
-                text_encoder_2.to(device)
-                logger.info("Models moved to device successfully")
-            except Exception as e:
-                logger.error(f"Error moving models to device: {e}")
-                logger.error(traceback.format_exc())
-                raise RuntimeError(f"Failed to move models to device: {e}")
-            
-            # Set precision
-            try:
-                logger.info(f"Setting UNet precision to {inference_args.precision}")
-                if inference_args.precision == "bf16":
-                    unet = unet.to(torch.bfloat16)
-                elif inference_args.precision == "fp16":
-                    unet = unet.to(torch.float16)
-                
-                logger.info(f"Setting VAE precision to {inference_args.vae_precision}")
-                if inference_args.vae_precision == "fp16":
-                    vae = vae.to(torch.float16)
-                
-                logger.info(f"Setting text encoder precision to {inference_args.text_encoder_precision}")
-                if inference_args.text_encoder_precision == "fp16":
-                    text_encoder = text_encoder.to(torch.float16)
-                
-                logger.info(f"Setting secondary text encoder precision to {inference_args.text_encoder_precision_2}")
-                if inference_args.text_encoder_precision_2 == "fp16":
-                    text_encoder_2 = text_encoder_2.to(torch.float16)
-                
-                logger.info("Precision set successfully")
-            except Exception as e:
-                logger.error(f"Error setting precision: {e}")
-                logger.error(traceback.format_exc())
-                logger.warning("Continuing with default precision")
-            
-            # Return the components
-            return {
-                "unet": unet,
-                "vae": vae,
-                "text_encoder": text_encoder,
-                "text_encoder_2": text_encoder_2,
-                "tokenizer": tokenizer,
-                "tokenizer_2": tokenizer_2,
-                "scheduler": scheduler,
-            }
-        except Exception as e:
-            logger.error(f"Unhandled error in load_components: {e}")
-            logger.error(traceback.format_exc())
-            raise
-
-
-# Factory function to get the appropriate model loader
-def get_model_loader(model_type: str) -> Type[ModelLoader]:
-    """
-    Get the appropriate model loader for the given model type.
-    
-    Args:
-        model_type: The type of model to load.
-        
-    Returns:
-        The model loader class.
-        
-    Raises:
-        ValueError: If the model type is not recognized.
-    """
-    model_loaders = {
-        "hunyuan": HunYuanModelLoader,
-        "HYVideo-T/2-cfgdistill": HunYuanModelLoader,
-    }
-    
-    if model_type not in model_loaders:
-        available_models = list(model_loaders.keys())
-        raise ValueError(
-            f"Model type '{model_type}' not recognized. "
-            f"Available model types: {available_models}"
-        )
-    
-    return model_loaders[model_type] 

fastvideo/models/__init__.py

@@ -1,11 +1,57 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from fastvideo.models.parameter import (BasevLLMParameter,
-                                           PackedvLLMParameter)
-from fastvideo.models.utils import set_random_seed
+import torch.nn as nn
+from typing import Dict, Any
+from fastvideo.loader.loader import get_model_loader
+from fastvideo.inference_args import InferenceArgs
+from fastvideo.logger import init_logger
+import os
+from transformers import PretrainedConfig
+
+logger = init_logger(__name__)
+
+def get_transformers_model(model_config: PretrainedConfig, model_path: str, inference_args: InferenceArgs) -> nn.Module:
+    """
+    Load a transformers model.
+    
+    Args:
+        architecture: The model architecture name
+        model_path: Path to the model directory
+        inference_args: Inference arguments
+        component_name: Optional component name (e.g., "text_encoder", "vae")
+            If provided, will load the config from the component subdirectory
+    
+    Returns:
+        The loaded model
+    """
+    model_loader = get_model_loader(inference_args)
+    return model_loader.load_model(model_config, model_path, inference_args)
+
+def get_diffusers_model(model_config: Dict[str, Any], model_path: str, inference_args: InferenceArgs) -> nn.Module:
+    """
+    Load a diffusers model.
+    """
+    raise NotImplementedError("Diffusers models are not supported yet.")
+
+def get_scheduler(module_path: str, architecture: str, inference_args: InferenceArgs) -> Dict:
+    """Create a scheduler based on the inference args. Can be overridden by subclasses."""
+    if hasattr(inference_args, 'denoise_type') and inference_args.denoise_type == "flow":
+        # TODO(will): add schedulers to register or create a new scheduler registry
+        # TODO(will): default to config file but allow override through
+        # inference args. Currently only uses inference args.
+        from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchDiscreteScheduler
+        return FlowMatchDiscreteScheduler(
+            shift=inference_args.flow_shift,
+            reverse=inference_args.flow_reverse,
+            solver=inference_args.flow_solver,
+        )
+    else:
+        raise ValueError(f"Invalid denoise type: {inference_args.denoise_type}")
 
 __all__ = [
     "set_random_seed",
     "BasevLLMParameter",
     "PackedvLLMParameter",
+    "get_model",
+    "get_scheduler",
 ]

fastvideo/models/dits/__init__.py

@@ -1,4 +0,0 @@
-
-from fastvideo.models.dits.hunyuanvideo import HunyuanVideoDiT
-
-__all__ = ["HunyuanVideoDiT"]

fastvideo/models/dits/hunyuanvideo.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 from typing import Optional, Tuple, List
-from fastvideo.attention.distributed_attn import DistributedAttention, LocalAttention
+from fastvideo.attention.flash_attn import DistributedAttention, LocalAttention
 from fastvideo.layers.linear import ReplicatedLinear
 from fastvideo.layers.layernorm import LayerNormScaleShift, ScaleResidual, ScaleResidualLayerNormScaleShift
 from fastvideo.layers.visual_embedding import PatchEmbed, TimestepEmbedder, ModulateProjection, unpatchify
@@ -329,7 +329,7 @@ class MMSingleStreamBlock(nn.Module):
 
 
 
-class HunyuanVideoDiT(BaseDiT):
+class HunyuanVideoTransformer3DModel(BaseDiT):
     """
     HunyuanVideo Transformer backbone adapted for distributed training.
     
@@ -799,9 +799,8 @@ class FinalLayer(nn.Module):
         )
         
     def forward(self, x, c):
-        # What the fuck HF? Why you change the scale and shift order here???
+        # What the heck HF? Why you change the scale and shift order here???
         scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1)
         x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
         x, _ = self.linear(x)
         return x
-

fastvideo/models/encoders/clip.py

@@ -5,7 +5,8 @@ from typing import Iterable, Optional, Set, Tuple, Union
 
 import torch
 import torch.nn as nn
-from transformers import CLIPVisionConfig
+from transformers import CLIPVisionConfig, CLIPTextConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling
 
 from vllm.attention.layer import MultiHeadAttention
 from fastvideo.distributed import divide, get_tensor_model_parallel_world_size
@@ -13,12 +14,15 @@ from fastvideo.layers.activation import get_act_fn
 from fastvideo.layers.linear import (ColumnParallelLinear,
                                                QKVParallelLinear,
                                                RowParallelLinear)
-from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+# TODO: support quantization
+# from vllm.model_executor.layers.quantization import QuantizationConfig
+from fastvideo.loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.interfaces import SupportsQuant
 
 from .vision import VisionEncoderInfo, resolve_visual_encoder_outputs
 
+class QuantizationConfig:
+    pass
 
 class CLIPEncoderInfo(VisionEncoderInfo[CLIPVisionConfig]):
 
@@ -88,6 +92,48 @@ class CLIPVisionEmbeddings(nn.Module):
         return embeddings
 
 
+class CLIPTextEmbeddings(nn.Module):
+
+    def __init__(self, config: CLIPTextConfig):
+        super().__init__()
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)
+        self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim)
+
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.register_buffer(
+            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+        )
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+    ) -> torch.Tensor:
+        seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]
+        max_position_embedding = self.position_embedding.weight.shape[0]
+
+        if seq_length > max_position_embedding:
+            raise ValueError(
+                f"Sequence length must be less than max_position_embeddings (got `sequence length`: "
+                f"{seq_length} and max_position_embeddings: {max_position_embedding}"
+            )
+
+        if position_ids is None:
+            position_ids = self.position_ids[:, :seq_length]
+
+        if inputs_embeds is None:
+            inputs_embeds = self.token_embedding(input_ids)
+
+        position_embeddings = self.position_embedding(position_ids)
+        embeddings = inputs_embeds + position_embeddings
+
+        return embeddings
+
+
 class CLIPAttention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
@@ -263,6 +309,190 @@ class CLIPEncoder(nn.Module):
         if return_all_hidden_states:
             return hidden_states_pool
         return hidden_states
+            
+
+
+class CLIPTextTransformer(nn.Module):
+
+    def __init__(self, 
+                 config: CLIPTextConfig,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 *,
+                 num_hidden_layers_override: Optional[int] = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.embeddings = CLIPTextEmbeddings(config)
+
+        self.encoder = CLIPEncoder(config,
+                                   quant_config=quant_config,
+                                   num_hidden_layers_override=num_hidden_layers_override,
+                                   prefix=prefix)
+
+        self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+        # For `pooled_output` computation
+        self.eos_token_id = config.eos_token_id
+
+        # For attention mask, it differs between `flash_attention_2` and other attention implementations
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        r"""
+        Returns:
+
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is None:
+            raise ValueError("You have to specify input_ids")
+
+        input_shape = input_ids.size()
+        input_ids = input_ids.view(-1, input_shape[-1])
+
+        hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)
+
+        # CLIP's text model uses causal mask, prepare it here.
+        # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324
+        causal_attention_mask = _create_4d_causal_attention_mask(
+            input_shape, hidden_states.dtype, device=hidden_states.device
+        )
+
+        # expand attention_mask
+        if attention_mask is not None and not self._use_flash_attention_2:
+            raise NotImplementedError("attention_mask is not supported for CLIPTextTransformer")
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
+
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            attention_mask=attention_mask,
+            causal_attention_mask=causal_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        last_hidden_state = encoder_outputs[0]
+        last_hidden_state = self.final_layer_norm(last_hidden_state)
+
+        if self.eos_token_id == 2:
+            # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here.
+            # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added
+            # ------------------------------------------------------------
+            # text_embeds.shape = [batch_size, sequence_length, transformer.width]
+            # take features from the eot embedding (eot_token is the highest number in each sequence)
+            # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14
+            pooled_output = last_hidden_state[
+                torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
+                input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1),
+            ]
+        else:
+            # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible)
+            pooled_output = last_hidden_state[
+                torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
+                # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`)
+                # Note: we assume each sequence (along batch dim.) contains an  `eos_token_id` (e.g. prepared by the tokenizer)
+                (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.eos_token_id)
+                .int()
+                .argmax(dim=-1),
+            ]
+
+        if not return_dict:
+            return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+class CLIPTextModel(nn.Module):
+
+    def __init__(
+        self,
+        config: CLIPTextConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+        self.text_model = CLIPTextTransformer(
+                            config=config,
+                            quant_config=quant_config,
+                            prefix=prefix)
+    
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        return self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+    def load_weights(self, weights: Iterable[Tuple[str,
+                                                  torch.Tensor]]) -> Set[str]:
+        
+        # Define mapping for stacked parameters
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: Set[str] = set()
+        for name, loaded_weight in weights:
+            # Handle q_proj, k_proj, v_proj -> qkv_proj mapping
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name in name:
+                    # Replace the weight name with the parameter name
+                    model_param_name = name.replace(weight_name, param_name)
+                    
+                    if model_param_name in params_dict:
+                        param = params_dict[model_param_name]
+                        weight_loader = param.weight_loader
+                        weight_loader(param, loaded_weight, shard_id)
+                        loaded_params.add(model_param_name)
+                    break
+            else:
+                # Use default weight loader for all other parameters
+                if name in params_dict:
+                    param = params_dict[name]
+                    weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                    weight_loader(param, loaded_weight)
+                    loaded_params.add(name)
+            
+        return loaded_params
 
 
 class CLIPVisionTransformer(nn.Module):

fastvideo/models/encoders/llama.py

@@ -0,0 +1,418 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# Adapted from
+# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
+# Copyright 2023 The vLLM team.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Inference-only LLaMA model compatible with HuggingFace weights."""
+from typing import Any, Dict, Iterable, Optional, Set, Tuple, Type, Union
+
+import torch
+from torch import nn
+from transformers import LlamaConfig
+from transformers.modeling_outputs import BaseModelOutputWithPast
+
+from fastvideo.attention import TextEncoderFlashAttention
+
+from fastvideo.distributed import get_tensor_model_parallel_world_size
+from fastvideo.layers.activation import SiluAndMul
+from fastvideo.layers.layernorm import RMSNorm
+from fastvideo.layers.linear import (MergedColumnParallelLinear,
+                                               QKVParallelLinear,
+                                               RowParallelLinear)
+# from vllm.model_executor.layers.quantization import QuantizationConfig
+
+from fastvideo.layers.rotary_embedding import get_rope
+from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from fastvideo.loader.weight_utils import (
+    default_weight_loader, maybe_remap_kv_scale_name)
+
+from .utils import (extract_layer_index)
+
+class QuantizationConfig:
+    pass
+
+
+class LlamaMLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: Optional[QuantizationConfig] = None,
+        bias: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size] * 2,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        if hidden_act != "silu":
+            raise ValueError(f"Unsupported activation: {hidden_act}. "
+                             "Only silu is supported for now.")
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        x, _ = self.gate_up_proj(x)
+        x = self.act_fn(x)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class LlamaAttention(nn.Module):
+
+    def __init__(self,
+                 config: LlamaConfig,
+                 hidden_size: int,
+                 num_heads: int,
+                 num_kv_heads: int,
+                 rope_theta: float = 10000,
+                 rope_scaling: Optional[Dict[str, Any]] = None,
+                 max_position_embeddings: int = 8192,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 bias: bool = False,
+                 bias_o_proj: bool = False,
+                 prefix: str = "") -> None:
+        super().__init__()
+        layer_idx = extract_layer_index(prefix)
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        # MistralConfig has an optional head_dim introduced by Mistral-Nemo
+        self.head_dim = getattr(config, "head_dim",
+                                self.hidden_size // self.total_num_heads)
+        # Phi models introduced a partial_rotary_factor parameter in the config
+        partial_rotary_factor = getattr(config, "partial_rotary_factor", 1)
+        self.rotary_dim = int(partial_rotary_factor * self.head_dim)
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=bias_o_proj,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        is_neox_style = True
+        is_gguf = quant_config and quant_config.get_name() == "gguf"
+        if is_gguf and config.model_type == "llama":
+            is_neox_style = False
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.rotary_dim,
+            max_position=max_position_embeddings,
+            base=rope_theta,
+            rope_scaling=rope_scaling,
+            is_neox_style=is_neox_style,
+        )
+
+        if hasattr(config, "interleaved_sliding_window"):
+            interleaved_sliding_window = config.interleaved_sliding_window
+            if isinstance(interleaved_sliding_window, int):
+                sliding_window = interleaved_sliding_window
+            elif isinstance(interleaved_sliding_window, list):
+                sw_idx = layer_idx % len(interleaved_sliding_window)
+                sliding_window = interleaved_sliding_window[sw_idx]
+            else:
+                raise ValueError(
+                    f"{type(interleaved_sliding_window)} is not supported.")
+        else:
+            sliding_window = None
+
+        self.attn = TextEncoderFlashAttention(
+            self.scaling,
+            # num_kv_heads=self.num_kv_heads,
+            # quant_config=quant_config,
+            # per_layer_sliding_window=sliding_window,
+            # prefix=f"{prefix}.attn",
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        attn_output = self.attn(q, k, v)
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class LlamaDecoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        config: LlamaConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        rope_theta = getattr(config, "rope_theta", 10000)
+        rope_scaling = getattr(config, "rope_scaling", None)
+        if rope_scaling is not None and getattr(
+                config, "original_max_position_embeddings", None):
+            rope_scaling["original_max_position_embeddings"] = (
+                config.original_max_position_embeddings)
+        max_position_embeddings = getattr(config, "max_position_embeddings",
+                                          8192)
+        # Support abacusai/Smaug-72B-v0.1 with attention_bias
+        # Support internlm/internlm-7b with bias
+        attention_bias = getattr(config, "attention_bias", False) or getattr(
+            config, "bias", False)
+        bias_o_proj = attention_bias
+        # support internlm/internlm3-8b with qkv_bias
+        if hasattr(config, 'qkv_bias'):
+            attention_bias = config.qkv_bias
+
+        self.self_attn = LlamaAttention(
+            config=config,
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=getattr(config, "num_key_value_heads",
+                                 config.num_attention_heads),
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            max_position_embeddings=max_position_embeddings,
+            quant_config=quant_config,
+            bias=attention_bias,
+            bias_o_proj=bias_o_proj,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.mlp = LlamaMLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            bias=getattr(config, "mlp_bias", False),
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size,
+                                       eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                                eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(
+                hidden_states, residual)
+        hidden_states = self.self_attn(positions=positions,
+                                       hidden_states=hidden_states)
+
+        # Fully Connected
+        hidden_states, residual = self.post_attention_layernorm(
+            hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+class LlamaModel(nn.Module):
+
+    def __init__(self,
+                 config: LlamaConfig,
+                 prefix: str = "",
+                 layer_type: Type[LlamaDecoderLayer] = LlamaDecoderLayer):
+        super().__init__()
+
+        quant_config = None
+        lora_config = None
+
+        self.config = config
+        self.quant_config = quant_config
+        lora_vocab = (lora_config.lora_extra_vocab_size *
+                      (lora_config.max_loras or 1)) if lora_config else 0
+        self.vocab_size = config.vocab_size + lora_vocab
+        self.org_vocab_size = config.vocab_size
+        
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            config.hidden_size,
+            org_num_embeddings=config.vocab_size,
+            quant_config=quant_config,
+        )
+        
+        self.layers = nn.ModuleList([
+            layer_type(
+                config=config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.layers.{i}"
+            ) for i in range(config.num_hidden_layers)
+        ])
+        
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor],
+        positions: torch.Tensor,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+    ) -> torch.Tensor:
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            hidden_states = self.get_input_embeddings(input_ids)
+        residual = None
+
+        all_hidden_states = () if output_hidden_states else None
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            hidden_states, residual = layer(positions, hidden_states, residual)
+
+        hidden_states, _ = self.norm(hidden_states, residual)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+        
+        # TODO(will): maybe unify the output format with other models and use
+        # our own class
+        output = BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            # past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            # attentions=all_self_attns,
+        )
+
+        return output
+
+    def load_weights(self, weights: Iterable[Tuple[str,
+                                                   torch.Tensor]]) -> Set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: Set[str] = set()
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+            if ("rotary_emb.cos_cached" in name
+                    or "rotary_emb.sin_cached" in name):
+                # Models trained using ColossalAI may include these tensors in
+                # the checkpoint. Skip them.
+                continue
+            if (self.quant_config is not None and
+                (scale_name := self.quant_config.get_cache_scale(name))):
+                # Loading kv cache quantization scales
+                param = params_dict[scale_name]
+                weight_loader = getattr(param, "weight_loader",
+                                        default_weight_loader)
+                loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
+                                 loaded_weight[0])
+                weight_loader(param, loaded_weight)
+                loaded_params.add(scale_name)
+                continue
+            if "scale" in name:
+                # Remapping the name of FP8 kv-scale.
+                name = maybe_remap_kv_scale_name(name, params_dict)
+                if name is None:
+                    continue
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader",
+                                        default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params

fastvideo/models/encoders/utils.py

@@ -0,0 +1,438 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import itertools
+from dataclasses import dataclass, field
+from typing import (Callable, Dict, Iterable, List, Literal, Mapping, Optional,
+                    Protocol, Set, Tuple, Union, overload)
+
+import torch
+import torch.nn as nn
+from torch.func import functional_call
+from transformers import PretrainedConfig
+
+from fastvideo.logger import init_logger
+from fastvideo.loader.weight_utils import default_weight_loader
+
+logger = init_logger(__name__)
+
+WeightsMapping = Mapping[str, Optional[str]]
+"""If a key maps to a value of `None`, the corresponding weight is ignored."""
+
+
+@dataclass
+class WeightsMapper:
+    """Maps the name of each weight if they match the following patterns."""
+
+    orig_to_new_substr: WeightsMapping = field(default_factory=dict)
+    orig_to_new_prefix: WeightsMapping = field(default_factory=dict)
+    orig_to_new_suffix: WeightsMapping = field(default_factory=dict)
+
+    def _map_name(self, key: str) -> Optional[str]:
+        for substr, new_key in self.orig_to_new_substr.items():
+            if substr in key:
+                if new_key is None:
+                    return None
+
+                key = key.replace(substr, new_key, 1)
+
+        for prefix, new_key in self.orig_to_new_prefix.items():
+            if key.startswith(prefix):
+                if new_key is None:
+                    return None
+
+                key = key.replace(prefix, new_key, 1)
+
+        for suffix, new_key in self.orig_to_new_suffix.items():
+            if key.endswith(suffix):
+                if new_key is None:
+                    return None
+
+                key = new_key.join(key.rsplit(suffix, 1))
+
+        return key
+
+    def apply(
+        self, weights: Iterable[Tuple[str, torch.Tensor]]
+    ) -> Iterable[Tuple[str, torch.Tensor]]:
+        return ((out_name, data) for name, data in weights
+                if (out_name := self._map_name(name)) is not None)
+
+
+class AutoWeightsLoader:
+    """
+    Helper class to load weights into a :class:`torch.nn.Module`. It is able
+    to automatically detect child modules and parameters while iterating over
+    the weights only once.
+
+    The weight loading logic for individual modules can be overridden
+    by defining a ``load_weights`` method.
+
+    Similarly, the weight loading logic for individual parameters can be
+    overridden by defining a ``weight_loader`` method.
+
+    Detailed weight loading information can be viewed by setting the
+    environment variable ``VLLM_LOGGING_LEVEL=DEBUG``.
+    """
+
+    def __init__(
+        self,
+        module: nn.Module,
+        *,
+        skip_prefixes: Optional[List[str]] = None,
+        ignore_unexpected_prefixes: Optional[List[str]] = None,
+    ) -> None:
+        super().__init__()
+
+        self.module = module
+        self.skip_prefixes = skip_prefixes or []
+        self.ignore_unexpected_prefixes = ignore_unexpected_prefixes or []
+
+    def _groupby_prefix(
+        self,
+        weights: Iterable[Tuple[str, torch.Tensor]],
+    ) -> Iterable[Tuple[str, Iterable[Tuple[str, torch.Tensor]]]]:
+        weights_by_parts = ((weight_name.split(".", 1), weight_data)
+                            for weight_name, weight_data in weights)
+
+        for prefix, group in itertools.groupby(weights_by_parts,
+                                               key=lambda x: x[0][0]):
+            yield (
+                prefix,
+                # Because maxsplit=1 in weight_name.split(...),
+                # the length of `parts` must either be 1 or 2
+                (("" if len(parts) == 1 else parts[1], weights_data)
+                 for parts, weights_data in group),
+            )
+
+    def _get_qualname(self, prefix: str, rest: str) -> str:
+        if prefix == "":
+            return rest
+        if rest == "":
+            return prefix
+
+        return ".".join((prefix, rest))
+
+    def _can_skip(self, qualname: str) -> bool:
+        return any(qualname.startswith(p) for p in self.skip_prefixes)
+
+    def _can_ignore_unexpected(self, qualname: str) -> bool:
+        return any(
+            qualname.startswith(p) for p in self.ignore_unexpected_prefixes)
+
+    def _load_param(
+        self,
+        base_prefix: str,
+        param: nn.Parameter,
+        weights: Iterable[Tuple[str, torch.Tensor]],
+    ) -> Iterable[str]:
+        for weight_name, weight_data in weights:
+            weight_qualname = self._get_qualname(base_prefix, weight_name)
+
+            if self._can_skip(weight_qualname):
+                logger.debug("Skipping weight %s", weight_qualname)
+
+                continue
+
+            if weight_name != "":
+                if self._can_ignore_unexpected(weight_qualname):
+                    logger.debug("Ignoring weight %s", weight_qualname)
+
+                    continue
+
+                raise ValueError(
+                    f"Attempted to load nested weight '{weight_qualname}' "
+                    f"into a single parameter '{base_prefix}'")
+
+            weight_loader = getattr(param, "weight_loader",
+                                    default_weight_loader)
+            weight_loader(param, weight_data)
+
+            logger.debug("Loaded weight %s with shape %s", weight_qualname,
+                         param.shape)
+
+            yield weight_qualname
+
+    def _load_module(
+        self,
+        base_prefix: str,
+        module: nn.Module,
+        weights: Iterable[Tuple[str, torch.Tensor]],
+    ) -> Iterable[str]:
+        if isinstance(module, PPMissingLayer):
+            return
+
+        # Avoid infinite recursion since this function is typically
+        # called inside load_weights of the module itself
+        if module != self.module:
+            module_load_weights = getattr(module, "load_weights", None)
+            if callable(module_load_weights):
+                loaded_params = module_load_weights(weights)
+                if loaded_params is None:
+                    logger.warning(
+                        "Unable to collect loaded parameters "
+                        "for module %s", module)
+                else:
+                    yield from map(
+                        lambda x: self._get_qualname(base_prefix, x),
+                        loaded_params,
+                    )
+
+        child_modules = dict(module.named_children())
+        child_params = dict(module.named_parameters(recurse=False))
+
+        for child_prefix, child_weights in self._groupby_prefix(weights):
+            prefix = self._get_qualname(base_prefix, child_prefix)
+
+            if child_prefix in child_modules:
+                if self._can_skip(prefix + "."):
+                    logger.debug("Skipping module %s", prefix)
+
+                    continue
+
+                yield from self._load_module(prefix,
+                                             child_modules[child_prefix],
+                                             child_weights)
+            elif child_prefix in child_params:
+                if self._can_skip(prefix):
+                    logger.debug("Skipping param %s", prefix)
+
+                    continue
+
+                yield from self._load_param(prefix, child_params[child_prefix],
+                                            child_weights)
+            else:
+                can_skip_module = self._can_skip(prefix + ".")
+                can_skip_param = self._can_skip(prefix)
+                if can_skip_module or can_skip_param:
+                    logger.debug("Skipping missing %s", prefix)
+
+                    continue
+
+                can_ignore_module = self._can_ignore_unexpected(prefix + ".")
+                can_ignore_param = self._can_ignore_unexpected(prefix)
+                if can_ignore_module or can_ignore_param:
+                    logger.debug("Ignoring missing %s", prefix)
+
+                    continue
+
+                msg = (f"There is no module or parameter named '{prefix}' "
+                       f"in {type(self.module).__name__}")
+                raise ValueError(msg)
+
+    def load_weights(
+        self,
+        weights: Iterable[Tuple[str, torch.Tensor]],
+        *,
+        mapper: Optional[WeightsMapper] = None,
+    ) -> Set[str]:
+        if mapper is not None:
+            weights = mapper.apply(weights)
+
+        autoloaded_weights = set(self._load_module("", self.module, weights))
+        return autoloaded_weights
+
+
+@overload
+def flatten_bn(x: torch.Tensor) -> torch.Tensor:
+    ...
+
+
+@overload
+def flatten_bn(x: List[torch.Tensor]) -> List[torch.Tensor]:
+    ...
+
+
+@overload
+def flatten_bn(
+    x: Union[List[torch.Tensor], torch.Tensor],
+    *,
+    concat: Literal[True],
+) -> torch.Tensor:
+    ...
+
+
+@overload
+def flatten_bn(
+    x: Union[List[torch.Tensor], torch.Tensor],
+    *,
+    concat: bool = False,
+) -> Union[List[torch.Tensor], torch.Tensor]:
+    ...
+
+
+def flatten_bn(
+    x: Union[List[torch.Tensor], torch.Tensor],
+    *,
+    concat: bool = False,
+) -> Union[List[torch.Tensor], torch.Tensor]:
+    """
+    Flatten the ``B`` and ``N`` dimensions of batched multimodal inputs.
+
+    The input tensor should have shape ``(B, N, ...)```.
+    """
+    if isinstance(x, torch.Tensor):
+        return x.flatten(0, 1)
+
+    if concat:
+        return torch.cat(x)
+
+    return [x_n for x_b in x for x_n in x_b]
+
+
+class LayerFn(Protocol):
+
+    def __call__(self, prefix: str) -> torch.nn.Module:
+        ...
+
+
+class PPMissingLayer(torch.nn.Identity):
+    """
+    A placeholder layer for missing layers in a pipeline parallel model.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+
+_CPU_OFFLOAD_BYTES = 0
+_CPU_OFFLOAD_MAX_BYTES = 0
+
+
+def set_cpu_offload_max_bytes(max_bytes: int) -> None:
+    global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES
+    _CPU_OFFLOAD_BYTES = 0
+    _CPU_OFFLOAD_MAX_BYTES = max_bytes
+
+
+def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module:
+    device = next(module.parameters()).device
+
+    if device == torch.device("cpu"):
+        return module
+
+    global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES
+    if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES:
+        return module
+
+    pin_memory = True
+
+    # offload parameters to CPU
+    # use pin_memory if possible, which helps cudagraph capture speed
+    offloaded_parameters = False
+    for p in module.parameters():
+        if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES:
+            # we use per-parameter offloading
+            # one module might have some parameters offloaded and some not
+            break
+
+        # `torch.empty_like` does not support `pin_memory` argument
+        cpu_data = torch.empty_strided(size=p.data.size(),
+                                       stride=p.data.stride(),
+                                       dtype=p.data.dtype,
+                                       layout=p.data.layout,
+                                       device='cpu',
+                                       pin_memory=pin_memory)
+        cpu_data.copy_(p.data)
+        p.data = cpu_data
+        _CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size()
+        offloaded_parameters = True
+
+    if offloaded_parameters:
+        original_forward = module.forward
+
+        def forward(*args, **kwargs):
+            module.forward = original_forward
+            device_state = {
+                # here we blindly call `to(device)`
+                # if the parameter is already on the device, it will be a no-op
+                k: v.to(device, non_blocking=True)
+                for k, v in module.state_dict().items()
+            }
+            output = functional_call(module,
+                                     device_state,
+                                     args=args,
+                                     kwargs=kwargs)
+            module.forward = forward
+            return output
+
+        module.forward = forward
+
+    return module
+
+
+def maybe_prefix(prefix: str, name: str) -> str:
+    """Add a prefix to a name if the prefix is non-empty.
+
+    Args:
+        prefix: The prefix to add. If empty, no prefix will be added.
+        name: The name to potentially prefix.
+
+    Returns:
+        The string "prefix.name" if prefix was non-empty, otherwise just "name".
+    """
+    return name if not prefix else f"{prefix}.{name}"
+
+
+def extract_layer_index(layer_name: str) -> int:
+    """
+    Extract the layer index from the module name.
+    Examples:
+    - "encoder.layers.0" -> 0
+    - "encoder.layers.1.self_attn" -> 1
+    - "2.self_attn" -> 2
+    - "model.encoder.layers.0.sub.1" -> ValueError
+    """
+    subnames = layer_name.split(".")
+    int_vals: List[int] = []
+    for subname in subnames:
+        try:
+            int_vals.append(int(subname))
+        except ValueError:
+            continue
+    assert len(int_vals) == 1, (f"layer name {layer_name} should"
+                                " only contain one integer")
+    return int_vals[0]
+
+
+def cast_overflow_tensors(
+    tensors: torch.Tensor,
+    offset: float = 1000,
+) -> torch.Tensor:
+    if tensors.isinf().any() or tensors.isnan().any():
+        clamp_value = torch.finfo(tensors.dtype).max - offset
+        tensors = torch.clamp(tensors, min=-clamp_value, max=clamp_value)
+    return tensors
+
+def make_layers(
+    num_hidden_layers: int,
+    layer_fn: LayerFn,
+    prefix: str,
+) -> torch.nn.ModuleList:
+    """Make a list of layers with the given layer function.
+    
+    Args:
+        num_hidden_layers: Number of layers to create.
+        layer_fn: Function that creates a layer given a prefix.
+        prefix: Prefix for layer names.
+        
+    Returns:
+        A ModuleList containing all layers.
+    """
+    modules = torch.nn.ModuleList([
+        maybe_offload_to_cpu(layer_fn(prefix=f"{prefix}.{idx}"))
+        for idx in range(num_hidden_layers)
+    ])
+    return modules
+
+
+def add_prefix(name: str, prefix: str) -> str:
+    """Add a weight path prefix to a module name.
+
+    Args:
+        name: base module name.
+        prefix: weight prefix str to added to the front of `name` concatenated with `.`.
+
+    Returns:
+        The string `prefix.name` if prefix is non-empty, otherwise just `name`.
+    """
+    return name if not prefix else f"{prefix}.{name}"

fastvideo/models/encoders/vision.py

@@ -6,11 +6,11 @@ from typing import Final, Generic, Optional, Protocol, TypeVar, Union
 import torch
 from transformers import PretrainedConfig
 
-import vllm.envs as envs
+import fastvideo.envs as envs
 from vllm.attention.selector import (backend_name_to_enum,
                                      get_global_forced_attn_backend)
-from vllm.logger import init_logger
-from vllm.platforms import _Backend, current_platform
+from fastvideo.logger import init_logger
+from fastvideo.platforms import _Backend, current_platform
 
 logger = init_logger(__name__)
 

fastvideo/models/flash_attn_no_pad.py

@@ -1,28 +0,0 @@
-from einops import rearrange
-from flash_attn import flash_attn_varlen_qkvpacked_func
-from flash_attn.bert_padding import pad_input, unpad_input
-
-
-def flash_attn_no_pad(qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None):
-    # adapted from https://github.com/Dao-AILab/flash-attention/blob/13403e81157ba37ca525890f2f0f2137edf75311/flash_attn/flash_attention.py#L27
-    batch_size = qkv.shape[0]
-    seqlen = qkv.shape[1]
-    nheads = qkv.shape[-2]
-    x = rearrange(qkv, "b s three h d -> b s (three h d)")
-    x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input(x, key_padding_mask)
-
-    x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads)
-    output_unpad = flash_attn_varlen_qkvpacked_func(
-        x_unpad,
-        cu_seqlens,
-        max_s,
-        dropout_p,
-        softmax_scale=softmax_scale,
-        causal=causal,
-    )
-    output = rearrange(
-        pad_input(rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, seqlen),
-        "b s (h d) -> b s h d",
-        h=nheads,
-    )
-    return output

fastvideo/models/hf_transformer_utils.py

@@ -0,0 +1,259 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for Huggingface Transformers."""
+
+import contextlib
+import os
+import warnings
+from pathlib import Path
+from typing import Dict, Optional, Type, Union, Any
+import json
+
+from huggingface_hub import snapshot_download
+from transformers import (
+    AutoConfig,
+    AutoProcessor,
+    AutoTokenizer,
+    PretrainedConfig,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerFast,
+)
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+
+# from fastvideo.models.configs import ChatGLMConfig, DbrxConfig, ExaoneConfig, Qwen2_5_VLConfig
+
+_CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
+    # ChatGLMConfig.model_type: ChatGLMConfig,
+    # DbrxConfig.model_type: DbrxConfig,
+    # ExaoneConfig.model_type: ExaoneConfig,
+    # Qwen2_5_VLConfig.model_type: Qwen2_5_VLConfig,
+}
+
+for name, cls in _CONFIG_REGISTRY.items():
+    with contextlib.suppress(ValueError):
+        AutoConfig.register(name, cls)
+
+
+def download_from_hf(model_path: str):
+    if os.path.exists(model_path):
+        return model_path
+
+    return snapshot_download(model_path, allow_patterns=["*.json", "*.bin", "*.model"])
+
+
+def get_hf_config(
+    model: str,
+    trust_remote_code: bool,
+    revision: Optional[str] = None,
+    model_override_args: Optional[dict] = None,
+    inference_args: Optional[dict] = None,
+    **kwargs,
+):
+    is_gguf = check_gguf_file(model)
+    if is_gguf:
+        kwargs["gguf_file"] = model
+        model = Path(model).parent
+
+    config = AutoConfig.from_pretrained(
+        model, trust_remote_code=trust_remote_code, revision=revision, **kwargs
+    )
+    if config.model_type in _CONFIG_REGISTRY:
+        config_class = _CONFIG_REGISTRY[config.model_type]
+        config = config_class.from_pretrained(model, revision=revision)
+        # NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`.
+        setattr(config, "_name_or_path", model)
+    if model_override_args:
+        config.update(model_override_args)
+
+    # Special architecture mapping check for GGUF models
+    if is_gguf:
+        if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
+            raise RuntimeError(f"Can't get gguf config for {config.model_type}.")
+        model_type = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type]
+        config.update({"architectures": [model_type]})
+
+    return config
+
+
+def get_diffusers_config(
+    model: str,
+    inference_args: Optional[dict] = None,
+) -> Dict[str, Any]:
+    """Gets a configuration for the given diffusers model.
+    
+    Args:
+        model: The model name or path.
+        inference_args: Optional inference arguments to override in the config.
+        
+    Returns:
+        The loaded configuration.
+    """
+    # Check if the model path exists
+    if os.path.exists(model):
+        config_file = os.path.join(model, "config.json")
+        if os.path.exists(config_file):
+            try:
+                # Load the config directly from the file
+                with open(config_file, "r") as f:
+                    config_dict = json.load(f)
+                
+                # TODO(will): apply any overrides from inference args
+                return config_dict
+            except Exception as e:
+                raise RuntimeError(f"Failed to load diffusers config from {config_file}: {e}")
+    else:
+        raise RuntimeError(f"Diffusers config file not found at {model}")
+
+
+# Models don't use the same configuration key for determining the maximum
+# context length.  Store them here so we can sanely check them.
+# NOTE: The ordering here is important. Some models have two of these and we
+# have a preference for which value gets used.
+CONTEXT_LENGTH_KEYS = [
+    "max_sequence_length",
+    "seq_length",
+    "max_seq_len",
+    "model_max_length",
+    "max_position_embeddings",
+]
+
+
+def get_context_length(config):
+    """Get the context length of a model from a huggingface model configs."""
+    text_config = config
+    rope_scaling = getattr(text_config, "rope_scaling", None)
+    if rope_scaling:
+        rope_scaling_factor = rope_scaling.get("factor", 1)
+        if "original_max_position_embeddings" in rope_scaling:
+            rope_scaling_factor = 1
+        if rope_scaling.get("rope_type", None) == "llama3":
+            rope_scaling_factor = 1
+    else:
+        rope_scaling_factor = 1
+
+    for key in CONTEXT_LENGTH_KEYS:
+        val = getattr(text_config, key, None)
+        if val is not None:
+            return int(rope_scaling_factor * val)
+    return 2048
+
+
+# A fast LLaMA tokenizer with the pre-processed `tokenizer.json` file.
+_FAST_LLAMA_TOKENIZER = "hf-internal-testing/llama-tokenizer"
+
+
+def get_tokenizer(
+    tokenizer_name: str,
+    *args,
+    tokenizer_mode: str = "auto",
+    trust_remote_code: bool = False,
+    tokenizer_revision: Optional[str] = None,
+    **kwargs,
+) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
+    """Gets a tokenizer for the given model name via Huggingface."""
+    if tokenizer_mode == "slow":
+        if kwargs.get("use_fast", False):
+            raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.")
+        kwargs["use_fast"] = False
+
+    is_gguf = check_gguf_file(tokenizer_name)
+    if is_gguf:
+        kwargs["gguf_file"] = tokenizer_name
+        tokenizer_name = Path(tokenizer_name).parent
+
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_name,
+            *args,
+            trust_remote_code=trust_remote_code,
+            tokenizer_revision=tokenizer_revision,
+            clean_up_tokenization_spaces=False,
+            **kwargs,
+        )
+    except TypeError as e:
+        # The LLaMA tokenizer causes a protobuf error in some environments.
+        err_msg = (
+            "Failed to load the tokenizer. If you are using a LLaMA V1 model "
+            f"consider using '{_FAST_LLAMA_TOKENIZER}' instead of the "
+            "original tokenizer."
+        )
+        raise RuntimeError(err_msg) from e
+    except ValueError as e:
+        # If the error pertains to the tokenizer class not existing or not
+        # currently being imported, suggest using the --trust-remote-code flag.
+        if not trust_remote_code and (
+            "does not exist or is not currently imported." in str(e)
+            or "requires you to execute the tokenizer file" in str(e)
+        ):
+            err_msg = (
+                "Failed to load the tokenizer. If the tokenizer is a custom "
+                "tokenizer not yet available in the HuggingFace transformers "
+                "library, consider setting `trust_remote_code=True` in LLM "
+                "or using the `--trust-remote-code` flag in the CLI."
+            )
+            raise RuntimeError(err_msg) from e
+        else:
+            raise e
+
+    if not isinstance(tokenizer, PreTrainedTokenizerFast):
+        warnings.warn(
+            "Using a slow tokenizer. This might cause a significant "
+            "slowdown. Consider using a fast tokenizer instead."
+        )
+
+    attach_additional_stop_token_ids(tokenizer)
+    return tokenizer
+
+
+def get_processor(
+    tokenizer_name: str,
+    *args,
+    tokenizer_mode: str = "auto",
+    trust_remote_code: bool = False,
+    tokenizer_revision: Optional[str] = None,
+    **kwargs,
+):
+    processor = AutoProcessor.from_pretrained(
+        tokenizer_name,
+        *args,
+        trust_remote_code=trust_remote_code,
+        tokenizer_revision=tokenizer_revision,
+        **kwargs,
+    )
+
+    attach_additional_stop_token_ids(processor.tokenizer)
+    return processor
+
+
+def attach_additional_stop_token_ids(tokenizer):
+    # Special handling for stop token <|eom_id|> generated by llama 3 tool use.
+    if "<|eom_id|>" in tokenizer.get_added_vocab():
+        tokenizer.additional_stop_token_ids = set(
+            [tokenizer.get_added_vocab()["<|eom_id|>"]]
+        )
+    else:
+        tokenizer.additional_stop_token_ids = None
+
+
+def check_gguf_file(model: Union[str, os.PathLike]) -> bool:
+    """Check if the file is a GGUF model."""
+    model = Path(model)
+    if not model.is_file():
+        return False
+    elif model.suffix == ".gguf":
+        return True
+
+    with open(model, "rb") as f:
+        header = f.read(4)
+    return header == b"GGUF"

fastvideo/models/hunyuan/modules/__init__.py

@@ -1,25 +1,25 @@
-from .models import HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer
+# from .models import HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer
 
 
-def load_model(args, in_channels, out_channels, factor_kwargs):
-    """load hunyuan video model
+# def load_model(args, in_channels, out_channels, factor_kwargs):
+#     """load hunyuan video model
 
-    Args:
-        args (dict): model args
-        in_channels (int): input channels number
-        out_channels (int): output channels number
-        factor_kwargs (dict): factor kwargs
+#     Args:
+#         args (dict): model args
+#         in_channels (int): input channels number
+#         out_channels (int): output channels number
+#         factor_kwargs (dict): factor kwargs
 
-    Returns:
-        model (nn.Module): The hunyuan video model
-    """
-    if args.model in HUNYUAN_VIDEO_CONFIG.keys():
-        model = HYVideoDiffusionTransformer(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            **HUNYUAN_VIDEO_CONFIG[args.model],
-            **factor_kwargs,
-        )
-        return model
-    else:
-        raise NotImplementedError()
+#     Returns:
+#         model (nn.Module): The hunyuan video model
+#     """
+#     if args.model in HUNYUAN_VIDEO_CONFIG.keys():
+#         model = HYVideoDiffusionTransformer(
+#             in_channels=in_channels,
+#             out_channels=out_channels,
+#             **HUNYUAN_VIDEO_CONFIG[args.model],
+#             **factor_kwargs,
+#         )
+#         return model
+#     else:
+#         raise NotImplementedError()

fastvideo/models/hunyuan_hf/modeling_hunyuan.py

@@ -1,836 +0,0 @@
-# Copyright 2024 The Hunyuan Team and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
-from diffusers.models.attention import FeedForward
-from diffusers.models.attention_processor import Attention, AttentionProcessor
-from diffusers.models.embeddings import (CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings,
-                                         get_1d_rotary_pos_embed)
-from diffusers.models.modeling_outputs import Transformer2DModelOutput
-from diffusers.models.modeling_utils import ModelMixin
-from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
-from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
-
-from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad
-from fastvideo.utils.communications import all_gather, all_to_all_4D
-from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-def shrink_head(encoder_state, dim):
-    local_heads = encoder_state.shape[dim] // nccl_info.sp_size
-    return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
-
-
-class HunyuanVideoAttnProcessor2_0:
-
-    def __init__(self):
-        if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError(
-                "HunyuanVideoAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.")
-
-    def __call__(
-        self,
-        attn: Attention,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        image_rotary_emb: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-
-        sequence_length = hidden_states.size(1)
-        encoder_sequence_length = encoder_hidden_states.size(1)
-        if attn.add_q_proj is None and encoder_hidden_states is not None:
-            hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
-
-        # 1. QKV projections
-        query = attn.to_q(hidden_states)
-        key = attn.to_k(hidden_states)
-        value = attn.to_v(hidden_states)
-
-        query = query.unflatten(2, (attn.heads, -1)).transpose(1, 2)
-        key = key.unflatten(2, (attn.heads, -1)).transpose(1, 2)
-        value = value.unflatten(2, (attn.heads, -1)).transpose(1, 2)
-
-        # 2. QK normalization
-        if attn.norm_q is not None:
-            query = attn.norm_q(query).to(value)
-        if attn.norm_k is not None:
-            key = attn.norm_k(key).to(value)
-
-        image_rotary_emb = (
-            shrink_head(image_rotary_emb[0], dim=0),
-            shrink_head(image_rotary_emb[1], dim=0),
-        )
-
-        # 3. Rotational positional embeddings applied to latent stream
-        if image_rotary_emb is not None:
-            from diffusers.models.embeddings import apply_rotary_emb
-
-            if attn.add_q_proj is None and encoder_hidden_states is not None:
-                query = torch.cat(
-                    [
-                        apply_rotary_emb(query[:, :, :-encoder_hidden_states.shape[1]], image_rotary_emb),
-                        query[:, :, -encoder_hidden_states.shape[1]:],
-                    ],
-                    dim=2,
-                )
-                key = torch.cat(
-                    [
-                        apply_rotary_emb(key[:, :, :-encoder_hidden_states.shape[1]], image_rotary_emb),
-                        key[:, :, -encoder_hidden_states.shape[1]:],
-                    ],
-                    dim=2,
-                )
-            else:
-                query = apply_rotary_emb(query, image_rotary_emb)
-                key = apply_rotary_emb(key, image_rotary_emb)
-
-        # 4. Encoder condition QKV projection and normalization
-        if attn.add_q_proj is not None and encoder_hidden_states is not None:
-            encoder_query = attn.add_q_proj(encoder_hidden_states)
-            encoder_key = attn.add_k_proj(encoder_hidden_states)
-            encoder_value = attn.add_v_proj(encoder_hidden_states)
-
-            encoder_query = encoder_query.unflatten(2, (attn.heads, -1)).transpose(1, 2)
-            encoder_key = encoder_key.unflatten(2, (attn.heads, -1)).transpose(1, 2)
-            encoder_value = encoder_value.unflatten(2, (attn.heads, -1)).transpose(1, 2)
-
-            if attn.norm_added_q is not None:
-                encoder_query = attn.norm_added_q(encoder_query).to(encoder_value)
-            if attn.norm_added_k is not None:
-                encoder_key = attn.norm_added_k(encoder_key).to(encoder_value)
-
-            query = torch.cat([query, encoder_query], dim=2)
-            key = torch.cat([key, encoder_key], dim=2)
-            value = torch.cat([value, encoder_value], dim=2)
-
-        if get_sequence_parallel_state():
-            query_img, query_txt = query[:, :, :sequence_length, :], query[:, :, sequence_length:, :]
-            key_img, key_txt = key[:, :, :sequence_length, :], key[:, :, sequence_length:, :]
-            value_img, value_txt = value[:, :, :sequence_length, :], value[:, :, sequence_length:, :]
-            query_img = all_to_all_4D(query_img, scatter_dim=1, gather_dim=2)  #
-            key_img = all_to_all_4D(key_img, scatter_dim=1, gather_dim=2)
-            value_img = all_to_all_4D(value_img, scatter_dim=1, gather_dim=2)
-
-            query_txt = shrink_head(query_txt, dim=1)
-            key_txt = shrink_head(key_txt, dim=1)
-            value_txt = shrink_head(value_txt, dim=1)
-            query = torch.cat([query_img, query_txt], dim=2)
-            key = torch.cat([key_img, key_txt], dim=2)
-            value = torch.cat([value_img, value_txt], dim=2)
-
-        query = query.unsqueeze(2)
-        key = key.unsqueeze(2)
-        value = value.unsqueeze(2)
-        qkv = torch.cat([query, key, value], dim=2)
-        qkv = qkv.transpose(1, 3)
-
-        # 5. Attention
-        attention_mask = attention_mask[:, 0, :]
-        seq_len = qkv.shape[1]
-        attn_len = attention_mask.shape[1]
-        attention_mask = F.pad(attention_mask, (seq_len - attn_len, 0), value=True)
-
-        hidden_states = flash_attn_no_pad(qkv, attention_mask, causal=False, dropout_p=0.0, softmax_scale=None)
-
-        if get_sequence_parallel_state():
-            hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
-                (sequence_length * nccl_info.sp_size, encoder_sequence_length), dim=1)
-            hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2)
-            encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous()
-            hidden_states = hidden_states.flatten(2, 3)
-            hidden_states = hidden_states.to(query.dtype)
-            encoder_hidden_states = encoder_hidden_states.flatten(2, 3)
-            encoder_hidden_states = encoder_hidden_states.to(query.dtype)
-        else:
-            hidden_states = hidden_states.flatten(2, 3)
-            hidden_states = hidden_states.to(query.dtype)
-
-            # 6. Output projection
-            if encoder_hidden_states is not None:
-                hidden_states, encoder_hidden_states = (
-                    hidden_states[:, :-encoder_hidden_states.shape[1]],
-                    hidden_states[:, -encoder_hidden_states.shape[1]:],
-                )
-
-        if encoder_hidden_states is not None:
-            if getattr(attn, "to_out", None) is not None:
-                hidden_states = attn.to_out[0](hidden_states)
-                hidden_states = attn.to_out[1](hidden_states)
-
-            if getattr(attn, "to_add_out", None) is not None:
-                encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
-
-        return hidden_states, encoder_hidden_states
-
-
-class HunyuanVideoPatchEmbed(nn.Module):
-
-    def __init__(
-        self,
-        patch_size: Union[int, Tuple[int, int, int]] = 16,
-        in_chans: int = 3,
-        embed_dim: int = 768,
-    ) -> None:
-        super().__init__()
-
-        patch_size = (patch_size, patch_size, patch_size) if isinstance(patch_size, int) else patch_size
-        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.proj(hidden_states)
-        hidden_states = hidden_states.flatten(2).transpose(1, 2)  # BCFHW -> BNC
-        return hidden_states
-
-
-class HunyuanVideoAdaNorm(nn.Module):
-
-    def __init__(self, in_features: int, out_features: Optional[int] = None) -> None:
-        super().__init__()
-
-        out_features = out_features or 2 * in_features
-        self.linear = nn.Linear(in_features, out_features)
-        self.nonlinearity = nn.SiLU()
-
-    def forward(self,
-                temb: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        temb = self.linear(self.nonlinearity(temb))
-        gate_msa, gate_mlp = temb.chunk(2, dim=1)
-        gate_msa, gate_mlp = gate_msa.unsqueeze(1), gate_mlp.unsqueeze(1)
-        return gate_msa, gate_mlp
-
-
-class HunyuanVideoIndividualTokenRefinerBlock(nn.Module):
-
-    def __init__(
-        self,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        mlp_width_ratio: str = 4.0,
-        mlp_drop_rate: float = 0.0,
-        attention_bias: bool = True,
-    ) -> None:
-        super().__init__()
-
-        hidden_size = num_attention_heads * attention_head_dim
-
-        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
-        self.attn = Attention(
-            query_dim=hidden_size,
-            cross_attention_dim=None,
-            heads=num_attention_heads,
-            dim_head=attention_head_dim,
-            bias=attention_bias,
-        )
-
-        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)
-        self.ff = FeedForward(hidden_size, mult=mlp_width_ratio, activation_fn="linear-silu", dropout=mlp_drop_rate)
-
-        self.norm_out = HunyuanVideoAdaNorm(hidden_size, 2 * hidden_size)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        norm_hidden_states = self.norm1(hidden_states)
-
-        attn_output = self.attn(
-            hidden_states=norm_hidden_states,
-            encoder_hidden_states=None,
-            attention_mask=attention_mask,
-        )
-
-        gate_msa, gate_mlp = self.norm_out(temb)
-        hidden_states = hidden_states + attn_output * gate_msa
-
-        ff_output = self.ff(self.norm2(hidden_states))
-        hidden_states = hidden_states + ff_output * gate_mlp
-
-        return hidden_states
-
-
-class HunyuanVideoIndividualTokenRefiner(nn.Module):
-
-    def __init__(
-        self,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        num_layers: int,
-        mlp_width_ratio: float = 4.0,
-        mlp_drop_rate: float = 0.0,
-        attention_bias: bool = True,
-    ) -> None:
-        super().__init__()
-
-        self.refiner_blocks = nn.ModuleList([
-            HunyuanVideoIndividualTokenRefinerBlock(
-                num_attention_heads=num_attention_heads,
-                attention_head_dim=attention_head_dim,
-                mlp_width_ratio=mlp_width_ratio,
-                mlp_drop_rate=mlp_drop_rate,
-                attention_bias=attention_bias,
-            ) for _ in range(num_layers)
-        ])
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-    ) -> None:
-        self_attn_mask = None
-        if attention_mask is not None:
-            batch_size = attention_mask.shape[0]
-            seq_len = attention_mask.shape[1]
-            attention_mask = attention_mask.to(hidden_states.device).bool()
-            self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1)
-            self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
-            self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
-            self_attn_mask[:, :, :, 0] = True
-
-        for block in self.refiner_blocks:
-            hidden_states = block(hidden_states, temb, self_attn_mask)
-
-        return hidden_states
-
-
-class HunyuanVideoTokenRefiner(nn.Module):
-
-    def __init__(
-        self,
-        in_channels: int,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        num_layers: int,
-        mlp_ratio: float = 4.0,
-        mlp_drop_rate: float = 0.0,
-        attention_bias: bool = True,
-    ) -> None:
-        super().__init__()
-
-        hidden_size = num_attention_heads * attention_head_dim
-
-        self.time_text_embed = CombinedTimestepTextProjEmbeddings(embedding_dim=hidden_size,
-                                                                  pooled_projection_dim=in_channels)
-        self.proj_in = nn.Linear(in_channels, hidden_size, bias=True)
-        self.token_refiner = HunyuanVideoIndividualTokenRefiner(
-            num_attention_heads=num_attention_heads,
-            attention_head_dim=attention_head_dim,
-            num_layers=num_layers,
-            mlp_width_ratio=mlp_ratio,
-            mlp_drop_rate=mlp_drop_rate,
-            attention_bias=attention_bias,
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        timestep: torch.LongTensor,
-        attention_mask: Optional[torch.LongTensor] = None,
-    ) -> torch.Tensor:
-        if attention_mask is None:
-            pooled_projections = hidden_states.mean(dim=1)
-        else:
-            original_dtype = hidden_states.dtype
-            mask_float = attention_mask.float().unsqueeze(-1)
-            pooled_projections = (hidden_states * mask_float).sum(dim=1) / mask_float.sum(dim=1)
-            pooled_projections = pooled_projections.to(original_dtype)
-
-        temb = self.time_text_embed(timestep, pooled_projections)
-        hidden_states = self.proj_in(hidden_states)
-        hidden_states = self.token_refiner(hidden_states, temb, attention_mask)
-
-        return hidden_states
-
-
-class HunyuanVideoRotaryPosEmbed(nn.Module):
-
-    def __init__(self, patch_size: int, patch_size_t: int, rope_dim: List[int], theta: float = 256.0) -> None:
-        super().__init__()
-
-        self.patch_size = patch_size
-        self.patch_size_t = patch_size_t
-        self.rope_dim = rope_dim
-        self.theta = theta
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        batch_size, num_channels, num_frames, height, width = hidden_states.shape
-        rope_sizes = [
-            num_frames * nccl_info.sp_size // self.patch_size_t, height // self.patch_size, width // self.patch_size
-        ]
-
-        axes_grids = []
-        for i in range(3):
-            # Note: The following line diverges from original behaviour. We create the grid on the device, whereas
-            # original implementation creates it on CPU and then moves it to device. This results in numerical
-            # differences in layerwise debugging outputs, but visually it is the same.
-            grid = torch.arange(0, rope_sizes[i], device=hidden_states.device, dtype=torch.float32)
-            axes_grids.append(grid)
-        grid = torch.meshgrid(*axes_grids, indexing="ij")  # [W, H, T]
-        grid = torch.stack(grid, dim=0)  # [3, W, H, T]
-
-        freqs = []
-        for i in range(3):
-            freq = get_1d_rotary_pos_embed(self.rope_dim[i], grid[i].reshape(-1), self.theta, use_real=True)
-            freqs.append(freq)
-
-        freqs_cos = torch.cat([f[0] for f in freqs], dim=1)  # (W * H * T, D / 2)
-        freqs_sin = torch.cat([f[1] for f in freqs], dim=1)  # (W * H * T, D / 2)
-        return freqs_cos, freqs_sin
-
-
-class HunyuanVideoSingleTransformerBlock(nn.Module):
-
-    def __init__(
-        self,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        mlp_ratio: float = 4.0,
-        qk_norm: str = "rms_norm",
-    ) -> None:
-        super().__init__()
-
-        hidden_size = num_attention_heads * attention_head_dim
-        mlp_dim = int(hidden_size * mlp_ratio)
-
-        self.attn = Attention(
-            query_dim=hidden_size,
-            cross_attention_dim=None,
-            dim_head=attention_head_dim,
-            heads=num_attention_heads,
-            out_dim=hidden_size,
-            bias=True,
-            processor=HunyuanVideoAttnProcessor2_0(),
-            qk_norm=qk_norm,
-            eps=1e-6,
-            pre_only=True,
-        )
-
-        self.norm = AdaLayerNormZeroSingle(hidden_size, norm_type="layer_norm")
-        self.proj_mlp = nn.Linear(hidden_size, mlp_dim)
-        self.act_mlp = nn.GELU(approximate="tanh")
-        self.proj_out = nn.Linear(hidden_size + mlp_dim, hidden_size)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-    ) -> torch.Tensor:
-        text_seq_length = encoder_hidden_states.shape[1]
-        hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
-
-        residual = hidden_states
-
-        # 1. Input normalization
-        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
-        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
-
-        norm_hidden_states, norm_encoder_hidden_states = (
-            norm_hidden_states[:, :-text_seq_length, :],
-            norm_hidden_states[:, -text_seq_length:, :],
-        )
-
-        # 2. Attention
-        attn_output, context_attn_output = self.attn(
-            hidden_states=norm_hidden_states,
-            encoder_hidden_states=norm_encoder_hidden_states,
-            attention_mask=attention_mask,
-            image_rotary_emb=image_rotary_emb,
-        )
-        attn_output = torch.cat([attn_output, context_attn_output], dim=1)
-
-        # 3. Modulation and residual connection
-        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
-        hidden_states = gate.unsqueeze(1) * self.proj_out(hidden_states)
-        hidden_states = hidden_states + residual
-
-        hidden_states, encoder_hidden_states = (
-            hidden_states[:, :-text_seq_length, :],
-            hidden_states[:, -text_seq_length:, :],
-        )
-        return hidden_states, encoder_hidden_states
-
-
-class HunyuanVideoTransformerBlock(nn.Module):
-
-    def __init__(
-        self,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        mlp_ratio: float,
-        qk_norm: str = "rms_norm",
-    ) -> None:
-        super().__init__()
-
-        hidden_size = num_attention_heads * attention_head_dim
-
-        self.norm1 = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
-        self.norm1_context = AdaLayerNormZero(hidden_size, norm_type="layer_norm")
-
-        self.attn = Attention(
-            query_dim=hidden_size,
-            cross_attention_dim=None,
-            added_kv_proj_dim=hidden_size,
-            dim_head=attention_head_dim,
-            heads=num_attention_heads,
-            out_dim=hidden_size,
-            context_pre_only=False,
-            bias=True,
-            processor=HunyuanVideoAttnProcessor2_0(),
-            qk_norm=qk_norm,
-            eps=1e-6,
-        )
-
-        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.ff = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
-
-        self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.ff_context = FeedForward(hidden_size, mult=mlp_ratio, activation_fn="gelu-approximate")
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        # 1. Input normalization
-        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
-        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
-            encoder_hidden_states, emb=temb)
-
-        # 2. Joint attention
-        attn_output, context_attn_output = self.attn(
-            hidden_states=norm_hidden_states,
-            encoder_hidden_states=norm_encoder_hidden_states,
-            attention_mask=attention_mask,
-            image_rotary_emb=freqs_cis,
-        )
-
-        # 3. Modulation and residual connection
-        hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1)
-        encoder_hidden_states = encoder_hidden_states + context_attn_output * c_gate_msa.unsqueeze(1)
-
-        norm_hidden_states = self.norm2(hidden_states)
-        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
-
-        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
-        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
-
-        # 4. Feed-forward
-        ff_output = self.ff(norm_hidden_states)
-        context_ff_output = self.ff_context(norm_encoder_hidden_states)
-
-        hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
-        encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
-
-        return hidden_states, encoder_hidden_states
-
-
-class HunyuanVideoTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
-    r"""
-    A Transformer model for video-like data used in [HunyuanVideo](https://huggingface.co/tencent/HunyuanVideo).
-
-    Args:
-        in_channels (`int`, defaults to `16`):
-            The number of channels in the input.
-        out_channels (`int`, defaults to `16`):
-            The number of channels in the output.
-        num_attention_heads (`int`, defaults to `24`):
-            The number of heads to use for multi-head attention.
-        attention_head_dim (`int`, defaults to `128`):
-            The number of channels in each head.
-        num_layers (`int`, defaults to `20`):
-            The number of layers of dual-stream blocks to use.
-        num_single_layers (`int`, defaults to `40`):
-            The number of layers of single-stream blocks to use.
-        num_refiner_layers (`int`, defaults to `2`):
-            The number of layers of refiner blocks to use.
-        mlp_ratio (`float`, defaults to `4.0`):
-            The ratio of the hidden layer size to the input size in the feedforward network.
-        patch_size (`int`, defaults to `2`):
-            The size of the spatial patches to use in the patch embedding layer.
-        patch_size_t (`int`, defaults to `1`):
-            The size of the tmeporal patches to use in the patch embedding layer.
-        qk_norm (`str`, defaults to `rms_norm`):
-            The normalization to use for the query and key projections in the attention layers.
-        guidance_embeds (`bool`, defaults to `True`):
-            Whether to use guidance embeddings in the model.
-        text_embed_dim (`int`, defaults to `4096`):
-            Input dimension of text embeddings from the text encoder.
-        pooled_projection_dim (`int`, defaults to `768`):
-            The dimension of the pooled projection of the text embeddings.
-        rope_theta (`float`, defaults to `256.0`):
-            The value of theta to use in the RoPE layer.
-        rope_axes_dim (`Tuple[int]`, defaults to `(16, 56, 56)`):
-            The dimensions of the axes to use in the RoPE layer.
-    """
-
-    _supports_gradient_checkpointing = True
-
-    @register_to_config
-    def __init__(
-            self,
-            in_channels: int = 16,
-            out_channels: int = 16,
-            num_attention_heads: int = 24,
-            attention_head_dim: int = 128,
-            num_layers: int = 20,
-            num_single_layers: int = 40,
-            num_refiner_layers: int = 2,
-            mlp_ratio: float = 4.0,
-            patch_size: int = 2,
-            patch_size_t: int = 1,
-            qk_norm: str = "rms_norm",
-            guidance_embeds: bool = True,
-            text_embed_dim: int = 4096,
-            pooled_projection_dim: int = 768,
-            rope_theta: float = 256.0,
-            rope_axes_dim: Tuple[int] = (16, 56, 56),
-    ) -> None:
-        super().__init__()
-
-        inner_dim = num_attention_heads * attention_head_dim
-        out_channels = out_channels or in_channels
-
-        # 1. Latent and condition embedders
-        self.x_embedder = HunyuanVideoPatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim)
-        self.context_embedder = HunyuanVideoTokenRefiner(text_embed_dim,
-                                                         num_attention_heads,
-                                                         attention_head_dim,
-                                                         num_layers=num_refiner_layers)
-        self.time_text_embed = CombinedTimestepGuidanceTextProjEmbeddings(inner_dim, pooled_projection_dim)
-
-        # 2. RoPE
-        self.rope = HunyuanVideoRotaryPosEmbed(patch_size, patch_size_t, rope_axes_dim, rope_theta)
-
-        # 3. Dual stream transformer blocks
-        self.transformer_blocks = nn.ModuleList([
-            HunyuanVideoTransformerBlock(num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm)
-            for _ in range(num_layers)
-        ])
-
-        # 4. Single stream transformer blocks
-        self.single_transformer_blocks = nn.ModuleList([
-            HunyuanVideoSingleTransformerBlock(num_attention_heads,
-                                               attention_head_dim,
-                                               mlp_ratio=mlp_ratio,
-                                               qk_norm=qk_norm) for _ in range(num_single_layers)
-        ])
-
-        # 5. Output projection
-        self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6)
-        self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * out_channels)
-
-        self.gradient_checkpointing = False
-
-    @property
-    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
-    def attn_processors(self) -> Dict[str, AttentionProcessor]:
-        r"""
-        Returns:
-            `dict` of attention processors: A dictionary containing all attention processors used in the model with
-            indexed by its weight name.
-        """
-        # set recursively
-        processors = {}
-
-        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
-            if hasattr(module, "get_processor"):
-                processors[f"{name}.processor"] = module.get_processor()
-
-            for sub_name, child in module.named_children():
-                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
-
-            return processors
-
-        for name, module in self.named_children():
-            fn_recursive_add_processors(name, module, processors)
-
-        return processors
-
-    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
-    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
-        r"""
-        Sets the attention processor to use to compute attention.
-
-        Parameters:
-            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
-                The instantiated processor class or a dictionary of processor classes that will be set as the processor
-                for **all** `Attention` layers.
-
-                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
-                processor. This is strongly recommended when setting trainable attention processors.
-
-        """
-        count = len(self.attn_processors.keys())
-
-        if isinstance(processor, dict) and len(processor) != count:
-            raise ValueError(
-                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
-                f" number of attention layers: {count}. Please make sure to pass {count} processor classes.")
-
-        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
-            if hasattr(module, "set_processor"):
-                if not isinstance(processor, dict):
-                    module.set_processor(processor)
-                else:
-                    module.set_processor(processor.pop(f"{name}.processor"))
-
-            for sub_name, child in module.named_children():
-                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
-
-        for name, module in self.named_children():
-            fn_recursive_attn_processor(name, module, processor)
-
-    def _set_gradient_checkpointing(self, module, value=False):
-        if hasattr(module, "gradient_checkpointing"):
-            module.gradient_checkpointing = value
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        timestep: torch.LongTensor,
-        encoder_attention_mask: torch.Tensor,
-        guidance: torch.Tensor = None,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        return_dict: bool = True,
-    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
-        if guidance is None:
-            guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16)
-
-        if attention_kwargs is not None:
-            attention_kwargs = attention_kwargs.copy()
-            lora_scale = attention_kwargs.pop("scale", 1.0)
-        else:
-            lora_scale = 1.0
-
-        if USE_PEFT_BACKEND:
-            # weight the lora layers by setting `lora_scale` for each PEFT layer
-            scale_lora_layers(self, lora_scale)
-        else:
-            if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
-                logger.warning("Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.")
-
-        batch_size, num_channels, num_frames, height, width = hidden_states.shape
-        p, p_t = self.config.patch_size, self.config.patch_size_t
-        post_patch_num_frames = num_frames // p_t
-        post_patch_height = height // p
-        post_patch_width = width // p
-
-        pooled_projections = encoder_hidden_states[:, 0, :self.config.pooled_projection_dim]
-        encoder_hidden_states = encoder_hidden_states[:, 1:]
-
-        # 1. RoPE
-        image_rotary_emb = self.rope(hidden_states)
-
-        # 2. Conditional embeddings
-        temb = self.time_text_embed(timestep, guidance, pooled_projections)
-        hidden_states = self.x_embedder(hidden_states)
-        encoder_hidden_states = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask)
-
-        # 3. Attention mask preparation
-        latent_sequence_length = hidden_states.shape[1]
-        condition_sequence_length = encoder_hidden_states.shape[1]
-        sequence_length = latent_sequence_length + condition_sequence_length
-        attention_mask = torch.zeros(batch_size,
-                                     sequence_length,
-                                     sequence_length,
-                                     device=hidden_states.device,
-                                     dtype=torch.bool)  # [B, N, N]
-
-        effective_condition_sequence_length = encoder_attention_mask.sum(dim=1, dtype=torch.int)
-        effective_sequence_length = latent_sequence_length + effective_condition_sequence_length
-
-        for i in range(batch_size):
-            attention_mask[i, :effective_sequence_length[i], :effective_sequence_length[i]] = True
-
-        # 4. Transformer blocks
-        if torch.is_grad_enabled() and self.gradient_checkpointing:
-
-            def create_custom_forward(module, return_dict=None):
-
-                def custom_forward(*inputs):
-                    if return_dict is not None:
-                        return module(*inputs, return_dict=return_dict)
-                    else:
-                        return module(*inputs)
-
-                return custom_forward
-
-            ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
-
-            for block in self.transformer_blocks:
-                hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(block),
-                    hidden_states,
-                    encoder_hidden_states,
-                    temb,
-                    attention_mask,
-                    image_rotary_emb,
-                    **ckpt_kwargs,
-                )
-
-            for block in self.single_transformer_blocks:
-                hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(block),
-                    hidden_states,
-                    encoder_hidden_states,
-                    temb,
-                    attention_mask,
-                    image_rotary_emb,
-                    **ckpt_kwargs,
-                )
-
-        else:
-            for block in self.transformer_blocks:
-                hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, attention_mask,
-                                                             image_rotary_emb)
-
-            for block in self.single_transformer_blocks:
-                hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, attention_mask,
-                                                             image_rotary_emb)
-
-        # 5. Output projection
-        hidden_states = self.norm_out(hidden_states, temb)
-        hidden_states = self.proj_out(hidden_states)
-
-        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width,
-                                              -1, p_t, p, p)
-        hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7)
-        hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
-
-        if USE_PEFT_BACKEND:
-            # remove `lora_scale` from each PEFT layer
-            unscale_lora_layers(self, lora_scale)
-
-        if not return_dict:
-            return (hidden_states, )
-
-        return Transformer2DModelOutput(sample=hidden_states)

fastvideo/models/hunyuan_hf/pipeline_hunyuan.py

@@ -1,691 +0,0 @@
-# Copyright 2024 The HunyuanVideo Team and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
-from diffusers.loaders import HunyuanVideoLoraLoaderMixin
-from diffusers.models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel
-from diffusers.pipelines.hunyuan_video.pipeline_output import HunyuanVideoPipelineOutput
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
-from diffusers.utils import logging, replace_example_docstring
-from diffusers.utils.torch_utils import randn_tensor
-from diffusers.video_processor import VideoProcessor
-from einops import rearrange
-from transformers import CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizerFast
-
-from fastvideo.utils.communications import all_gather
-from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```python
-        >>> import torch
-        >>> from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
-        >>> from diffusers.utils import export_to_video
-
-        >>> model_id = "tencent/HunyuanVideo"
-        >>> transformer = HunyuanVideoTransformer3DModel.from_pretrained(
-        ...     model_id, subfolder="transformer", torch_dtype=torch.bfloat16
-        ... )
-        >>> pipe = HunyuanVideoPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch.float16)
-        >>> pipe.vae.enable_tiling()
-        >>> pipe.to("cuda")
-
-        >>> output = pipe(
-        ...     prompt="A cat walks on the grass, realistic",
-        ...     height=320,
-        ...     width=512,
-        ...     num_frames=61,
-        ...     num_inference_steps=30,
-        ... ).frames[0]
-        >>> export_to_video(output, "output.mp4", fps=15)
-        ```
-"""
-
-DEFAULT_PROMPT_TEMPLATE = {
-    "template": ("<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: "
-                 "1. The main content and theme of the video."
-                 "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects."
-                 "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects."
-                 "4. background environment, light, style and atmosphere."
-                 "5. camera angles, movements, and transitions used in the video:<|eot_id|>"
-                 "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>"),
-    "crop_start":
-    95,
-}
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-    r"""
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`List[int]`, *optional*):
-            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
-            `num_inference_steps` and `sigmas` must be `None`.
-        sigmas (`List[float]`, *optional*):
-            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
-            `num_inference_steps` and `timesteps` must be `None`.
-
-    Returns:
-        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
-        second element is the number of inference steps.
-    """
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler.")
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler.")
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-
-
-class HunyuanVideoPipeline(DiffusionPipeline, HunyuanVideoLoraLoaderMixin):
-    r"""
-    Pipeline for text-to-video generation using HunyuanVideo.
-
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
-    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-    Args:
-        text_encoder ([`LlamaModel`]):
-            [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers).
-        tokenizer_2 (`LlamaTokenizer`):
-            Tokenizer from [Llava Llama3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers).
-        transformer ([`HunyuanVideoTransformer3DModel`]):
-            Conditional Transformer to denoise the encoded image latents.
-        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
-            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
-        vae ([`AutoencoderKLHunyuanVideo`]):
-            Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations.
-        text_encoder_2 ([`CLIPTextModel`]):
-            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
-            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
-        tokenizer_2 (`CLIPTokenizer`):
-            Tokenizer of class
-            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
-    """
-
-    model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
-    _callback_tensor_inputs = ["latents", "prompt_embeds"]
-
-    def __init__(
-        self,
-        text_encoder: LlamaModel,
-        tokenizer: LlamaTokenizerFast,
-        transformer: HunyuanVideoTransformer3DModel,
-        vae: AutoencoderKLHunyuanVideo,
-        scheduler: FlowMatchEulerDiscreteScheduler,
-        text_encoder_2: CLIPTextModel,
-        tokenizer_2: CLIPTokenizer,
-    ):
-        super().__init__()
-
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-            transformer=transformer,
-            scheduler=scheduler,
-            text_encoder_2=text_encoder_2,
-            tokenizer_2=tokenizer_2,
-        )
-
-        self.vae_scale_factor_temporal = (self.vae.temporal_compression_ratio
-                                          if hasattr(self, "vae") and self.vae is not None else 4)
-        self.vae_scale_factor_spatial = (self.vae.spatial_compression_ratio
-                                         if hasattr(self, "vae") and self.vae is not None else 8)
-        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
-
-    def _get_llama_prompt_embeds(
-        self,
-        prompt: Union[str, List[str]],
-        prompt_template: Dict[str, Any],
-        num_videos_per_prompt: int = 1,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-        max_sequence_length: int = 256,
-        num_hidden_layers_to_skip: int = 2,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        device = device or self._execution_device
-        dtype = dtype or self.text_encoder.dtype
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt)
-
-        prompt = [prompt_template["template"].format(p) for p in prompt]
-
-        crop_start = prompt_template.get("crop_start", None)
-        if crop_start is None:
-            prompt_template_input = self.tokenizer(
-                prompt_template["template"],
-                padding="max_length",
-                return_tensors="pt",
-                return_length=False,
-                return_overflowing_tokens=False,
-                return_attention_mask=False,
-            )
-            crop_start = prompt_template_input["input_ids"].shape[-1]
-            # Remove <|eot_id|> token and placeholder {}
-            crop_start -= 2
-
-        max_sequence_length += crop_start
-        text_inputs = self.tokenizer(
-            prompt,
-            max_length=max_sequence_length,
-            padding="max_length",
-            truncation=True,
-            return_tensors="pt",
-            return_length=False,
-            return_overflowing_tokens=False,
-            return_attention_mask=True,
-        )
-        text_input_ids = text_inputs.input_ids.to(device=device)
-        prompt_attention_mask = text_inputs.attention_mask.to(device=device)
-
-        prompt_embeds = self.text_encoder(
-            input_ids=text_input_ids,
-            attention_mask=prompt_attention_mask,
-            output_hidden_states=True,
-        ).hidden_states[-(num_hidden_layers_to_skip + 1)]
-        prompt_embeds = prompt_embeds.to(dtype=dtype)
-
-        if crop_start is not None and crop_start > 0:
-            prompt_embeds = prompt_embeds[:, crop_start:]
-            prompt_attention_mask = prompt_attention_mask[:, crop_start:]
-
-        # duplicate text embeddings for each generation per prompt, using mps friendly method
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
-        prompt_attention_mask = prompt_attention_mask.repeat(1, num_videos_per_prompt)
-        prompt_attention_mask = prompt_attention_mask.view(batch_size * num_videos_per_prompt, seq_len)
-
-        return prompt_embeds, prompt_attention_mask
-
-    def _get_clip_prompt_embeds(
-        self,
-        prompt: Union[str, List[str]],
-        num_videos_per_prompt: int = 1,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-        max_sequence_length: int = 77,
-    ) -> torch.Tensor:
-        device = device or self._execution_device
-        dtype = dtype or self.text_encoder_2.dtype
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt)
-
-        text_inputs = self.tokenizer_2(
-            prompt,
-            padding="max_length",
-            max_length=max_sequence_length,
-            truncation=True,
-            return_tensors="pt",
-        )
-
-        text_input_ids = text_inputs.input_ids
-        untruncated_ids = self.tokenizer_2(prompt, padding="longest", return_tensors="pt").input_ids
-        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
-            removed_text = self.tokenizer_2.batch_decode(untruncated_ids[:, max_sequence_length - 1:-1])
-            logger.warning("The following part of your input was truncated because CLIP can only handle sequences up to"
-                           f" {max_sequence_length} tokens: {removed_text}")
-
-        prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False).pooler_output
-
-        # duplicate text embeddings for each generation per prompt, using mps friendly method
-        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt)
-        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, -1)
-
-        return prompt_embeds
-
-    def encode_prompt(
-        self,
-        prompt: Union[str, List[str]],
-        prompt_2: Union[str, List[str]] = None,
-        prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
-        num_videos_per_prompt: int = 1,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        pooled_prompt_embeds: Optional[torch.Tensor] = None,
-        prompt_attention_mask: Optional[torch.Tensor] = None,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-        max_sequence_length: int = 256,
-    ):
-
-        if prompt_embeds is None:
-            prompt_embeds, prompt_attention_mask = self._get_llama_prompt_embeds(
-                prompt,
-                prompt_template,
-                num_videos_per_prompt,
-                device=device,
-                dtype=dtype,
-                max_sequence_length=max_sequence_length,
-            )
-
-        if pooled_prompt_embeds is None:
-            if prompt_2 is None and pooled_prompt_embeds is None:
-                prompt_2 = prompt
-            pooled_prompt_embeds = self._get_clip_prompt_embeds(
-                prompt,
-                num_videos_per_prompt,
-                device=device,
-                dtype=dtype,
-                max_sequence_length=77,
-            )
-
-        return prompt_embeds, pooled_prompt_embeds, prompt_attention_mask
-
-    def check_inputs(
-        self,
-        prompt,
-        prompt_2,
-        height,
-        width,
-        prompt_embeds=None,
-        callback_on_step_end_tensor_inputs=None,
-        prompt_template=None,
-    ):
-        if height % 16 != 0 or width % 16 != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
-
-        if callback_on_step_end_tensor_inputs is not None and not all(k in self._callback_tensor_inputs
-                                                                      for k in callback_on_step_end_tensor_inputs):
-            raise ValueError(
-                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
-            )
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two.")
-        elif prompt_2 is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two.")
-        elif prompt is None and prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.")
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
-            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
-
-        if prompt_template is not None:
-            if not isinstance(prompt_template, dict):
-                raise ValueError(f"`prompt_template` has to be of type `dict` but is {type(prompt_template)}")
-            if "template" not in prompt_template:
-                raise ValueError(
-                    f"`prompt_template` has to contain a key `template` but only found {prompt_template.keys()}")
-
-    def prepare_latents(
-        self,
-        batch_size: int,
-        num_channels_latents: 32,
-        height: int = 720,
-        width: int = 1280,
-        num_frames: int = 129,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        if latents is not None:
-            return latents.to(device=device, dtype=dtype)
-
-        shape = (
-            batch_size,
-            num_channels_latents,
-            num_frames,
-            int(height) // self.vae_scale_factor_spatial,
-            int(width) // self.vae_scale_factor_spatial,
-        )
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators.")
-
-        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        return latents
-
-    def enable_vae_slicing(self):
-        r"""
-        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
-        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
-        """
-        self.vae.enable_slicing()
-
-    def disable_vae_slicing(self):
-        r"""
-        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_slicing()
-
-    def enable_vae_tiling(self):
-        r"""
-        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
-        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
-        processing larger images.
-        """
-        self.vae.enable_tiling()
-
-    def disable_vae_tiling(self):
-        r"""
-        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_tiling()
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @property
-    def attention_kwargs(self):
-        return self._attention_kwargs
-
-    @property
-    def interrupt(self):
-        return self._interrupt
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        prompt_2: Union[str, List[str]] = None,
-        height: int = 720,
-        width: int = 1280,
-        num_frames: int = 129,
-        num_inference_steps: int = 50,
-        sigmas: List[float] = None,
-        guidance_scale: float = 6.0,
-        num_videos_per_prompt: Optional[int] = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        pooled_prompt_embeds: Optional[torch.Tensor] = None,
-        prompt_attention_mask: Optional[torch.Tensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback,
-                                             MultiPipelineCallbacks]] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
-        prompt_template: Dict[str, Any] = DEFAULT_PROMPT_TEMPLATE,
-        max_sequence_length: int = 256,
-    ):
-        r"""
-        The call function to the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            prompt_2 (`str` or `List[str]`, *optional*):
-                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
-                will be used instead.
-            height (`int`, defaults to `720`):
-                The height in pixels of the generated image.
-            width (`int`, defaults to `1280`):
-                The width in pixels of the generated image.
-            num_frames (`int`, defaults to `129`):
-                The number of frames in the generated video.
-            num_inference_steps (`int`, defaults to `50`):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            sigmas (`List[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            guidance_scale (`float`, defaults to `6.0`):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality. Note that the only available HunyuanVideo model is
-                CFG-distilled, which means that traditional guidance between unconditional and conditional latent is
-                not applied.
-            num_videos_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
-                generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor is generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
-                provided, text embeddings are generated from the `prompt` input argument.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`HunyuanVideoPipelineOutput`] instead of a plain tuple.
-            attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            clip_skip (`int`, *optional*):
-                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
-                the output of the pre-final layer will be used for computing the prompt embeddings.
-            callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
-                A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
-                each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
-                DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
-                list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-
-        Examples:
-
-        Returns:
-            [`~HunyuanVideoPipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] is returned, otherwise a `tuple` is returned
-                where the first element is a list with the generated images and the second element is a list of `bool`s
-                indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content.
-        """
-
-        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
-            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt,
-            prompt_2,
-            height,
-            width,
-            prompt_embeds,
-            callback_on_step_end_tensor_inputs,
-            prompt_template,
-        )
-
-        self._guidance_scale = guidance_scale
-        self._attention_kwargs = attention_kwargs
-        self._interrupt = False
-
-        device = self._execution_device
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        # 3. Encode input prompt
-        prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self.encode_prompt(
-            prompt=prompt,
-            prompt_2=prompt,
-            prompt_template=prompt_template,
-            num_videos_per_prompt=num_videos_per_prompt,
-            prompt_embeds=prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            prompt_attention_mask=prompt_attention_mask,
-            device=device,
-            max_sequence_length=max_sequence_length,
-        )
-
-        transformer_dtype = self.transformer.dtype
-        prompt_embeds = prompt_embeds.to(transformer_dtype)
-        prompt_attention_mask = prompt_attention_mask.to(transformer_dtype)
-        if pooled_prompt_embeds is not None:
-            pooled_prompt_embeds = pooled_prompt_embeds.to(transformer_dtype)
-
-        # 4. Prepare timesteps
-        sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] if sigmas is None else sigmas
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler,
-            num_inference_steps,
-            device,
-            sigmas=sigmas,
-        )
-
-        # 5. Prepare latent variables
-        num_channels_latents = self.transformer.config.in_channels
-        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
-
-        latents = self.prepare_latents(
-            batch_size * num_videos_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            num_latent_frames,
-            torch.float32,
-            device,
-            generator,
-            latents,
-        )
-        # check sequence_parallel
-        world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
-        if get_sequence_parallel_state():
-            latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
-            latents = latents[:, :, rank, :, :, :]
-
-        # 6. Prepare guidance condition
-        guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0
-
-        # 7. Denoising loop
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
-        self._num_timesteps = len(timesteps)
-
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                latent_model_input = latents.to(transformer_dtype)
-                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latents.shape[0]).to(latents.dtype)
-                if pooled_prompt_embeds.shape[-1] != prompt_embeds.shape[-1]:
-                    pooled_prompt_embeds_padding = F.pad(
-                        pooled_prompt_embeds,
-                        (0, prompt_embeds.shape[2] - pooled_prompt_embeds.shape[1]),
-                        value=0,
-                    ).unsqueeze(1)
-                encoder_hidden_states = torch.cat([pooled_prompt_embeds_padding, prompt_embeds], dim=1)
-
-                noise_pred = self.transformer(
-                    hidden_states=latent_model_input,
-                    encoder_hidden_states=encoder_hidden_states,  # [1, 257, 4096]
-                    timestep=timestep,
-                    encoder_attention_mask=prompt_attention_mask,
-                    guidance=guidance,
-                    attention_kwargs=attention_kwargs,
-                    return_dict=False,
-                )[0]
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-
-        if get_sequence_parallel_state():
-            latents = all_gather(latents, dim=2)
-
-        if not output_type == "latent":
-            latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor
-            video = self.vae.decode(latents, return_dict=False)[0]
-            video = self.video_processor.postprocess_video(video, output_type=output_type)
-        else:
-            video = latents
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return (video, )
-
-        return HunyuanVideoPipelineOutput(frames=video)

fastvideo/models/loader.py

@@ -0,0 +1,304 @@
+from abc import ABC, abstractmethod
+
+import torch
+from fastvideo.inference_args import InferenceArgs
+from fastvideo.logger import init_logger 
+import os
+from fastvideo.distributed.parallel_state import initialize_sequence_parallel_group
+import glob
+from fastvideo.loader.fsdp_load import load_fsdp_model
+from fastvideo.loader.loader import get_model_loader
+from transformers import PretrainedConfig, AutoTokenizer
+from fastvideo.models.hf_transformer_utils import get_hf_config, get_diffusers_config
+from fastvideo.models import get_scheduler
+from fastvideo.models.registry import ModelRegistry
+from safetensors.torch import load_file as safetensors_load_file
+
+import json
+
+
+logger = init_logger(__name__)
+
+
+class ComponentLoader(ABC):
+    """Base class for loading a specific type of model component."""
+    
+    def __init__(self, device=None):
+        self.device = device
+    
+    @abstractmethod
+    def load(self, model_path: str, architecture: str, inference_args: InferenceArgs):
+        """
+        Load the component based on the model path, architecture, and inference args.
+        
+        Args:
+            model_path: Path to the component model
+            architecture: Architecture of the component model
+            inference_args: Inference arguments
+            
+        Returns:
+            The loaded component
+        """
+        raise NotImplementedError
+    
+    @classmethod
+    def for_module_type(cls, module_type: str, transformers_or_diffusers: str) -> 'ComponentLoader':
+        """
+        Factory method to create a component loader for a specific module type.
+        
+        Args:
+            module_type: Type of module (e.g., "vae", "text_encoder", "transformer", "scheduler")
+            transformers_or_diffusers: Whether the module is from transformers or diffusers
+            
+        Returns:
+            A component loader for the specified module type
+        """
+        # Map of module types to their loader classes and expected library
+        module_loaders = {
+            "scheduler": (SchedulerLoader, "diffusers"),
+            "transformer": (TransformerLoader, "diffusers"),
+            "vae": (VAELoader, "diffusers"),
+            "text_encoder": (TextEncoderLoader, "transformers"),
+            "text_encoder_2": (TextEncoderLoader, "transformers"),
+            "tokenizer": (TokenizerLoader, "transformers"),
+            "tokenizer_2": (TokenizerLoader, "transformers"),
+        }
+        
+        if module_type in module_loaders:
+            loader_cls, expected_library = module_loaders[module_type]
+            # Assert that the library matches what's expected for this module type
+            assert transformers_or_diffusers == expected_library, f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}"
+            return loader_cls()
+        
+        # For unknown module types, use a generic loader
+        logger.warning(f"No specific loader found for module type: {module_type}. Using generic loader.")
+        return GenericComponentLoader(transformers_or_diffusers)
+
+class TextEncoderLoader(ComponentLoader):
+    """Loader for text encoders."""
+
+    def load(self, model_path: str, architecture: str, inference_args: InferenceArgs):
+        """Load the text encoders based on the model path, architecture, and inference args."""
+        use_v1 = True
+        if not use_v1:
+            return self.load_v0(model_path, architecture, inference_args)
+        else:
+            return self.load_v1(model_path, architecture, inference_args)
+    
+    def load_v1(self, model_path: str, architecture: str, inference_args: InferenceArgs):
+        """Load the text encoders based on the model path, architecture, and inference args."""
+        model_config: PretrainedConfig = get_hf_config(
+            model=model_path,
+            trust_remote_code=inference_args.trust_remote_code,
+            revision=inference_args.revision,
+            model_override_args=None,
+            inference_args=inference_args,
+        )
+        logger.info(f"HF Model config: {model_config}")
+        model_loader = get_model_loader(inference_args)
+        model = model_loader.load_model(model_path, model_config, inference_args)
+        return model
+    
+    def load_v0(self, model_path: str, architecture: str, inference_args: InferenceArgs):        
+        from fastvideo.pipelines.hunyuan.constants import PROMPT_TEMPLATE
+        from fastvideo.models.hunyuan.text_encoder import TextEncoder
+        # Text encoder
+        if inference_args.prompt_template_video is not None:
+            crop_start = PROMPT_TEMPLATE[inference_args.prompt_template_video].get("crop_start", 0)
+        elif inference_args.prompt_template is not None:
+            crop_start = PROMPT_TEMPLATE[inference_args.prompt_template].get("crop_start", 0)
+        else:
+            crop_start = 0
+        max_length = inference_args.text_len + crop_start
+
+        # prompt_template
+        prompt_template = (PROMPT_TEMPLATE[inference_args.prompt_template] if inference_args.prompt_template is not None else None)
+
+        # prompt_template_video
+        prompt_template_video = (PROMPT_TEMPLATE[inference_args.prompt_template_video]
+                                 if inference_args.prompt_template_video is not None else None)
+
+        text_encoder = TextEncoder(
+            text_encoder_type=inference_args.text_encoder,
+            text_encoder_path=model_path,
+            max_length=max_length,
+            text_encoder_precision=inference_args.text_encoder_precision,
+            tokenizer_type=inference_args.tokenizer,
+            prompt_template=prompt_template,
+            prompt_template_video=prompt_template_video,
+            hidden_state_skip_layer=inference_args.hidden_state_skip_layer,
+            apply_final_norm=inference_args.apply_final_norm,
+            reproduce=inference_args.reproduce,
+            logger=logger,
+            device=self.device if not inference_args.use_cpu_offload else "cpu",
+        )
+        return text_encoder
+
+class TokenizerLoader(ComponentLoader):
+    """Loader for tokenizers."""
+    
+    def load(self, model_path: str, architecture: str, inference_args: InferenceArgs):
+        """Load the tokenizer based on the model path, architecture, and inference args."""
+        logger.info(f"Loading tokenizer from {model_path}")
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            # TODO(will): pass these tokenizer kwargs from inference args? Maybe
+            # other method of config?
+            padding_size='right',
+        )
+        logger.info(f"Loaded tokenizer: {tokenizer.__class__.__name__}")
+        return tokenizer
+
+class VAELoader(ComponentLoader):
+    """Loader for VAE."""
+    
+    def load(self, model_path: str, architecture: str, inference_args: InferenceArgs):
+        """Load the VAE based on the model path, architecture, and inference args."""
+        # TODO(will): move this to a constants file
+        from fastvideo.pipelines.hunyuan.constants import PRECISION_TO_TYPE
+
+        config = get_diffusers_config(model=model_path)
+        
+        class_name = config.pop("_class_name")
+        assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported."
+        config.pop("_diffusers_version")
+
+        vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+
+        vae = vae_cls(**config).to(inference_args.device)
+        
+        # Find all safetensors files
+        safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
+        # TODO(PY)
+        assert len(safetensors_list) == 1, f"Found {len(safetensors_list)} safetensors files in {path}"
+        loaded = safetensors_load_file(safetensors_list[0])
+        vae.load_state_dict(loaded)
+        dtype = PRECISION_TO_TYPE[inference_args.vae_precision]
+        vae = vae.eval().to(dtype)
+
+        # TODO(will):  should we define hunyuan vae config class?
+        vae_kwargs = {
+            "s_ratio": config["spatial_compression_ratio"],
+            "t_ratio": config["temporal_compression_ratio"],
+        }
+
+        vae.kwargs = vae_kwargs
+        
+        return vae
+
+class TransformerLoader(ComponentLoader):
+    """Loader for transformer."""
+    
+    def load(self, model_path: str, architecture: str, inference_args: InferenceArgs):
+        """Load the transformer based on the model path, architecture, and inference args."""
+        
+        model_config = get_diffusers_config(model=model_path)
+        cls_name = model_config.pop("_class_name")
+        if cls_name is None:
+            raise ValueError(f"Model config does not contain a _class_name attribute. "
+                         "Only diffusers format is supported.")
+        model_config.pop("_diffusers_version")
+
+        model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
+
+        # Find all safetensors files
+        safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
+        if not safetensors_list:
+            raise ValueError(f"No safetensors files found in {model_path}")
+        
+        logger.info(f"Loading model from {len(safetensors_list)} safetensors files in {model_path}")
+
+        # initialize_sequence_parallel_group(inference_args.sp_size)
+        
+        # Load the model using FSDP loader
+        logger.info(f"Loading model from {cls_name}")
+        model = load_fsdp_model(
+            model_cls=model_cls,
+            init_params=model_config,
+            weight_dir_list=safetensors_list,
+            device=inference_args.device,
+            cpu_offload=inference_args.use_cpu_offload
+        )
+        
+        total_params = sum(p.numel() for p in model.parameters())
+        logger.info(f"Loaded model with {total_params / 1e9:.2f}B parameters")
+        
+        model.eval()
+        return model
+
+class SchedulerLoader(ComponentLoader):
+    """Loader for scheduler."""
+    
+    def load(self, model_path: str, architecture: str, inference_args: InferenceArgs):
+        """Load the scheduler based on the model path, architecture, and inference args."""
+        
+        scheduler = get_scheduler(
+            module_path=model_path,
+            architecture=architecture,
+            inference_args=inference_args,
+        )
+        logger.info(f"Scheduler loaded: {scheduler}")
+        return scheduler
+
+class GenericComponentLoader(ComponentLoader):
+    """Generic loader for components that don't have a specific loader."""
+    
+    def __init__(self, library="transformers"):
+        super().__init__()
+        self.library = library
+    
+    def load(self, model_path: str, architecture: str, inference_args: InferenceArgs):
+        """Load a generic component based on the model path, architecture, and inference args."""
+        logger.warning(f"Using generic loader for {model_path} with library {self.library}")
+        
+        if self.library == "transformers":
+            from transformers import AutoModel
+            
+            model = AutoModel.from_pretrained(
+                model_path,
+                trust_remote_code=inference_args.trust_remote_code,
+                revision=inference_args.revision,
+            )
+            logger.info(f"Loaded generic transformers model: {model.__class__.__name__}")
+            return model
+        elif self.library == "diffusers":
+            logger.warning(f"Generic loading for diffusers components is not fully implemented")
+            from fastvideo.models.hf_transformer_utils import get_diffusers_config
+            
+            model_config = get_diffusers_config(model=model_path)
+            logger.info(f"Diffusers Model config: {model_config}")
+            # This is a placeholder - in a real implementation, you'd need to handle this properly
+            return None
+        else:
+            raise ValueError(f"Unsupported library: {self.library}")
+
+class PipelineComponentLoader:
+    """
+    Utility class for loading pipeline components.
+    This replaces the chain of if-else statements in load_pipeline_module.
+    """
+    
+    @staticmethod
+    def load_module(module_name: str, component_model_path: str, transformers_or_diffusers: str, 
+                   architecture: str, inference_args: InferenceArgs):
+        """
+        Load a pipeline module.
+        
+        Args:
+            module_name: Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler")
+            component_model_path: Path to the component model
+            transformers_or_diffusers: Whether the module is from transformers or diffusers
+            architecture: Architecture of the component model
+            inference_args: Inference arguments
+            
+        Returns:
+            The loaded module
+        """
+        logger.info(f"Loading {module_name} using {transformers_or_diffusers} from {component_model_path}")
+        
+        # Get the appropriate loader for this module type
+        loader = ComponentLoader.for_module_type(module_name, transformers_or_diffusers)
+        
+        # Load the module
+        return loader.load(component_model_path, architecture, inference_args)

fastvideo/models/mochi_hf/convert_diffusers_to_mochi.py

@@ -1,361 +0,0 @@
-import argparse
-import os
-
-import torch
-from safetensors.torch import save_file
-
-parser = argparse.ArgumentParser()
-parser.add_argument("--diffusers_path", required=True, type=str)
-parser.add_argument("--transformer_path", type=str, default=None, help="Path to save transformer model")
-parser.add_argument("--vae_encoder_path", type=str, default=None, help="Path to save VAE encoder model")
-parser.add_argument("--vae_decoder_path", type=str, default=None, help="Path to save VAE decoder model")
-
-args = parser.parse_args()
-
-
-def reverse_scale_shift(weight, dim):
-    scale, shift = weight.chunk(2, dim=0)
-    new_weight = torch.cat([shift, scale], dim=0)
-    return new_weight
-
-
-def reverse_proj_gate(weight):
-    gate, proj = weight.chunk(2, dim=0)
-    new_weight = torch.cat([proj, gate], dim=0)
-    return new_weight
-
-
-def convert_diffusers_transformer_to_mochi(state_dict):
-    original_state_dict = state_dict.copy()
-    new_state_dict = {}
-
-    # Convert patch_embed
-    new_state_dict["x_embedder.proj.weight"] = original_state_dict.pop("patch_embed.proj.weight")
-    new_state_dict["x_embedder.proj.bias"] = original_state_dict.pop("patch_embed.proj.bias")
-
-    # Convert time_embed
-    new_state_dict["t_embedder.mlp.0.weight"] = original_state_dict.pop("time_embed.timestep_embedder.linear_1.weight")
-    new_state_dict["t_embedder.mlp.0.bias"] = original_state_dict.pop("time_embed.timestep_embedder.linear_1.bias")
-    new_state_dict["t_embedder.mlp.2.weight"] = original_state_dict.pop("time_embed.timestep_embedder.linear_2.weight")
-    new_state_dict["t_embedder.mlp.2.bias"] = original_state_dict.pop("time_embed.timestep_embedder.linear_2.bias")
-    new_state_dict["t5_y_embedder.to_kv.weight"] = original_state_dict.pop("time_embed.pooler.to_kv.weight")
-    new_state_dict["t5_y_embedder.to_kv.bias"] = original_state_dict.pop("time_embed.pooler.to_kv.bias")
-    new_state_dict["t5_y_embedder.to_q.weight"] = original_state_dict.pop("time_embed.pooler.to_q.weight")
-    new_state_dict["t5_y_embedder.to_q.bias"] = original_state_dict.pop("time_embed.pooler.to_q.bias")
-    new_state_dict["t5_y_embedder.to_out.weight"] = original_state_dict.pop("time_embed.pooler.to_out.weight")
-    new_state_dict["t5_y_embedder.to_out.bias"] = original_state_dict.pop("time_embed.pooler.to_out.bias")
-    new_state_dict["t5_yproj.weight"] = original_state_dict.pop("time_embed.caption_proj.weight")
-    new_state_dict["t5_yproj.bias"] = original_state_dict.pop("time_embed.caption_proj.bias")
-
-    # Convert transformer blocks
-    num_layers = 48
-    for i in range(num_layers):
-        block_prefix = f"transformer_blocks.{i}."
-        new_prefix = f"blocks.{i}."
-
-        # norm1
-        new_state_dict[new_prefix + "mod_x.weight"] = original_state_dict.pop(block_prefix + "norm1.linear.weight")
-        new_state_dict[new_prefix + "mod_x.bias"] = original_state_dict.pop(block_prefix + "norm1.linear.bias")
-
-        if i < num_layers - 1:
-            new_state_dict[new_prefix + "mod_y.weight"] = original_state_dict.pop(block_prefix +
-                                                                                  "norm1_context.linear.weight")
-            new_state_dict[new_prefix + "mod_y.bias"] = original_state_dict.pop(block_prefix +
-                                                                                "norm1_context.linear.bias")
-        else:
-            new_state_dict[new_prefix + "mod_y.weight"] = original_state_dict.pop(block_prefix +
-                                                                                  "norm1_context.linear_1.weight")
-            new_state_dict[new_prefix + "mod_y.bias"] = original_state_dict.pop(block_prefix +
-                                                                                "norm1_context.linear_1.bias")
-
-        # Visual attention
-        q = original_state_dict.pop(block_prefix + "attn1.to_q.weight")
-        k = original_state_dict.pop(block_prefix + "attn1.to_k.weight")
-        v = original_state_dict.pop(block_prefix + "attn1.to_v.weight")
-        qkv_weight = torch.cat([q, k, v], dim=0)
-        new_state_dict[new_prefix + "attn.qkv_x.weight"] = qkv_weight
-
-        new_state_dict[new_prefix + "attn.q_norm_x.weight"] = original_state_dict.pop(block_prefix +
-                                                                                      "attn1.norm_q.weight")
-        new_state_dict[new_prefix + "attn.k_norm_x.weight"] = original_state_dict.pop(block_prefix +
-                                                                                      "attn1.norm_k.weight")
-        new_state_dict[new_prefix + "attn.proj_x.weight"] = original_state_dict.pop(block_prefix +
-                                                                                    "attn1.to_out.0.weight")
-        new_state_dict[new_prefix + "attn.proj_x.bias"] = original_state_dict.pop(block_prefix + "attn1.to_out.0.bias")
-
-        # Context attention
-        q = original_state_dict.pop(block_prefix + "attn1.add_q_proj.weight")
-        k = original_state_dict.pop(block_prefix + "attn1.add_k_proj.weight")
-        v = original_state_dict.pop(block_prefix + "attn1.add_v_proj.weight")
-        qkv_weight = torch.cat([q, k, v], dim=0)
-        new_state_dict[new_prefix + "attn.qkv_y.weight"] = qkv_weight
-
-        new_state_dict[new_prefix + "attn.q_norm_y.weight"] = original_state_dict.pop(block_prefix +
-                                                                                      "attn1.norm_added_q.weight")
-        new_state_dict[new_prefix + "attn.k_norm_y.weight"] = original_state_dict.pop(block_prefix +
-                                                                                      "attn1.norm_added_k.weight")
-        if i < num_layers - 1:
-            new_state_dict[new_prefix + "attn.proj_y.weight"] = original_state_dict.pop(block_prefix +
-                                                                                        "attn1.to_add_out.weight")
-            new_state_dict[new_prefix + "attn.proj_y.bias"] = original_state_dict.pop(block_prefix +
-                                                                                      "attn1.to_add_out.bias")
-
-        # MLP
-        new_state_dict[new_prefix + "mlp_x.w1.weight"] = reverse_proj_gate(
-            original_state_dict.pop(block_prefix + "ff.net.0.proj.weight"))
-        new_state_dict[new_prefix + "mlp_x.w2.weight"] = original_state_dict.pop(block_prefix + "ff.net.2.weight")
-        if i < num_layers - 1:
-            new_state_dict[new_prefix + "mlp_y.w1.weight"] = reverse_proj_gate(
-                original_state_dict.pop(block_prefix + "ff_context.net.0.proj.weight"))
-            new_state_dict[new_prefix + "mlp_y.w2.weight"] = original_state_dict.pop(block_prefix +
-                                                                                     "ff_context.net.2.weight")
-
-    # Output layers
-    new_state_dict["final_layer.mod.weight"] = reverse_scale_shift(original_state_dict.pop("norm_out.linear.weight"),
-                                                                   dim=0)
-    new_state_dict["final_layer.mod.bias"] = reverse_scale_shift(original_state_dict.pop("norm_out.linear.bias"), dim=0)
-    new_state_dict["final_layer.linear.weight"] = original_state_dict.pop("proj_out.weight")
-    new_state_dict["final_layer.linear.bias"] = original_state_dict.pop("proj_out.bias")
-
-    new_state_dict["pos_frequencies"] = original_state_dict.pop("pos_frequencies")
-
-    print("Remaining Keys:", original_state_dict.keys())
-
-    return new_state_dict
-
-
-def convert_diffusers_vae_to_mochi(state_dict):
-    original_state_dict = state_dict.copy()
-    encoder_state_dict = {}
-    decoder_state_dict = {}
-
-    # Convert encoder
-    prefix = "encoder."
-
-    encoder_state_dict["layers.0.weight"] = original_state_dict.pop(f"{prefix}proj_in.weight")
-    encoder_state_dict["layers.0.bias"] = original_state_dict.pop(f"{prefix}proj_in.bias")
-
-    # Convert block_in
-    for i in range(3):
-        encoder_state_dict[f"layers.{i+1}.stack.0.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight")
-        encoder_state_dict[f"layers.{i+1}.stack.0.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias")
-        encoder_state_dict[f"layers.{i+1}.stack.2.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv1.conv.weight")
-        encoder_state_dict[f"layers.{i+1}.stack.2.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv1.conv.bias")
-        encoder_state_dict[f"layers.{i+1}.stack.3.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight")
-        encoder_state_dict[f"layers.{i+1}.stack.3.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias")
-        encoder_state_dict[f"layers.{i+1}.stack.5.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv2.conv.weight")
-        encoder_state_dict[f"layers.{i+1}.stack.5.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv2.conv.bias")
-
-    # Convert down_blocks
-    down_block_layers = [3, 4, 6]
-    for block in range(3):
-        encoder_state_dict[f"layers.{block+4}.layers.0.weight"] = original_state_dict.pop(
-            f"{prefix}down_blocks.{block}.conv_in.conv.weight")
-        encoder_state_dict[f"layers.{block+4}.layers.0.bias"] = original_state_dict.pop(
-            f"{prefix}down_blocks.{block}.conv_in.conv.bias")
-
-        for i in range(down_block_layers[block]):
-            # Convert resnets
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.0.weight"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.weight")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.0.bias"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.bias")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.2.weight"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.weight")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.2.bias"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.bias")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.3.weight"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.weight")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.3.bias"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.bias")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.5.weight"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.weight")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.5.bias"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.bias")
-
-            # Convert attentions
-            q = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_q.weight")
-            k = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_k.weight")
-            v = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_v.weight")
-            qkv_weight = torch.cat([q, k, v], dim=0)
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.qkv.weight"] = qkv_weight
-
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.out.weight"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.weight")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.out.bias"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.bias")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.norm.weight"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.weight")
-            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.norm.bias"] = original_state_dict.pop(
-                f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.bias")
-
-    # Convert block_out
-    for i in range(3):
-        encoder_state_dict[f"layers.{i+7}.stack.0.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight")
-        encoder_state_dict[f"layers.{i+7}.stack.0.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias")
-        encoder_state_dict[f"layers.{i+7}.stack.2.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv1.conv.weight")
-        encoder_state_dict[f"layers.{i+7}.stack.2.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv1.conv.bias")
-        encoder_state_dict[f"layers.{i+7}.stack.3.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight")
-        encoder_state_dict[f"layers.{i+7}.stack.3.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias")
-        encoder_state_dict[f"layers.{i+7}.stack.5.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv2.conv.weight")
-        encoder_state_dict[f"layers.{i+7}.stack.5.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv2.conv.bias")
-
-        q = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_q.weight")
-        k = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_k.weight")
-        v = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_v.weight")
-        qkv_weight = torch.cat([q, k, v], dim=0)
-        encoder_state_dict[f"layers.{i+7}.attn_block.attn.qkv.weight"] = qkv_weight
-
-        encoder_state_dict[f"layers.{i+7}.attn_block.attn.out.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.attentions.{i}.to_out.0.weight")
-        encoder_state_dict[f"layers.{i+7}.attn_block.attn.out.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.attentions.{i}.to_out.0.bias")
-        encoder_state_dict[f"layers.{i+7}.attn_block.norm.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.norms.{i}.norm_layer.weight")
-        encoder_state_dict[f"layers.{i+7}.attn_block.norm.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.norms.{i}.norm_layer.bias")
-
-    # Convert output layers
-    encoder_state_dict["output_norm.weight"] = original_state_dict.pop(f"{prefix}norm_out.norm_layer.weight")
-    encoder_state_dict["output_norm.bias"] = original_state_dict.pop(f"{prefix}norm_out.norm_layer.bias")
-    encoder_state_dict["output_proj.weight"] = original_state_dict.pop(f"{prefix}proj_out.weight")
-
-    # Convert decoder
-    prefix = "decoder."
-
-    decoder_state_dict["blocks.0.0.weight"] = original_state_dict.pop(f"{prefix}conv_in.weight")
-    decoder_state_dict["blocks.0.0.bias"] = original_state_dict.pop(f"{prefix}conv_in.bias")
-
-    # Convert block_in
-    for i in range(3):
-        decoder_state_dict[f"blocks.0.{i+1}.stack.0.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.0.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.2.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv1.conv.weight")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.2.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv1.conv.bias")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.3.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.3.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.5.weight"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv2.conv.weight")
-        decoder_state_dict[f"blocks.0.{i+1}.stack.5.bias"] = original_state_dict.pop(
-            f"{prefix}block_in.resnets.{i}.conv2.conv.bias")
-
-    # Convert up_blocks
-    up_block_layers = [6, 4, 3]
-    for block in range(3):
-        for i in range(up_block_layers[block]):
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.0.weight"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.weight")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.0.bias"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.bias")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.2.weight"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.weight")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.2.bias"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.bias")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.3.weight"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.weight")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.3.bias"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.bias")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.5.weight"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.weight")
-            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.5.bias"] = original_state_dict.pop(
-                f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.bias")
-        decoder_state_dict[f"blocks.{block+1}.proj.weight"] = original_state_dict.pop(
-            f"{prefix}up_blocks.{block}.proj.weight")
-        decoder_state_dict[f"blocks.{block+1}.proj.bias"] = original_state_dict.pop(
-            f"{prefix}up_blocks.{block}.proj.bias")
-
-    # Convert block_out
-    for i in range(3):
-        decoder_state_dict[f"blocks.4.{i}.stack.0.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight")
-        decoder_state_dict[f"blocks.4.{i}.stack.0.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias")
-        decoder_state_dict[f"blocks.4.{i}.stack.2.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv1.conv.weight")
-        decoder_state_dict[f"blocks.4.{i}.stack.2.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv1.conv.bias")
-        decoder_state_dict[f"blocks.4.{i}.stack.3.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight")
-        decoder_state_dict[f"blocks.4.{i}.stack.3.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias")
-        decoder_state_dict[f"blocks.4.{i}.stack.5.weight"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv2.conv.weight")
-        decoder_state_dict[f"blocks.4.{i}.stack.5.bias"] = original_state_dict.pop(
-            f"{prefix}block_out.resnets.{i}.conv2.conv.bias")
-
-    # Convert output layers
-    decoder_state_dict["output_proj.weight"] = original_state_dict.pop(f"{prefix}proj_out.weight")
-    decoder_state_dict["output_proj.bias"] = original_state_dict.pop(f"{prefix}proj_out.bias")
-
-    return encoder_state_dict, decoder_state_dict
-
-
-def ensure_safetensors_extension(path):
-    if not path.endswith(".safetensors"):
-        path = path + ".safetensors"
-    return path
-
-
-def ensure_directory_exists(path):
-    directory = os.path.dirname(path)
-    if directory:
-        os.makedirs(directory, exist_ok=True)
-
-
-def main(args):
-    from diffusers import MochiPipeline
-
-    pipe = MochiPipeline.from_pretrained(args.diffusers_path)
-
-    if args.transformer_path:
-        transformer_path = ensure_safetensors_extension(args.transformer_path)
-        ensure_directory_exists(transformer_path)
-
-        print("Converting transformer model...")
-        transformer_state_dict = convert_diffusers_transformer_to_mochi(pipe.transformer.state_dict())
-        save_file(transformer_state_dict, transformer_path)
-        print(f"Saved transformer to {transformer_path}")
-
-    if args.vae_encoder_path and args.vae_decoder_path:
-        encoder_path = ensure_safetensors_extension(args.vae_encoder_path)
-        decoder_path = ensure_safetensors_extension(args.vae_decoder_path)
-
-        ensure_directory_exists(encoder_path)
-        ensure_directory_exists(decoder_path)
-
-        print("Converting VAE models...")
-        encoder_state_dict, decoder_state_dict = convert_diffusers_vae_to_mochi(pipe.vae.state_dict())
-
-        save_file(encoder_state_dict, encoder_path)
-        print(f"Saved VAE encoder to {encoder_path}")
-
-        save_file(decoder_state_dict, decoder_path)
-        print(f"Saved VAE decoder to {decoder_path}")
-    elif args.vae_encoder_path or args.vae_decoder_path:
-        print("Warning: Both VAE encoder and decoder paths must be specified to convert VAE models.")
-
-
-if __name__ == "__main__":
-    main(args)

fastvideo/models/mochi_hf/mochi_latents_utils.py

@@ -1,45 +0,0 @@
-import torch
-
-mochi_latents_mean = torch.tensor([
-    -0.06730895953510081,
-    -0.038011381506090416,
-    -0.07477820912866141,
-    -0.05565264470995561,
-    0.012767231469026969,
-    -0.04703542746246419,
-    0.043896967884726704,
-    -0.09346305707025976,
-    -0.09918314763016893,
-    -0.008729793427399178,
-    -0.011931556316503654,
-    -0.0321993391887285,
-]).view(1, 12, 1, 1, 1)
-mochi_latents_std = torch.tensor([
-    0.9263795028493863,
-    0.9248894543193766,
-    0.9393059390890617,
-    0.959253732819592,
-    0.8244560132752793,
-    0.917259975397747,
-    0.9294154431013696,
-    1.3720942357788521,
-    0.881393668867029,
-    0.9168315692124348,
-    0.9185249279345552,
-    0.9274757570805041,
-]).view(1, 12, 1, 1, 1)
-mochi_scaling_factor = 1.0
-
-
-def normalize_dit_input(model_type, latents):
-    if model_type == "mochi":
-        latents_mean = mochi_latents_mean.to(latents.device, latents.dtype)
-        latents_std = mochi_latents_std.to(latents.device, latents.dtype)
-        latents = (latents - latents_mean) / latents_std
-        return latents
-    elif model_type == "hunyuan_hf":
-        return latents * 0.476986
-    elif model_type == "hunyuan":
-        return latents * 0.476986
-    else:
-        raise NotImplementedError(f"model_type {model_type} not supported")

fastvideo/models/mochi_hf/modeling_mochi.py

@@ -1,663 +0,0 @@
-# Copyright 2024 The Genmo team and The HuggingFace Team.
-# All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import Any, Dict, Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.loaders import PeftAdapterMixin
-from diffusers.models.attention import FeedForward as HF_FeedForward
-from diffusers.models.attention_processor import Attention
-from diffusers.models.embeddings import MochiCombinedTimestepCaptionEmbedding, PatchEmbed
-from diffusers.models.modeling_utils import ModelMixin
-from diffusers.models.normalization import AdaLayerNormContinuous
-from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
-from diffusers.utils.torch_utils import maybe_allow_in_graph
-from liger_kernel.ops.swiglu import LigerSiLUMulFunction
-
-from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad
-from fastvideo.models.mochi_hf.norm import (MochiLayerNormContinuous, MochiModulatedRMSNorm, MochiRMSNorm,
-                                            MochiRMSNormZero)
-from fastvideo.utils.communications import all_gather, all_to_all_4D
-from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-class FeedForward(HF_FeedForward):
-
-    def __init__(
-        self,
-        dim: int,
-        dim_out: Optional[int] = None,
-        mult: int = 4,
-        dropout: float = 0.0,
-        activation_fn: str = "geglu",
-        final_dropout: bool = False,
-        inner_dim=None,
-        bias: bool = True,
-    ):
-        super().__init__(dim, dim_out, mult, dropout, activation_fn, final_dropout, inner_dim, bias)
-        assert activation_fn == "swiglu"
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.net[0].proj(hidden_states)
-        hidden_states, gate = hidden_states.chunk(2, dim=-1)
-
-        return self.net[2](LigerSiLUMulFunction.apply(gate, hidden_states))
-
-
-class MochiAttention(nn.Module):
-
-    def __init__(
-        self,
-        query_dim: int,
-        processor: "MochiAttnProcessor2_0",
-        heads: int = 8,
-        dim_head: int = 64,
-        dropout: float = 0.0,
-        bias: bool = False,
-        added_kv_proj_dim: Optional[int] = None,
-        added_proj_bias: Optional[bool] = True,
-        out_dim: int = None,
-        out_context_dim: int = None,
-        out_bias: bool = True,
-        context_pre_only: bool = False,
-        eps: float = 1e-5,
-    ):
-        super().__init__()
-        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
-        self.out_dim = out_dim if out_dim is not None else query_dim
-        self.out_context_dim = out_context_dim if out_context_dim else query_dim
-        self.context_pre_only = context_pre_only
-
-        self.heads = out_dim // dim_head if out_dim is not None else heads
-
-        self.norm_q = MochiRMSNorm(dim_head, eps)
-        self.norm_k = MochiRMSNorm(dim_head, eps)
-        self.norm_added_q = MochiRMSNorm(dim_head, eps)
-        self.norm_added_k = MochiRMSNorm(dim_head, eps)
-
-        self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
-        self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias)
-        self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias)
-
-        self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
-        self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
-        if self.context_pre_only is not None:
-            self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
-
-        self.to_out = nn.ModuleList([])
-        self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
-        self.to_out.append(nn.Dropout(dropout))
-
-        if not self.context_pre_only:
-            self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias)
-
-        self.processor = processor
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        **kwargs,
-    ):
-        return self.processor(
-            self,
-            hidden_states,
-            encoder_hidden_states=encoder_hidden_states,
-            attention_mask=attention_mask,
-            **kwargs,
-        )
-
-
-class MochiAttnProcessor2_0:
-    """Attention processor used in Mochi."""
-
-    def __init__(self):
-        if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError("MochiAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.")
-
-    def __call__(
-        self,
-        attn: Attention,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        encoder_attention_mask: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        image_rotary_emb: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        # [b, s, h * d]
-        query = attn.to_q(hidden_states)
-        key = attn.to_k(hidden_states)
-        value = attn.to_v(hidden_states)
-
-        # [b, s, h=24, d=128]
-        query = query.unflatten(2, (attn.heads, -1))
-        key = key.unflatten(2, (attn.heads, -1))
-        value = value.unflatten(2, (attn.heads, -1))
-
-        if attn.norm_q is not None:
-            query = attn.norm_q(query)
-        if attn.norm_k is not None:
-            key = attn.norm_k(key)
-        # [b, 256, h * d]
-        encoder_query = attn.add_q_proj(encoder_hidden_states)
-        encoder_key = attn.add_k_proj(encoder_hidden_states)
-        encoder_value = attn.add_v_proj(encoder_hidden_states)
-
-        # [b, 256, h=24, d=128]
-        encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
-        encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
-        encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
-
-        if attn.norm_added_q is not None:
-            encoder_query = attn.norm_added_q(encoder_query)
-        if attn.norm_added_k is not None:
-            encoder_key = attn.norm_added_k(encoder_key)
-
-        if image_rotary_emb is not None:
-            freqs_cos, freqs_sin = image_rotary_emb[0], image_rotary_emb[1]
-        # shard the head dimension
-        if get_sequence_parallel_state():
-            # B, S, H, D to (S, B,) H, D
-            # batch_size, seq_len, attn_heads, head_dim
-            query = all_to_all_4D(query, scatter_dim=2, gather_dim=1)
-            key = all_to_all_4D(key, scatter_dim=2, gather_dim=1)
-            value = all_to_all_4D(value, scatter_dim=2, gather_dim=1)
-
-            def shrink_head(encoder_state, dim):
-                local_heads = encoder_state.shape[dim] // nccl_info.sp_size
-                return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
-
-            encoder_query = shrink_head(encoder_query, dim=2)
-            encoder_key = shrink_head(encoder_key, dim=2)
-            encoder_value = shrink_head(encoder_value, dim=2)
-            if image_rotary_emb is not None:
-                freqs_cos = shrink_head(freqs_cos, dim=1)
-                freqs_sin = shrink_head(freqs_sin, dim=1)
-
-        if image_rotary_emb is not None:
-
-            def apply_rotary_emb(x, freqs_cos, freqs_sin):
-                x_even = x[..., 0::2].float()
-                x_odd = x[..., 1::2].float()
-                cos = (x_even * freqs_cos - x_odd * freqs_sin).to(x.dtype)
-                sin = (x_even * freqs_sin + x_odd * freqs_cos).to(x.dtype)
-
-                return torch.stack([cos, sin], dim=-1).flatten(-2)
-
-            query = apply_rotary_emb(query, freqs_cos, freqs_sin)
-            key = apply_rotary_emb(key, freqs_cos, freqs_sin)
-
-        # query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
-        # encoder_query, encoder_key, encoder_value = (
-        #     encoder_query.transpose(1, 2),
-        #     encoder_key.transpose(1, 2),
-        #     encoder_value.transpose(1, 2),
-        # )
-        # [b, s, h, d]
-        sequence_length = query.size(1)
-        encoder_sequence_length = encoder_query.size(1)
-
-        # H
-        query = torch.cat([query, encoder_query], dim=1).unsqueeze(2)
-        key = torch.cat([key, encoder_key], dim=1).unsqueeze(2)
-        value = torch.cat([value, encoder_value], dim=1).unsqueeze(2)
-        # B, S, 3, H, D
-        qkv = torch.cat([query, key, value], dim=2)
-
-        attn_mask = encoder_attention_mask[:, :].bool()
-        attn_mask = F.pad(attn_mask, (sequence_length, 0), value=True)
-        hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None)
-
-        # hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask = None, dropout_p=0.0, is_causal=False)
-
-        # valid_lengths = encoder_attention_mask.sum(dim=1) + sequence_length
-        # def no_padding_mask(score, b, h, q_idx, kv_idx):
-        #     return torch.where(kv_idx < valid_lengths[b],score,  -float("inf"))
-
-        # hidden_states = flex_attention(query, key, value, score_mod=no_padding_mask)
-        if get_sequence_parallel_state():
-            hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
-                (sequence_length, encoder_sequence_length), dim=1)
-            # B, S, H, D
-            hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2)
-            encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous()
-            hidden_states = hidden_states.flatten(2, 3)
-            hidden_states = hidden_states.to(query.dtype)
-            encoder_hidden_states = encoder_hidden_states.flatten(2, 3)
-            encoder_hidden_states = encoder_hidden_states.to(query.dtype)
-        else:
-            hidden_states = hidden_states.flatten(2, 3)
-            hidden_states = hidden_states.to(query.dtype)
-
-            hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
-                (sequence_length, encoder_sequence_length), dim=1)
-
-        # linear proj
-        hidden_states = attn.to_out[0](hidden_states)
-        # dropout
-        hidden_states = attn.to_out[1](hidden_states)
-
-        if hasattr(attn, "to_add_out"):
-            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
-
-        return hidden_states, encoder_hidden_states
-
-
-@maybe_allow_in_graph
-class MochiTransformerBlock(nn.Module):
-    r"""
-    Transformer block used in [Mochi](https://huggingface.co/genmo/mochi-1-preview).
-
-    Args:
-        dim (`int`):
-            The number of channels in the input and output.
-        num_attention_heads (`int`):
-            The number of heads to use for multi-head attention.
-        attention_head_dim (`int`):
-            The number of channels in each head.
-        qk_norm (`str`, defaults to `"rms_norm"`):
-            The normalization layer to use.
-        activation_fn (`str`, defaults to `"swiglu"`):
-            Activation function to use in feed-forward.
-        context_pre_only (`bool`, defaults to `False`):
-            Whether or not to process context-related conditions with additional layers.
-        eps (`float`, defaults to `1e-6`):
-            Epsilon value for normalization layers.
-    """
-
-    def __init__(
-        self,
-        dim: int,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        pooled_projection_dim: int,
-        qk_norm: str = "rms_norm",
-        activation_fn: str = "swiglu",
-        context_pre_only: bool = False,
-        eps: float = 1e-6,
-    ) -> None:
-        super().__init__()
-
-        self.context_pre_only = context_pre_only
-        self.ff_inner_dim = (4 * dim * 2) // 3
-        self.ff_context_inner_dim = (4 * pooled_projection_dim * 2) // 3
-
-        self.norm1 = MochiRMSNormZero(dim, 4 * dim, eps=eps, elementwise_affine=False)
-
-        if not context_pre_only:
-            self.norm1_context = MochiRMSNormZero(dim, 4 * pooled_projection_dim, eps=eps, elementwise_affine=False)
-        else:
-            self.norm1_context = MochiLayerNormContinuous(
-                embedding_dim=pooled_projection_dim,
-                conditioning_embedding_dim=dim,
-                eps=eps,
-            )
-
-        self.attn1 = MochiAttention(
-            query_dim=dim,
-            heads=num_attention_heads,
-            dim_head=attention_head_dim,
-            bias=False,
-            added_kv_proj_dim=pooled_projection_dim,
-            added_proj_bias=False,
-            out_dim=dim,
-            out_context_dim=pooled_projection_dim,
-            context_pre_only=context_pre_only,
-            processor=MochiAttnProcessor2_0(),
-            eps=1e-5,
-        )
-
-        # TODO(aryan): norm_context layers are not needed when `context_pre_only` is True
-        self.norm2 = MochiModulatedRMSNorm(eps=eps)
-        self.norm2_context = (MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None)
-
-        self.norm3 = MochiModulatedRMSNorm(eps)
-        self.norm3_context = (MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None)
-
-        self.ff = FeedForward(dim, inner_dim=self.ff_inner_dim, activation_fn=activation_fn, bias=False)
-        self.ff_context = None
-        if not context_pre_only:
-            self.ff_context = FeedForward(
-                pooled_projection_dim,
-                inner_dim=self.ff_context_inner_dim,
-                activation_fn=activation_fn,
-                bias=False,
-            )
-
-        self.norm4 = MochiModulatedRMSNorm(eps=eps)
-        self.norm4_context = MochiModulatedRMSNorm(eps=eps)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        encoder_attention_mask: torch.Tensor,
-        temb: torch.Tensor,
-        image_rotary_emb: Optional[torch.Tensor] = None,
-        output_attn=False,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
-
-        if not self.context_pre_only:
-            (
-                norm_encoder_hidden_states,
-                enc_gate_msa,
-                enc_scale_mlp,
-                enc_gate_mlp,
-            ) = self.norm1_context(encoder_hidden_states, temb)
-        else:
-            norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, temb)
-
-        attn_hidden_states, context_attn_hidden_states = self.attn1(
-            hidden_states=norm_hidden_states,
-            encoder_hidden_states=norm_encoder_hidden_states,
-            image_rotary_emb=image_rotary_emb,
-            encoder_attention_mask=encoder_attention_mask,
-        )
-
-        hidden_states = hidden_states + self.norm2(attn_hidden_states, torch.tanh(gate_msa).unsqueeze(1))
-        norm_hidden_states = self.norm3(hidden_states, (1 + scale_mlp.unsqueeze(1).to(torch.float32)))
-        ff_output = self.ff(norm_hidden_states)
-        hidden_states = hidden_states + self.norm4(ff_output, torch.tanh(gate_mlp).unsqueeze(1))
-
-        if not self.context_pre_only:
-            encoder_hidden_states = encoder_hidden_states + self.norm2_context(context_attn_hidden_states,
-                                                                               torch.tanh(enc_gate_msa).unsqueeze(1))
-            norm_encoder_hidden_states = self.norm3_context(
-                encoder_hidden_states,
-                (1 + enc_scale_mlp.unsqueeze(1).to(torch.float32)),
-            )
-            context_ff_output = self.ff_context(norm_encoder_hidden_states)
-            encoder_hidden_states = encoder_hidden_states + self.norm4_context(context_ff_output,
-                                                                               torch.tanh(enc_gate_mlp).unsqueeze(1))
-
-        if not output_attn:
-            attn_hidden_states = None
-        return hidden_states, encoder_hidden_states, attn_hidden_states
-
-
-class MochiRoPE(nn.Module):
-    r"""
-    RoPE implementation used in [Mochi](https://huggingface.co/genmo/mochi-1-preview).
-
-    Args:
-        base_height (`int`, defaults to `192`):
-            Base height used to compute interpolation scale for rotary positional embeddings.
-        base_width (`int`, defaults to `192`):
-            Base width used to compute interpolation scale for rotary positional embeddings.
-    """
-
-    def __init__(self, base_height: int = 192, base_width: int = 192) -> None:
-        super().__init__()
-
-        self.target_area = base_height * base_width
-
-    def _centers(self, start, stop, num, device, dtype) -> torch.Tensor:
-        edges = torch.linspace(start, stop, num + 1, device=device, dtype=dtype)
-        return (edges[:-1] + edges[1:]) / 2
-
-    def _get_positions(
-        self,
-        num_frames: int,
-        height: int,
-        width: int,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ) -> torch.Tensor:
-        scale = (self.target_area / (height * width))**0.5
-        t = torch.arange(num_frames * nccl_info.sp_size, device=device, dtype=dtype)
-        h = self._centers(-height * scale / 2, height * scale / 2, height, device, dtype)
-        w = self._centers(-width * scale / 2, width * scale / 2, width, device, dtype)
-
-        grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij")
-
-        positions = torch.stack([grid_t, grid_h, grid_w], dim=-1).view(-1, 3)
-        return positions
-
-    def _create_rope(self, freqs: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
-        with torch.autocast(freqs.device.type, enabled=False):
-            # Always run ROPE freqs computation in FP32
-            freqs = torch.einsum(
-                "nd,dhf->nhf",  # codespell:ignore
-                pos.to(torch.float32),  # codespell:ignore
-                freqs.to(torch.float32))
-        freqs_cos = torch.cos(freqs)
-        freqs_sin = torch.sin(freqs)
-        return freqs_cos, freqs_sin
-
-    def forward(
-        self,
-        pos_frequencies: torch.Tensor,
-        num_frames: int,
-        height: int,
-        width: int,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        pos = self._get_positions(num_frames, height, width, device, dtype)
-        rope_cos, rope_sin = self._create_rope(pos_frequencies, pos)
-        return rope_cos, rope_sin
-
-
-@maybe_allow_in_graph
-class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
-    r"""
-    A Transformer model for video-like data introduced in [Mochi](https://huggingface.co/genmo/mochi-1-preview).
-
-    Args:
-        patch_size (`int`, defaults to `2`):
-            The size of the patches to use in the patch embedding layer.
-        num_attention_heads (`int`, defaults to `24`):
-            The number of heads to use for multi-head attention.
-        attention_head_dim (`int`, defaults to `128`):
-            The number of channels in each head.
-        num_layers (`int`, defaults to `48`):
-            The number of layers of Transformer blocks to use.
-        in_channels (`int`, defaults to `12`):
-            The number of channels in the input.
-        out_channels (`int`, *optional*, defaults to `None`):
-            The number of channels in the output.
-        qk_norm (`str`, defaults to `"rms_norm"`):
-            The normalization layer to use.
-        text_embed_dim (`int`, defaults to `4096`):
-            Input dimension of text embeddings from the text encoder.
-        time_embed_dim (`int`, defaults to `256`):
-            Output dimension of timestep embeddings.
-        activation_fn (`str`, defaults to `"swiglu"`):
-            Activation function to use in feed-forward.
-        max_sequence_length (`int`, defaults to `256`):
-            The maximum sequence length of text embeddings supported.
-    """
-
-    _supports_gradient_checkpointing = True
-
-    @register_to_config
-    def __init__(
-        self,
-        patch_size: int = 2,
-        num_attention_heads: int = 24,
-        attention_head_dim: int = 128,
-        num_layers: int = 48,
-        pooled_projection_dim: int = 1536,
-        in_channels: int = 12,
-        out_channels: Optional[int] = None,
-        qk_norm: str = "rms_norm",
-        text_embed_dim: int = 4096,
-        time_embed_dim: int = 256,
-        activation_fn: str = "swiglu",
-        max_sequence_length: int = 256,
-    ) -> None:
-        super().__init__()
-
-        inner_dim = num_attention_heads * attention_head_dim
-        out_channels = out_channels or in_channels
-
-        self.patch_embed = PatchEmbed(
-            patch_size=patch_size,
-            in_channels=in_channels,
-            embed_dim=inner_dim,
-            pos_embed_type=None,
-        )
-
-        self.time_embed = MochiCombinedTimestepCaptionEmbedding(
-            embedding_dim=inner_dim,
-            pooled_projection_dim=pooled_projection_dim,
-            text_embed_dim=text_embed_dim,
-            time_embed_dim=time_embed_dim,
-            num_attention_heads=8,
-        )
-
-        self.pos_frequencies = nn.Parameter(torch.full((3, num_attention_heads, attention_head_dim // 2), 0.0))
-        self.rope = MochiRoPE()
-
-        self.transformer_blocks = nn.ModuleList([
-            MochiTransformerBlock(
-                dim=inner_dim,
-                num_attention_heads=num_attention_heads,
-                attention_head_dim=attention_head_dim,
-                pooled_projection_dim=pooled_projection_dim,
-                qk_norm=qk_norm,
-                activation_fn=activation_fn,
-                context_pre_only=i == num_layers - 1,
-            ) for i in range(num_layers)
-        ])
-
-        self.norm_out = AdaLayerNormContinuous(
-            inner_dim,
-            inner_dim,
-            elementwise_affine=False,
-            eps=1e-6,
-            norm_type="layer_norm",
-        )
-        self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels)
-
-        self.gradient_checkpointing = False
-
-    def _set_gradient_checkpointing(self, module, value=False):
-        if hasattr(module, "gradient_checkpointing"):
-            module.gradient_checkpointing = value
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        timestep: torch.LongTensor,
-        encoder_attention_mask: torch.Tensor,
-        output_features=False,
-        output_features_stride=8,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        return_dict: bool = False,
-    ) -> torch.Tensor:
-        assert (return_dict is False), "return_dict is not supported in MochiTransformer3DModel"
-
-        if attention_kwargs is not None:
-            attention_kwargs = attention_kwargs.copy()
-            lora_scale = attention_kwargs.pop("scale", 1.0)
-        else:
-            lora_scale = 1.0
-
-        if USE_PEFT_BACKEND:
-            # weight the lora layers by setting `lora_scale` for each PEFT layer
-            scale_lora_layers(self, lora_scale)
-        else:
-            if (attention_kwargs is not None and attention_kwargs.get("scale", None) is not None):
-                logger.warning("Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.")
-
-        batch_size, num_channels, num_frames, height, width = hidden_states.shape
-        p = self.config.patch_size
-
-        post_patch_height = height // p
-        post_patch_width = width // p
-        # Peiyuan: This is hacked to force mochi to follow the behaviour of SD3 and Flux
-        timestep = 1000 - timestep
-        temb, encoder_hidden_states = self.time_embed(
-            timestep,
-            encoder_hidden_states,
-            encoder_attention_mask,
-            hidden_dtype=hidden_states.dtype,
-        )
-
-        hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)
-        hidden_states = self.patch_embed(hidden_states)
-        hidden_states = hidden_states.unflatten(0, (batch_size, -1)).flatten(1, 2)
-
-        image_rotary_emb = self.rope(
-            self.pos_frequencies,
-            num_frames,
-            post_patch_height,
-            post_patch_width,
-            device=hidden_states.device,
-            dtype=torch.float32,
-        )
-        attn_outputs_list = []
-        for i, block in enumerate(self.transformer_blocks):
-            if self.gradient_checkpointing:
-
-                def create_custom_forward(module):
-
-                    def custom_forward(*inputs):
-                        return module(*inputs)
-
-                    return custom_forward
-
-                ckpt_kwargs: Dict[str, Any] = ({"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {})
-                (
-                    hidden_states,
-                    encoder_hidden_states,
-                    attn_outputs,
-                ) = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(block),
-                    hidden_states,
-                    encoder_hidden_states,
-                    encoder_attention_mask,
-                    temb,
-                    image_rotary_emb,
-                    output_features,
-                    **ckpt_kwargs,
-                )
-            else:
-                hidden_states, encoder_hidden_states, attn_outputs = block(
-                    hidden_states=hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_attention_mask=encoder_attention_mask,
-                    temb=temb,
-                    image_rotary_emb=image_rotary_emb,
-                    output_attn=output_features,
-                )
-            if i % output_features_stride == 0:
-                attn_outputs_list.append(attn_outputs)
-
-        hidden_states = self.norm_out(hidden_states, temb)
-        hidden_states = self.proj_out(hidden_states)
-
-        hidden_states = hidden_states.reshape(batch_size, num_frames, post_patch_height, post_patch_width, p, p, -1)
-        hidden_states = hidden_states.permute(0, 6, 1, 2, 4, 3, 5)
-        output = hidden_states.reshape(batch_size, -1, num_frames, height, width)
-
-        if USE_PEFT_BACKEND:
-            # remove `lora_scale` from each PEFT layer
-            unscale_lora_layers(self, lora_scale)
-
-        if not output_features:
-            attn_outputs_list = None
-        else:
-            attn_outputs_list = torch.stack(attn_outputs_list, dim=0)
-        # Peiyuan: This is hacked to force mochi to follow the behaviour of SD3 and Flux
-        return (-output, attn_outputs_list)

fastvideo/models/mochi_hf/norm.py

@@ -1,128 +0,0 @@
-# Copyright 2024 The Genmo team and The HuggingFace Team.
-# All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Tuple
-
-import torch
-import torch.nn as nn
-
-
-class MochiModulatedRMSNorm(nn.Module):
-
-    def __init__(self, eps: float):
-        super().__init__()
-
-        self.eps = eps
-
-    def forward(self, hidden_states, scale=None):
-        hidden_states_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
-        if scale is not None:
-            hidden_states = hidden_states * scale
-
-        hidden_states = hidden_states.to(hidden_states_dtype)
-
-        return hidden_states
-
-
-class MochiRMSNorm(nn.Module):
-
-    def __init__(self, dim, eps: float, elementwise_affine=True):
-        super().__init__()
-
-        self.eps = eps
-        if elementwise_affine:
-            self.weight = nn.Parameter(torch.ones(dim))
-        else:
-            self.weight = None
-
-    def forward(self, hidden_states):
-        hidden_states_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
-        if self.weight is not None:
-            # convert into half-precision if necessary
-            if self.weight.dtype in [torch.float16, torch.bfloat16]:
-                hidden_states = hidden_states.to(self.weight.dtype)
-            hidden_states = hidden_states * self.weight
-        hidden_states = hidden_states.to(hidden_states_dtype)
-
-        return hidden_states
-
-
-class MochiLayerNormContinuous(nn.Module):
-
-    def __init__(
-        self,
-        embedding_dim: int,
-        conditioning_embedding_dim: int,
-        eps=1e-5,
-        bias=True,
-    ):
-        super().__init__()
-
-        # AdaLN
-        self.silu = nn.SiLU()
-        self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias)
-        self.norm = MochiModulatedRMSNorm(eps=eps)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        conditioning_embedding: torch.Tensor,
-    ) -> torch.Tensor:
-        input_dtype = x.dtype
-
-        # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT)
-        scale = self.linear_1(self.silu(conditioning_embedding).to(x.dtype))
-        x = self.norm(x, (1 + scale.unsqueeze(1).to(torch.float32)))
-
-        return x.to(input_dtype)
-
-
-class MochiRMSNormZero(nn.Module):
-    r"""
-    Adaptive RMS Norm used in Mochi.
-    Parameters:
-        embedding_dim (`int`): The size of each embedding vector.
-    """
-
-    def __init__(
-        self,
-        embedding_dim: int,
-        hidden_dim: int,
-        eps: float = 1e-5,
-        elementwise_affine: bool = False,
-    ) -> None:
-        super().__init__()
-
-        self.silu = nn.SiLU()
-        self.linear = nn.Linear(embedding_dim, hidden_dim)
-        self.norm = MochiModulatedRMSNorm(eps=eps)
-
-    def forward(self, hidden_states: torch.Tensor,
-                emb: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        hidden_states_dtype = hidden_states.dtype
-
-        emb = self.linear(self.silu(emb))
-        scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1)
-
-        hidden_states = self.norm(hidden_states, (1 + scale_msa[:, None].to(torch.float32)))
-        hidden_states = hidden_states.to(hidden_states_dtype)
-
-        return hidden_states, gate_msa, scale_mlp, gate_mlp

fastvideo/models/mochi_hf/pipeline_mochi.py

@@ -1,757 +0,0 @@
-# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import copy
-import inspect
-from typing import Any, Callable, Dict, List, Optional, Union
-
-import numpy as np
-import torch
-from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
-from diffusers.loaders import Mochi1LoraLoaderMixin
-from diffusers.models.autoencoders import AutoencoderKL
-from diffusers.pipelines.mochi.pipeline_output import MochiPipelineOutput
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
-from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
-from diffusers.utils.torch_utils import randn_tensor
-from diffusers.video_processor import VideoProcessor
-from einops import rearrange
-from transformers import T5EncoderModel, T5TokenizerFast
-
-from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
-from fastvideo.utils.communications import all_gather
-from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-
-if is_torch_xla_available():
-    import torch_xla.core.xla_model as xm
-
-    XLA_AVAILABLE = True
-else:
-    XLA_AVAILABLE = False
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```py
-        >>> import torch
-        >>> from diffusers import MochiPipeline
-        >>> from diffusers.utils import export_to_video
-
-        >>> pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview", torch_dtype=torch.bfloat16)
-        >>> pipe.to("cuda")
-        >>> prompt = "Close-up of a chameleon's eye, with its scaly skin changing color. Ultra high resolution 4k."
-        >>> frames = pipe(prompt, num_inference_steps=28, guidance_scale=3.5).frames[0]
-        >>> export_to_video(frames, "mochi.mp4")
-        ```
-"""
-
-
-def calculate_shift(
-    image_seq_len,
-    base_seq_len: int = 256,
-    max_seq_len: int = 4096,
-    base_shift: float = 0.5,
-    max_shift: float = 1.16,
-):
-    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
-    b = base_shift - m * base_seq_len
-    mu = image_seq_len * m + b
-    return mu
-
-
-# from: https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77
-def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None):
-    if linear_steps is None:
-        linear_steps = num_steps // 2
-    linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)]
-    threshold_noise_step_diff = linear_steps - threshold_noise * num_steps
-    quadratic_steps = num_steps - linear_steps
-    quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
-    linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2)
-    const = quadratic_coef * (linear_steps**2)
-    quadratic_sigma_schedule = [
-        quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, num_steps)
-    ]
-    sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule
-    sigma_schedule = [1.0 - x for x in sigma_schedule]
-    return sigma_schedule
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-    r"""
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`List[int]`, *optional*):
-            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
-            `num_inference_steps` and `sigmas` must be `None`.
-        sigmas (`List[float]`, *optional*):
-            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
-            `num_inference_steps` and `timesteps` must be `None`.
-
-    Returns:
-        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
-        second element is the number of inference steps.
-    """
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler.")
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler.")
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-
-
-class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
-    r"""
-    The mochi pipeline for text-to-video generation.
-
-    Reference: https://github.com/genmoai/models
-
-    Args:
-        transformer ([`MochiTransformer3DModel`]):
-            Conditional Transformer architecture to denoise the encoded video latents.
-        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
-            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
-        vae ([`AutoencoderKL`]):
-            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
-        text_encoder ([`T5EncoderModel`]):
-            [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
-            the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
-        tokenizer (`CLIPTokenizer`):
-            Tokenizer of class
-            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
-        tokenizer (`T5TokenizerFast`):
-            Second Tokenizer of class
-            [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
-    """
-
-    model_cpu_offload_seq = "text_encoder->transformer->vae"
-    _optional_components = []
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
-
-    def __init__(
-        self,
-        scheduler: FlowMatchEulerDiscreteScheduler,
-        vae: AutoencoderKL,
-        text_encoder: T5EncoderModel,
-        tokenizer: T5TokenizerFast,
-        transformer: MochiTransformer3DModel,
-    ):
-        super().__init__()
-
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-            transformer=transformer,
-            scheduler=scheduler,
-        )
-        self.vae_spatial_scale_factor = 8
-        self.vae_temporal_scale_factor = 6
-        self.patch_size = 2
-
-        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_scale_factor)
-        self.tokenizer_max_length = (self.tokenizer.model_max_length
-                                     if hasattr(self, "tokenizer") and self.tokenizer is not None else 77)
-        self.default_height = 480
-        self.default_width = 848
-
-    # Adapted from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline._get_t5_prompt_embeds
-    def _get_t5_prompt_embeds(
-        self,
-        prompt: Union[str, List[str]] = None,
-        num_videos_per_prompt: int = 1,
-        max_sequence_length: int = 256,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        device = device or self._execution_device
-        dtype = dtype or self.text_encoder.dtype
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt)
-
-        text_inputs = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=max_sequence_length,
-            truncation=True,
-            add_special_tokens=True,
-            return_tensors="pt",
-        )
-        text_input_ids = text_inputs.input_ids
-        prompt_attention_mask = text_inputs.attention_mask
-        prompt_attention_mask = prompt_attention_mask.bool().to(device)
-
-        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
-
-        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
-            removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1:-1])
-            logger.warning("The following part of your input was truncated because `max_sequence_length` is set to "
-                           f" {max_sequence_length} tokens: {removed_text}")
-
-        prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)[0]
-        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
-
-        # duplicate text embeddings for each generation per prompt, using mps friendly method
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
-
-        prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
-        prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1)
-
-        return prompt_embeds, prompt_attention_mask
-
-    # Adapted from diffusers.pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipeline.encode_prompt
-    def encode_prompt(
-        self,
-        prompt: Union[str, List[str]],
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        do_classifier_free_guidance: bool = True,
-        num_videos_per_prompt: int = 1,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        prompt_attention_mask: Optional[torch.Tensor] = None,
-        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
-        max_sequence_length: int = 256,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        r"""
-        Encodes the prompt into text encoder hidden states.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                prompt to be encoded
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than `1`).
-            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
-                Whether to use classifier free guidance or not.
-            num_videos_per_prompt (`int`, *optional*, defaults to 1):
-                Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            device: (`torch.device`, *optional*):
-                torch device
-            dtype: (`torch.dtype`, *optional*):
-                torch dtype
-        """
-        device = device or self._execution_device
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        if prompt is not None:
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        if prompt_embeds is None:
-            prompt_embeds, prompt_attention_mask = self._get_t5_prompt_embeds(
-                prompt=prompt,
-                num_videos_per_prompt=num_videos_per_prompt,
-                max_sequence_length=max_sequence_length,
-                device=device,
-                dtype=dtype,
-            )
-
-        if do_classifier_free_guidance and negative_prompt_embeds is None:
-            negative_prompt = negative_prompt or ""
-            negative_prompt = (batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt)
-
-            if prompt is not None and type(prompt) is not type(negative_prompt):
-                raise TypeError(
-                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
-                    f" {type(prompt)}.")
-            elif batch_size != len(negative_prompt):
-                raise ValueError(
-                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
-                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
-                    " the batch size of `prompt`.")
-
-            (
-                negative_prompt_embeds,
-                negative_prompt_attention_mask,
-            ) = self._get_t5_prompt_embeds(
-                prompt=negative_prompt,
-                num_videos_per_prompt=num_videos_per_prompt,
-                max_sequence_length=max_sequence_length,
-                device=device,
-                dtype=dtype,
-            )
-
-        return (
-            prompt_embeds,
-            prompt_attention_mask,
-            negative_prompt_embeds,
-            negative_prompt_attention_mask,
-        )
-
-    def check_inputs(
-        self,
-        prompt,
-        height,
-        width,
-        callback_on_step_end_tensor_inputs=None,
-        prompt_embeds=None,
-        negative_prompt_embeds=None,
-        prompt_attention_mask=None,
-        negative_prompt_attention_mask=None,
-    ):
-        if height % 8 != 0 or width % 8 != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
-
-        if callback_on_step_end_tensor_inputs is not None and not all(k in self._callback_tensor_inputs
-                                                                      for k in callback_on_step_end_tensor_inputs):
-            raise ValueError(
-                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
-            )
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two.")
-        elif prompt is None and prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.")
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-
-        if prompt_embeds is not None and prompt_attention_mask is None:
-            raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
-
-        if (negative_prompt_embeds is not None and negative_prompt_attention_mask is None):
-            raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
-
-        if prompt_embeds is not None and negative_prompt_embeds is not None:
-            if prompt_embeds.shape != negative_prompt_embeds.shape:
-                raise ValueError(
-                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
-                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
-                    f" {negative_prompt_embeds.shape}.")
-            if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
-                raise ValueError(
-                    "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
-                    f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
-                    f" {negative_prompt_attention_mask.shape}.")
-
-    def enable_vae_slicing(self):
-        r"""
-        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
-        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
-        """
-        self.vae.enable_slicing()
-
-    def disable_vae_slicing(self):
-        r"""
-        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_slicing()
-
-    def enable_vae_tiling(self):
-        r"""
-        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
-        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
-        processing larger images.
-        """
-        self.vae.enable_tiling()
-
-    def disable_vae_tiling(self):
-        r"""
-        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_tiling()
-
-    def prepare_latents(
-        self,
-        batch_size,
-        num_channels_latents,
-        height,
-        width,
-        num_frames,
-        dtype,
-        device,
-        generator,
-        latents=None,
-    ):
-        height = height // self.vae_spatial_scale_factor
-        width = width // self.vae_spatial_scale_factor
-        num_frames = (num_frames - 1) // self.vae_temporal_scale_factor + 1
-
-        shape = (batch_size, num_channels_latents, num_frames, height, width)
-
-        if latents is not None:
-            return latents.to(device=device, dtype=dtype)
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators.")
-
-        latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
-        latents = latents.to(dtype)
-        return latents
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def do_classifier_free_guidance(self):
-        return self._guidance_scale > 1.0
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @property
-    def attention_kwargs(self):
-        return self._attention_kwargs
-
-    @property
-    def interrupt(self):
-        return self._interrupt
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        num_frames: int = 19,
-        num_inference_steps: int = 64,
-        timesteps: List[int] = None,
-        guidance_scale: float = 4.5,
-        num_videos_per_prompt: Optional[int] = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        prompt_attention_mask: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
-        max_sequence_length: int = 256,
-        return_all_states=False,
-    ):
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The height in pixels of the generated image. This is set to 1024 by default for the best results.
-            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The width in pixels of the generated image. This is set to 1024 by default for the best results.
-            num_frames (`int`, defaults to 16):
-                The number of video frames to generate
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            timesteps (`List[int]`, *optional*):
-                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
-                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
-                passed will be used. Must be in descending order.
-            guidance_scale (`float`, defaults to `4.5`):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality.
-            num_videos_per_prompt (`int`, *optional*, defaults to 1):
-                The number of videos to generate per prompt.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will ge generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            prompt_attention_mask (`torch.Tensor`, *optional*):
-                Pre-generated attention mask for text embeddings.
-            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not
-                provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
-            negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
-                Pre-generated attention mask for negative text embeddings.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.mochi.MochiPipelineOutput`] instead of a plain tuple.
-            attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            callback_on_step_end (`Callable`, *optional*):
-                A function that calls at the end of each denoising steps during the inference. The function is called
-                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
-                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
-                `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int` defaults to `256`):
-                Maximum sequence length to use with the `prompt`.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.mochi.MochiPipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`~pipelines.mochi.MochiPipelineOutput`] is returned, otherwise a `tuple`
-                is returned where the first element is a list with the generated images.
-        """
-
-        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
-            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
-
-        height = height or self.default_height
-        width = width or self.default_width
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt=prompt,
-            height=height,
-            width=width,
-            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            prompt_attention_mask=prompt_attention_mask,
-            negative_prompt_attention_mask=negative_prompt_attention_mask,
-        )
-
-        self._guidance_scale = guidance_scale
-        self._attention_kwargs = attention_kwargs
-        self._interrupt = False
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        device = self._execution_device
-
-        # 3. Prepare text embeddings
-        (
-            prompt_embeds,
-            prompt_attention_mask,
-            negative_prompt_embeds,
-            negative_prompt_attention_mask,
-        ) = self.encode_prompt(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            do_classifier_free_guidance=self.do_classifier_free_guidance,
-            num_videos_per_prompt=num_videos_per_prompt,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            prompt_attention_mask=prompt_attention_mask,
-            negative_prompt_attention_mask=negative_prompt_attention_mask,
-            max_sequence_length=max_sequence_length,
-            device=device,
-        )
-        if self.do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
-            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
-
-        # 4. Prepare latent variables
-        num_channels_latents = self.transformer.config.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_videos_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            num_frames,
-            prompt_embeds.dtype,
-            device,
-            generator,
-            latents,
-        )
-        world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
-        if get_sequence_parallel_state():
-            latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
-            latents = latents[:, :, rank, :, :, :]
-
-        original_noise = copy.deepcopy(latents)
-        # 5. Prepare timestep
-        # from https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77
-        threshold_noise = 0.025
-        sigmas = linear_quadratic_schedule(num_inference_steps, threshold_noise)
-        sigmas = np.array(sigmas)
-        # check if of type FlowMatchEulerDiscreteScheduler
-        if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
-            timesteps, num_inference_steps = retrieve_timesteps(
-                self.scheduler,
-                num_inference_steps,
-                device,
-                timesteps,
-                sigmas,
-            )
-        else:
-            timesteps, num_inference_steps = retrieve_timesteps(
-                self.scheduler,
-                num_inference_steps,
-                device,
-            )
-        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
-        self._num_timesteps = len(timesteps)
-
-        # 6. Denoising loop
-        self._progress_bar_config = {"disable": nccl_info.rank_within_group != 0}
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                latent_model_input = (torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents)
-                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype)
-
-                noise_pred = self.transformer(
-                    hidden_states=latent_model_input,
-                    encoder_hidden_states=prompt_embeds,
-                    timestep=timestep,
-                    encoder_attention_mask=prompt_attention_mask,
-                    attention_kwargs=attention_kwargs,
-                    return_dict=False,
-                )[0]
-
-                # Mochi CFG + Sampling runs in FP32
-                noise_pred = noise_pred.to(torch.float32)
-                if self.do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents_dtype = latents.dtype
-                latents = self.scheduler.step(noise_pred, t, latents.to(torch.float32), return_dict=False)[0]
-                latents = latents.to(latents_dtype)
-
-                if latents.dtype != latents_dtype:
-                    if torch.backends.mps.is_available():
-                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
-                        latents = latents.to(latents_dtype)
-
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-
-                if XLA_AVAILABLE:
-                    xm.mark_step()
-
-        if get_sequence_parallel_state():
-            latents = all_gather(latents, dim=2)
-            # latents_shape = list(latents.shape)
-            # full_shape = [latents_shape[0] * world_size] + latents_shape[1:]
-            # all_latents = torch.zeros(full_shape, dtype=latents.dtype, device=latents.device)
-            # torch.distributed.all_gather_into_tensor(all_latents, latents)
-            # latents_list = list(all_latents.chunk(world_size, dim=0))
-            # latents = torch.cat(latents_list, dim=2)
-
-        if output_type == "latent":
-            video = latents
-        else:
-            # unscale/denormalize the latents
-            # denormalize with the mean and std if available and not None
-            has_latents_mean = (hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None)
-            has_latents_std = (hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None)
-            if has_latents_mean and has_latents_std:
-                latents_mean = (torch.tensor(self.vae.config.latents_mean).view(1, 12, 1, 1,
-                                                                                1).to(latents.device, latents.dtype))
-                latents_std = (torch.tensor(self.vae.config.latents_std).view(1, 12, 1, 1,
-                                                                              1).to(latents.device, latents.dtype))
-                latents = (latents * latents_std / self.vae.config.scaling_factor + latents_mean)
-            else:
-                latents = latents / self.vae.config.scaling_factor
-
-            video = self.vae.decode(latents, return_dict=False)[0]
-            video = self.video_processor.postprocess_video(video, output_type=output_type)
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-        if return_all_states:
-            # Pay extra attention here:
-            # prompt_embeds with shape torch.Size([2, 256]), where prompt_embeds[1] is the prompt_embeds for the actual prompt
-            # prompt_embeds[0] is for negative prompt
-            return original_noise, video, latents, prompt_embeds, prompt_attention_mask
-
-        if not return_dict:
-            return (video, )
-
-        return MochiPipelineOutput(frames=video)

fastvideo/models/parameter.py

@@ -431,4 +431,4 @@ def _adjust_shard_indexes_for_packing(shard_size, shard_offset, packed_factor,
             shard_size=shard_size,
             shard_offset=shard_offset,
             marlin_tile_size=marlin_tile_size)
-    return shard_size, shard_offset
+    return shard_size, shard_offset

fastvideo/models/registry.py

@@ -13,18 +13,34 @@ import cloudpickle
 from torch import nn
 from fastvideo.logger import logger
 
-# huggingface class name: (fastvideo module name, fastvideo class name)
-_TEXT_TO_VIDEO_MODELS = {
-    "HunyuanVideoTransformer3DModel": ("hunyuanvideo", "HunyuanVideoDiT"),
+# huggingface class name: (component_name, fastvideo module name, fastvideo class name)
+_TEXT_TO_VIDEO_DIT_MODELS = {
+    "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"),
 }
 
-_IMAGE_TO_VIDEO_MODELS = {
-    
+_IMAGE_TO_VIDEO_DIT_MODELS = {
+    # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"),
+}
+
+_TEXT_ENCODER_MODELS = {
+    "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"),
+    "LlamaModel":    ("encoders", "llama", "LlamaModel"),
+}
+
+_IMAGE_ENCODER_MODELS = {
+    # "HunyuanVideoTransformer3DModel": ("image_encoder", "hunyuanvideo", "HunyuanVideoImageEncoder"),
+}
+
+_VAE_MODELS = {
+    "AutoencoderKLHunyuanVideo": ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"),
 }
 
 _FAST_VIDEO_MODELS = {
-    **_TEXT_TO_VIDEO_MODELS,
-    **_IMAGE_TO_VIDEO_MODELS,
+    **_TEXT_TO_VIDEO_DIT_MODELS,
+    **_IMAGE_TO_VIDEO_DIT_MODELS,
+    **_TEXT_ENCODER_MODELS,
+    **_IMAGE_ENCODER_MODELS,
+    **_VAE_MODELS,
 }
 
 _SUBPROCESS_COMMAND = [
@@ -111,6 +127,7 @@ class _LazyRegisteredModel(_BaseRegisteredModel):
     Represents a model that has not been imported in the main process.
     """
     module_name: str
+    component_name: str
     class_name: str
 
     # Performed in another process to avoid initializing CUDA
@@ -128,7 +145,7 @@ def _try_load_model_cls(
     model_arch: str,
     model: _BaseRegisteredModel,
 ) -> Optional[Type[nn.Module]]:
-    from vllm.platforms import current_platform
+    from fastvideo.platforms import current_platform
     current_platform.verify_model_arch(model_arch)
     try:
         return model.load_model_cls()
@@ -266,8 +283,9 @@ class _ModelRegistry:
 ModelRegistry = _ModelRegistry({
     model_arch:
     _LazyRegisteredModel(
-        module_name=f"fastvideo.models.dits.{mod_relname}",
+        module_name=f"fastvideo.models.{component_name}.{mod_relname}",
+        component_name=component_name,
         class_name=cls_name,
     )
-    for model_arch, (mod_relname, cls_name) in _FAST_VIDEO_MODELS.items()
+    for model_arch, (component_name, mod_relname, cls_name) in _FAST_VIDEO_MODELS.items()
 })

fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py

@@ -1,3 +1,22 @@
+# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+
 from dataclasses import dataclass
 from typing import Optional, Tuple, Union
 
@@ -36,6 +55,8 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
         timestep_spacing (`str`, defaults to `"linspace"`):
             The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
             Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        shift (`float`, defaults to 1.0):
+            The shift value for the timestep schedule.
         reverse (`bool`, defaults to `True`):
             Whether to reverse the timestep schedule.
     """
@@ -47,9 +68,10 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
     def __init__(
         self,
         num_train_timesteps: int = 1000,
-        reverse: bool = False,
+        shift: float = 1.0,
+        reverse: bool = True,
         solver: str = "euler",
-        device: Union[str, torch.device] = None,
+        n_tokens: Optional[int] = None,
     ):
         sigmas = torch.linspace(1, 0, num_train_timesteps + 1)
 
@@ -63,8 +85,6 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
         self._step_index = None
         self._begin_index = None
 
-        self.device = device
-
         self.supported_solver = ["euler"]
         if solver not in self.supported_solver:
             raise ValueError(f"Solver {solver} not supported. Supported solvers: {self.supported_solver}")
@@ -100,8 +120,8 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
     def set_timesteps(
         self,
         num_inference_steps: int,
-        time_shift: float = 13.0,
         device: Union[str, torch.device] = None,
+        n_tokens: int = None,
     ):
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -114,17 +134,16 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
             n_tokens (`int`, *optional*):
                 Number of tokens in the input sequence.
         """
-        device = device or self.device
         self.num_inference_steps = num_inference_steps
 
-        sigmas = torch.linspace(1, 0, num_inference_steps + 1, device=device)
-        sigmas = self.sd3_time_shift(sigmas, time_shift)
+        sigmas = torch.linspace(1, 0, num_inference_steps + 1)
+        sigmas = self.sd3_time_shift(sigmas)
 
         if not self.config.reverse:
             sigmas = 1 - sigmas
 
         self.sigmas = sigmas
-        self.timesteps = sigmas[:-1]
+        self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(dtype=torch.float32, device=device)
 
         # Reset step index
         self._step_index = None
@@ -154,15 +173,15 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
     def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
         return sample
 
-    def sd3_time_shift(self, t: torch.Tensor, time_shift: float = 13.0):
-        return (time_shift * t) / (1 + (time_shift - 1) * t)
+    def sd3_time_shift(self, t: torch.Tensor):
+        return (self.config.shift * t) / (1 + (self.config.shift - 1) * t)
 
     def step(
         self,
         model_output: torch.FloatTensor,
         timestep: Union[float, torch.FloatTensor],
         sample: torch.FloatTensor,
-        return_dict: bool = False,
+        return_dict: bool = True,
     ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]:
         """
         Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
@@ -212,9 +231,9 @@ class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
         self._step_index += 1
 
         if not return_dict:
-            return prev_sample
+            return (prev_sample, )
 
         return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample)
 
     def __len__(self):
-        return self.config.num_train_timesteps
+        return self.config.num_train_timesteps

fastvideo/models/stepvideo/__init__.py

@@ -1,7 +0,0 @@
-import os
-
-os.environ["NCCL_DEBUG"] = "ERROR"
-
-from .diffusion.scheduler import *
-from .diffusion.video_pipeline import *
-from .modules.model import *

fastvideo/models/stepvideo/__version__.py

@@ -1 +0,0 @@
-__version__ = "0.1.0"

fastvideo/models/stepvideo/config.py

@@ -1,174 +0,0 @@
-import argparse
-
-
-def parse_args(namespace=None):
-    parser = argparse.ArgumentParser(description="StepVideo inference script")
-
-    parser = add_extra_models_args(parser)
-    parser = add_denoise_schedule_args(parser)
-    parser = add_inference_args(parser)
-    parser = add_parallel_args(parser)
-
-    args = parser.parse_args(namespace=namespace)
-
-    return args
-
-
-def add_extra_models_args(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group(title="Extra models args, including vae, text encoders and tokenizers)")
-
-    group.add_argument(
-        "--vae_url",
-        type=str,
-        default='127.0.0.1',
-        help="vae url.",
-    )
-    group.add_argument(
-        "--caption_url",
-        type=str,
-        default='127.0.0.1',
-        help="caption url.",
-    )
-
-    return parser
-
-
-def add_denoise_schedule_args(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group(title="Denoise schedule args")
-
-    # Flow Matching
-    group.add_argument(
-        "--time_shift",
-        type=float,
-        default=7.0,
-        help="Shift factor for flow matching schedulers.",
-    )
-    group.add_argument(
-        "--flow_reverse",
-        action="store_true",
-        help="If reverse, learning/sampling from t=1 -> t=0.",
-    )
-    group.add_argument(
-        "--flow_solver",
-        type=str,
-        default="euler",
-        help="Solver for flow matching.",
-    )
-
-    return parser
-
-
-def add_inference_args(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group(title="Inference args")
-
-    # ======================== Model loads ========================
-    group.add_argument(
-        "--model_dir",
-        type=str,
-        default="./ckpts",
-        help="Root path of all the models, including t2v models and extra models.",
-    )
-    group.add_argument(
-        "--model_resolution",
-        type=str,
-        default="540p",
-        choices=["540p"],
-        help="Root path of all the models, including t2v models and extra models.",
-    )
-    group.add_argument(
-        "--use-cpu-offload",
-        action="store_true",
-        help="Use CPU offload for the model load.",
-    )
-
-    # ======================== Inference general setting ========================
-    group.add_argument(
-        "--batch_size",
-        type=int,
-        default=1,
-        help="Batch size for inference and evaluation.",
-    )
-    group.add_argument(
-        "--infer_steps",
-        type=int,
-        default=50,
-        help="Number of denoising steps for inference.",
-    )
-    group.add_argument(
-        "--save_path",
-        type=str,
-        default="./results",
-        help="Path to save the generated samples.",
-    )
-    group.add_argument(
-        "--name_suffix",
-        type=str,
-        default="",
-        help="Suffix for the names of saved samples.",
-    )
-    group.add_argument(
-        "--num_videos",
-        type=int,
-        default=1,
-        help="Number of videos to generate for each prompt.",
-    )
-    # ---sample size---
-    group.add_argument(
-        "--num_frames",
-        type=int,
-        default=204,
-        help="How many frames to sample from a video. ",
-    )
-    group.add_argument(
-        "--height",
-        type=int,
-        default=544,
-        help="The height of video sample",
-    )
-    group.add_argument(
-        "--width",
-        type=int,
-        default=992,
-        help="The width of video sample",
-    )
-    # --- prompt ---
-    group.add_argument(
-        "--prompt",
-        type=str,
-        default=None,
-        help="Prompt for sampling during evaluation.",
-    )
-    group.add_argument("--seed", type=int, default=1234, help="Seed for evaluation.")
-
-    # Classifier-Free Guidance
-    group.add_argument("--pos_magic",
-                       type=str,
-                       default="超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。",
-                       help="Positive magic prompt for sampling.")
-    group.add_argument("--neg_magic",
-                       type=str,
-                       default="画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。",
-                       help="Negative magic prompt for sampling.")
-    group.add_argument("--cfg_scale", type=float, default=9.0, help="Classifier free guidance scale.")
-
-    return parser
-
-
-def add_parallel_args(parser: argparse.ArgumentParser):
-    group = parser.add_argument_group(title="Parallel args")
-
-    # ======================== Model loads ========================
-    group.add_argument(
-        "--ulysses_degree",
-        type=int,
-        default=8,
-        help="Ulysses degree.",
-    )
-    group.add_argument(
-        "--ring_degree",
-        type=int,
-        default=1,
-        help="Ulysses degree.",
-    )
-
-    return parser

fastvideo/models/stepvideo/diffusion/video_pipeline.py

@@ -1,325 +0,0 @@
-# Copyright 2025 StepFun Inc. All Rights Reserved.
-
-import asyncio
-import pickle
-from dataclasses import dataclass
-from typing import Dict, List, Optional, Union
-
-import numpy as np
-import torch
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from diffusers.utils import BaseOutput
-
-from fastvideo.models.stepvideo.diffusion.scheduler import FlowMatchDiscreteScheduler
-from fastvideo.models.stepvideo.modules.model import StepVideoModel
-from fastvideo.models.stepvideo.utils import VideoProcessor
-
-
-def call_api_gen(url, api, port=8080):
-    url = f"http://{url}:{port}/{api}-api"
-    import aiohttp
-
-    async def _fn(samples, *args, **kwargs):
-        if api == 'vae':
-            data = {
-                "samples": samples,
-            }
-        elif api == 'caption':
-            data = {
-                "prompts": samples,
-            }
-        else:
-            raise Exception(f"Not supported api: {api}...")
-
-        async with aiohttp.ClientSession() as sess:
-            data_bytes = pickle.dumps(data)
-            async with sess.get(url, data=data_bytes, timeout=12000) as response:
-                result = bytearray()
-                while not response.content.at_eof():
-                    chunk = await response.content.read(1024)
-                    result += chunk
-                response_data = pickle.loads(result)
-        return response_data
-
-    return _fn
-
-
-@dataclass
-class StepVideoPipelineOutput(BaseOutput):
-    video: Union[torch.Tensor, np.ndarray]
-
-
-class StepVideoPipeline(DiffusionPipeline):
-    r"""
-    Pipeline for text-to-video generation using StepVideo.
-
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
-    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-    Args:
-        transformer ([`StepVideoModel`]):
-            Conditional Transformer to denoise the encoded image latents.
-        scheduler ([`FlowMatchDiscreteScheduler`]):
-            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
-        vae_url:
-            remote vae server's url.
-        caption_url:
-            remote caption (stepllm and clip) server's url.
-    """
-
-    def __init__(
-        self,
-        transformer: StepVideoModel,
-        scheduler: FlowMatchDiscreteScheduler,
-        vae_url: str = '127.0.0.1',
-        caption_url: str = '127.0.0.1',
-        save_path: str = './results',
-        name_suffix: str = '',
-    ):
-        super().__init__()
-
-        self.register_modules(
-            transformer=transformer,
-            scheduler=scheduler,
-        )
-
-        self.vae_scale_factor_temporal = self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 8
-        self.vae_scale_factor_spatial = self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 16
-        self.video_processor = VideoProcessor(save_path, name_suffix)
-
-        self.vae_url = vae_url
-        self.caption_url = caption_url
-        self.setup_api(self.vae_url, self.caption_url)
-
-    def setup_api(self, vae_url, caption_url):
-        self.vae_url = vae_url
-        self.caption_url = caption_url
-        self.caption = call_api_gen(caption_url, 'caption')
-        self.vae = call_api_gen(vae_url, 'vae')
-        return self
-
-    def encode_prompt(
-        self,
-        prompt: str,
-        neg_magic: str = '',
-        pos_magic: str = '',
-    ):
-        device = self._execution_device
-        prompts = [prompt + pos_magic]
-        bs = len(prompts)
-        prompts += [neg_magic] * bs
-
-        data = asyncio.run(self.caption(prompts))
-        prompt_embeds, prompt_attention_mask, clip_embedding = data['y'].to(device), data['y_mask'].to(
-            device), data['clip_embedding'].to(device)
-
-        return prompt_embeds, clip_embedding, prompt_attention_mask
-
-    def decode_vae(self, samples):
-        samples = asyncio.run(self.vae(samples.cpu()))
-        return samples
-
-    def check_inputs(self, num_frames, width, height):
-        num_frames = max(num_frames // 17 * 17, 1)
-        width = max(width // 16 * 16, 16)
-        height = max(height // 16 * 16, 16)
-        return num_frames, width, height
-
-    def prepare_latents(
-        self,
-        batch_size: int,
-        num_channels_latents: 64,
-        height: int = 544,
-        width: int = 992,
-        num_frames: int = 204,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        if latents is not None:
-            return latents.to(device=device, dtype=dtype)
-
-        num_frames, width, height = self.check_inputs(num_frames, width, height)
-        shape = (
-            batch_size,
-            max(num_frames // 17 * 3, 1),
-            num_channels_latents,
-            int(height) // self.vae_scale_factor_spatial,
-            int(width) // self.vae_scale_factor_spatial,
-        )  # b,f,c,h,w
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators.")
-
-        if generator is None:
-            generator = torch.Generator(device=self._execution_device)
-
-        latents = torch.randn(shape, generator=generator, device=device, dtype=dtype)
-        return latents
-
-    @torch.inference_mode()
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        height: int = 544,
-        width: int = 992,
-        num_frames: int = 204,
-        num_inference_steps: int = 50,
-        guidance_scale: float = 9.0,
-        time_shift: float = 13.0,
-        neg_magic: str = "",
-        pos_magic: str = "",
-        num_videos_per_prompt: Optional[int] = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        output_type: Optional[str] = "mp4",
-        output_file_name: Optional[str] = "",
-        return_dict: bool = True,
-        mask_strategy: Optional[Dict[str, list]] = None,
-    ):
-        r"""
-        The call function to the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            height (`int`, defaults to `544`):
-                The height in pixels of the generated image.
-            width (`int`, defaults to `992`):
-                The width in pixels of the generated image.
-            num_frames (`int`, defaults to `204`):
-                The number of frames in the generated video.
-            num_inference_steps (`int`, defaults to `50`):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            guidance_scale (`float`, defaults to `9.0`):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality. 
-            num_videos_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
-                generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor is generated by sampling using the supplied random `generator`.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
-            output_file_name(`str`, *optional*`):
-                The output mp4 file name.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`StepVideoPipelineOutput`] instead of a plain tuple.
-
-        Examples:
-
-        Returns:
-            [`~StepVideoPipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`StepVideoPipelineOutput`] is returned, otherwise a `tuple` is returned
-                where the first element is a list with the generated images and the second element is a list of `bool`s
-                indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content.
-        """
-
-        # 1. Check inputs. Raise error if not correct
-        device = self._execution_device
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        do_classifier_free_guidance = guidance_scale > 1.0
-
-        # 3. Encode input prompt
-        prompt_embeds, prompt_embeds_2, prompt_attention_mask = self.encode_prompt(
-            prompt=prompt,
-            neg_magic=neg_magic,
-            pos_magic=pos_magic,
-        )
-
-        transformer_dtype = self.transformer.dtype
-        prompt_embeds = prompt_embeds.to(transformer_dtype)
-        prompt_attention_mask = prompt_attention_mask.to(transformer_dtype)
-        prompt_embeds_2 = prompt_embeds_2.to(transformer_dtype)
-
-        # 4. Prepare timesteps
-        self.scheduler.set_timesteps(num_inference_steps=num_inference_steps, time_shift=time_shift, device=device)
-
-        # 5. Prepare latent variables
-        num_channels_latents = self.transformer.config.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_videos_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            num_frames,
-            torch.bfloat16,
-            device,
-            generator,
-            latents,
-        )
-
-        def dict_to_3d_list(best_masks, t_max=50, l_max=48, h_max=48):
-            result = [[[None for _ in range(h_max)] for _ in range(l_max)] for _ in range(t_max)]
-            if best_masks is None:
-                return result
-            for key, value in best_masks.items():
-                timestep, layer, head = map(int, key.split('_'))
-                result[timestep][layer][head] = value
-            return result
-
-        mask_strategy = dict_to_3d_list(mask_strategy)
-
-        #best_mask_selections = None
-        # 7. Denoising loop
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(self.scheduler.timesteps):
-                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-                latent_model_input = latent_model_input.to(transformer_dtype)
-                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
-
-                noise_pred = self.transformer(
-                    hidden_states=latent_model_input,
-                    timestep=timestep,
-                    encoder_hidden_states=prompt_embeds,
-                    encoder_attention_mask=prompt_attention_mask,
-                    encoder_hidden_states_2=prompt_embeds_2,
-                    return_dict=False,
-                    mask_strategy=mask_strategy[i],
-                )
-                # perform guidance
-                if do_classifier_free_guidance:
-                    noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(model_output=noise_pred, timestep=t, sample=latents)
-
-                progress_bar.update()
-
-        if not torch.distributed.is_initialized() or int(torch.distributed.get_rank()) == 0:
-            if not output_type == "latent":
-                video = self.decode_vae(latents)
-                video = self.video_processor.postprocess_video(video,
-                                                               output_file_name=output_file_name,
-                                                               output_type=output_type)
-            else:
-                video = latents
-
-            # Offload all models
-            self.maybe_free_model_hooks()
-
-            if not return_dict:
-                return (video, )
-
-            return StepVideoPipelineOutput(video=video)

fastvideo/models/stepvideo/modules/attentions.py

@@ -1,96 +0,0 @@
-import torch
-import torch.nn as nn
-from einops import rearrange
-from flash_attn import flash_attn_func
-
-try:
-    from st_attn import sliding_tile_attention
-except ImportError:
-    print("Could not load Sliding Tile Attention.")
-    sliding_tile_attention = None
-
-from fastvideo.utils.communications import all_to_all_4D
-from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-
-
-class Attention(nn.Module):
-
-    def __init__(self):
-        super().__init__()
-
-    def attn_processor(self, attn_type):
-        if attn_type == 'torch':
-            return self.torch_attn_func
-        elif attn_type == 'parallel':
-            return self.parallel_attn_func
-        else:
-            raise Exception('Not supported attention type...')
-
-    def tile(self, x, sp_size):
-        x = rearrange(x, "b (sp t h w) head d -> b (t sp h w) head d", sp=sp_size, t=36 // sp_size, h=48, w=48)
-        return rearrange(x,
-                         "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
-                         n_t=6,
-                         n_h=6,
-                         n_w=6,
-                         ts_t=6,
-                         ts_h=8,
-                         ts_w=8)
-
-    def untile(self, x, sp_size):
-        x = rearrange(x,
-                      "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
-                      n_t=6,
-                      n_h=6,
-                      n_w=6,
-                      ts_t=6,
-                      ts_h=8,
-                      ts_w=8)
-        return rearrange(x, "b (t sp h w) head d -> b (sp t h w) head d", sp=sp_size, t=36 // sp_size, h=48, w=48)
-
-    def torch_attn_func(self, q, k, v, attn_mask=None, causal=False, drop_rate=0.0, **kwargs):
-
-        if attn_mask is not None and attn_mask.dtype != torch.bool:
-            attn_mask = attn_mask.to(q.dtype)
-
-        if attn_mask is not None and attn_mask.ndim == 3:  ## no head
-            n_heads = q.shape[2]
-            attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)
-
-        q, k, v = map(lambda x: rearrange(x, 'b s h d -> b h s d'), (q, k, v))
-        x = torch.nn.functional.scaled_dot_product_attention(q,
-                                                             k,
-                                                             v,
-                                                             attn_mask=attn_mask,
-                                                             dropout_p=drop_rate,
-                                                             is_causal=causal)
-        x = rearrange(x, 'b h s d -> b s h d')
-        return x
-
-    def parallel_attn_func(self, q, k, v, causal=False, mask_strategy=None, **kwargs):
-        if get_sequence_parallel_state():
-            q = all_to_all_4D(q, scatter_dim=2, gather_dim=1)
-            k = all_to_all_4D(k, scatter_dim=2, gather_dim=1)
-            v = all_to_all_4D(v, scatter_dim=2, gather_dim=1)
-
-        if mask_strategy[0] is not None:
-            q = self.tile(q, nccl_info.sp_size).transpose(1, 2).contiguous()
-            k = self.tile(k, nccl_info.sp_size).transpose(1, 2).contiguous()
-            v = self.tile(v, nccl_info.sp_size).transpose(1, 2).contiguous()
-
-            head_num = q.size(1)  # 48 // sp_size
-            current_rank = nccl_info.rank_within_group
-
-            start_head = current_rank * head_num
-            windows = [mask_strategy[head_idx + start_head] for head_idx in range(head_num)]
-
-            x = sliding_tile_attention(q, k, v, windows, 0, False).transpose(1, 2).contiguous()
-            x = self.untile(x, nccl_info.sp_size)
-        else:
-            x = flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False)
-
-        if get_sequence_parallel_state():
-            x = all_to_all_4D(x, scatter_dim=1, gather_dim=2)
-
-        x = x.to(q.dtype)
-        return x

fastvideo/models/stepvideo/modules/blocks.py

@@ -1,296 +0,0 @@
-# Copyright 2025 StepFun Inc. All Rights Reserved.
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-# ==============================================================================
-from typing import Optional
-
-import torch
-import torch.nn as nn
-from einops import rearrange
-
-from fastvideo.models.stepvideo.modules.attentions import Attention
-from fastvideo.models.stepvideo.modules.normalization import RMSNorm
-from fastvideo.models.stepvideo.modules.rope import RoPE3D
-
-
-class SelfAttention(Attention):
-
-    def __init__(self, hidden_dim, head_dim, bias=False, with_rope=True, with_qk_norm=True, attn_type='torch'):
-        super().__init__()
-        self.head_dim = head_dim
-        self.n_heads = hidden_dim // head_dim
-
-        self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=bias)
-        self.wo = nn.Linear(hidden_dim, hidden_dim, bias=bias)
-
-        self.with_rope = with_rope
-        self.with_qk_norm = with_qk_norm
-        if self.with_qk_norm:
-            self.q_norm = RMSNorm(head_dim, elementwise_affine=True)
-            self.k_norm = RMSNorm(head_dim, elementwise_affine=True)
-
-        if self.with_rope:
-            self.rope_3d = RoPE3D(freq=1e4, F0=1.0, scaling_factor=1.0)
-            self.rope_ch_split = [64, 32, 32]
-
-        self.core_attention = self.attn_processor(attn_type=attn_type)
-        self.parallel = attn_type == 'parallel'
-
-    def apply_rope3d(self, x, fhw_positions, rope_ch_split, parallel=True):
-        x = self.rope_3d(x, fhw_positions, rope_ch_split, parallel)
-        return x
-
-    def forward(self, x, cu_seqlens=None, max_seqlen=None, rope_positions=None, attn_mask=None, mask_strategy=None):
-        xqkv = self.wqkv(x)
-        xqkv = xqkv.view(*x.shape[:-1], self.n_heads, 3 * self.head_dim)
-
-        xq, xk, xv = torch.split(xqkv, [self.head_dim] * 3, dim=-1)  ## seq_len, n, dim
-
-        if self.with_qk_norm:
-            xq = self.q_norm(xq)
-            xk = self.k_norm(xk)
-
-        if self.with_rope:
-            xq = self.apply_rope3d(xq, rope_positions, self.rope_ch_split, parallel=self.parallel)
-            xk = self.apply_rope3d(xk, rope_positions, self.rope_ch_split, parallel=self.parallel)
-
-        output = self.core_attention(xq,
-                                     xk,
-                                     xv,
-                                     cu_seqlens=cu_seqlens,
-                                     max_seqlen=max_seqlen,
-                                     attn_mask=attn_mask,
-                                     mask_strategy=mask_strategy)
-        output = rearrange(output, 'b s h d -> b s (h d)')
-        output = self.wo(output)
-
-        return output
-
-
-class CrossAttention(Attention):
-
-    def __init__(self, hidden_dim, head_dim, bias=False, with_qk_norm=True, attn_type='torch'):
-        super().__init__()
-        self.head_dim = head_dim
-        self.n_heads = hidden_dim // head_dim
-
-        self.wq = nn.Linear(hidden_dim, hidden_dim, bias=bias)
-        self.wkv = nn.Linear(hidden_dim, hidden_dim * 2, bias=bias)
-        self.wo = nn.Linear(hidden_dim, hidden_dim, bias=bias)
-
-        self.with_qk_norm = with_qk_norm
-        if self.with_qk_norm:
-            self.q_norm = RMSNorm(head_dim, elementwise_affine=True)
-            self.k_norm = RMSNorm(head_dim, elementwise_affine=True)
-
-        self.core_attention = self.attn_processor(attn_type=attn_type)
-
-    def forward(self, x: torch.Tensor, encoder_hidden_states: torch.Tensor, attn_mask=None):
-        xq = self.wq(x)
-        xq = xq.view(*xq.shape[:-1], self.n_heads, self.head_dim)
-
-        xkv = self.wkv(encoder_hidden_states)
-        xkv = xkv.view(*xkv.shape[:-1], self.n_heads, 2 * self.head_dim)
-
-        xk, xv = torch.split(xkv, [self.head_dim] * 2, dim=-1)  ## seq_len, n, dim
-
-        if self.with_qk_norm:
-            xq = self.q_norm(xq)
-            xk = self.k_norm(xk)
-
-        output = self.core_attention(xq, xk, xv, attn_mask=attn_mask)
-
-        output = rearrange(output, 'b s h d -> b s (h d)')
-        output = self.wo(output)
-
-        return output
-
-
-class GELU(nn.Module):
-    r"""
-    GELU activation function with tanh approximation support with `approximate="tanh"`.
-
-    Parameters:
-        dim_in (`int`): The number of channels in the input.
-        dim_out (`int`): The number of channels in the output.
-        approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation.
-        bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
-    """
-
-    def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True):
-        super().__init__()
-        self.proj = nn.Linear(dim_in, dim_out, bias=bias)
-        self.approximate = approximate
-
-    def gelu(self, gate: torch.Tensor) -> torch.Tensor:
-        return torch.nn.functional.gelu(gate, approximate=self.approximate)
-
-    def forward(self, hidden_states):
-        hidden_states = self.proj(hidden_states)
-        hidden_states = self.gelu(hidden_states)
-        return hidden_states
-
-
-class FeedForward(nn.Module):
-
-    def __init__(
-        self,
-        dim: int,
-        inner_dim: Optional[int] = None,
-        dim_out: Optional[int] = None,
-        mult: int = 4,
-        bias: bool = False,
-    ):
-        super().__init__()
-        inner_dim = dim * mult if inner_dim is None else inner_dim
-        dim_out = dim if dim_out is None else dim_out
-        self.net = nn.ModuleList([
-            GELU(dim, inner_dim, approximate="tanh", bias=bias),
-            nn.Identity(),
-            nn.Linear(inner_dim, dim_out, bias=bias)
-        ])
-
-    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
-        for module in self.net:
-            hidden_states = module(hidden_states)
-        return hidden_states
-
-
-def modulate(x, scale, shift):
-    x = x * (1 + scale) + shift
-    return x
-
-
-def gate(x, gate):
-    x = gate * x
-    return x
-
-
-class StepVideoTransformerBlock(nn.Module):
-    r"""
-    A basic Transformer block.
-
-    Parameters:
-        dim (`int`): The number of channels in the input and output.
-        num_attention_heads (`int`): The number of heads to use for multi-head attention.
-        attention_head_dim (`int`): The number of channels in each head.
-        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
-        cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
-        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
-        num_embeds_ada_norm (:
-            obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
-        attention_bias (:
-            obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
-        only_cross_attention (`bool`, *optional*):
-            Whether to use only cross-attention layers. In this case two cross attention layers are used.
-        double_self_attention (`bool`, *optional*):
-            Whether to use two self-attention layers. In this case no cross attention layers are used.
-        upcast_attention (`bool`, *optional*):
-            Whether to upcast the attention computation to float32. This is useful for mixed precision training.
-        norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
-            Whether to use learnable elementwise affine parameters for normalization.
-        norm_type (`str`, *optional*, defaults to `"layer_norm"`):
-            The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
-        final_dropout (`bool` *optional*, defaults to False):
-            Whether to apply a final dropout after the last feed-forward layer.
-        attention_type (`str`, *optional*, defaults to `"default"`):
-            The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`.
-        positional_embeddings (`str`, *optional*, defaults to `None`):
-            The type of positional embeddings to apply to.
-        num_positional_embeddings (`int`, *optional*, defaults to `None`):
-            The maximum number of positional embeddings to apply.
-    """
-
-    def __init__(self,
-                 dim: int,
-                 attention_head_dim: int,
-                 norm_eps: float = 1e-5,
-                 ff_inner_dim: Optional[int] = None,
-                 ff_bias: bool = False,
-                 attention_type: str = 'parallel'):
-        super().__init__()
-        self.dim = dim
-        self.norm1 = nn.LayerNorm(dim, eps=norm_eps)
-        self.attn1 = SelfAttention(dim,
-                                   attention_head_dim,
-                                   bias=False,
-                                   with_rope=True,
-                                   with_qk_norm=True,
-                                   attn_type=attention_type)
-
-        self.norm2 = nn.LayerNorm(dim, eps=norm_eps)
-        self.attn2 = CrossAttention(dim, attention_head_dim, bias=False, with_qk_norm=True, attn_type='torch')
-
-        self.ff = FeedForward(dim=dim, inner_dim=ff_inner_dim, dim_out=dim, bias=ff_bias)
-
-        self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
-
-    @torch.no_grad()
-    def forward(self,
-                q: torch.Tensor,
-                kv: Optional[torch.Tensor] = None,
-                timestep: Optional[torch.LongTensor] = None,
-                attn_mask=None,
-                rope_positions: list = None,
-                mask_strategy=None) -> torch.Tensor:
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (torch.clone(chunk) for chunk in (
-            self.scale_shift_table[None] + timestep.reshape(-1, 6, self.dim)).chunk(6, dim=1))
-
-        scale_shift_q = modulate(self.norm1(q), scale_msa, shift_msa)
-
-        attn_q = self.attn1(scale_shift_q, rope_positions=rope_positions, mask_strategy=mask_strategy)
-
-        q = gate(attn_q, gate_msa) + q
-
-        attn_q = self.attn2(q, kv, attn_mask)
-
-        q = attn_q + q
-
-        scale_shift_q = modulate(self.norm2(q), scale_mlp, shift_mlp)
-
-        ff_output = self.ff(scale_shift_q)
-
-        q = gate(ff_output, gate_mlp) + q
-
-        return q
-
-
-class PatchEmbed(nn.Module):
-    """2D Image to Patch Embedding"""
-
-    def __init__(
-        self,
-        patch_size=64,
-        in_channels=3,
-        embed_dim=768,
-        layer_norm=False,
-        flatten=True,
-        bias=True,
-    ):
-        super().__init__()
-
-        self.flatten = flatten
-        self.layer_norm = layer_norm
-
-        self.proj = nn.Conv2d(in_channels,
-                              embed_dim,
-                              kernel_size=(patch_size, patch_size),
-                              stride=patch_size,
-                              bias=bias)
-
-    def forward(self, latent):
-        latent = self.proj(latent).to(latent.dtype)
-        if self.flatten:
-            latent = latent.flatten(2).transpose(1, 2)  # BCHW -> BNC
-        if self.layer_norm:
-            latent = self.norm(latent)
-
-        return latent

fastvideo/models/stepvideo/modules/model.py

@@ -1,198 +0,0 @@
-# Copyright 2025 StepFun Inc. All Rights Reserved.
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-# ==============================================================================
-from typing import Dict, Optional
-
-import torch
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.models.modeling_utils import ModelMixin
-from einops import rearrange, repeat
-from torch import nn
-
-from fastvideo.models.stepvideo.modules.blocks import PatchEmbed, StepVideoTransformerBlock
-from fastvideo.models.stepvideo.modules.normalization import AdaLayerNormSingle, PixArtAlphaTextProjection
-from fastvideo.models.stepvideo.parallel import parallel_forward
-from fastvideo.models.stepvideo.utils import with_empty_init
-
-
-class StepVideoModel(ModelMixin, ConfigMixin):
-    _no_split_modules = ["StepVideoTransformerBlock", "PatchEmbed"]
-
-    @with_empty_init
-    @register_to_config
-    def __init__(
-        self,
-        num_attention_heads: int = 48,
-        attention_head_dim: int = 128,
-        in_channels: int = 64,
-        out_channels: Optional[int] = 64,
-        num_layers: int = 48,
-        dropout: float = 0.0,
-        patch_size: int = 1,
-        norm_type: str = "ada_norm_single",
-        norm_elementwise_affine: bool = False,
-        norm_eps: float = 1e-6,
-        use_additional_conditions: Optional[bool] = False,
-        caption_channels: Optional[int] | list | tuple = [6144, 1024],
-        attention_type: Optional[str] = "parallel",
-    ):
-        super().__init__()
-
-        # Set some common variables used across the board.
-        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
-        self.out_channels = in_channels if out_channels is None else out_channels
-
-        self.use_additional_conditions = use_additional_conditions
-
-        self.pos_embed = PatchEmbed(
-            patch_size=patch_size,
-            in_channels=self.config.in_channels,
-            embed_dim=self.inner_dim,
-        )
-
-        self.transformer_blocks = nn.ModuleList([
-            StepVideoTransformerBlock(dim=self.inner_dim,
-                                      attention_head_dim=self.config.attention_head_dim,
-                                      attention_type=attention_type) for _ in range(self.config.num_layers)
-        ])
-
-        # 3. Output blocks.
-        self.norm_out = nn.LayerNorm(self.inner_dim, eps=norm_eps, elementwise_affine=norm_elementwise_affine)
-        self.scale_shift_table = nn.Parameter(torch.randn(2, self.inner_dim) / self.inner_dim**0.5)
-        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels)
-        self.patch_size = patch_size
-
-        self.adaln_single = AdaLayerNormSingle(self.inner_dim, use_additional_conditions=self.use_additional_conditions)
-
-        if isinstance(self.config.caption_channels, int):
-            caption_channel = self.config.caption_channels
-        else:
-            caption_channel, clip_channel = self.config.caption_channels
-            self.clip_projection = nn.Linear(clip_channel, self.inner_dim)
-
-        self.caption_norm = nn.LayerNorm(caption_channel, eps=norm_eps, elementwise_affine=norm_elementwise_affine)
-
-        self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channel, hidden_size=self.inner_dim)
-
-        self.parallel = attention_type == 'parallel'
-
-    def patchfy(self, hidden_states):
-        hidden_states = rearrange(hidden_states, 'b f c h w -> (b f) c h w')
-        hidden_states = self.pos_embed(hidden_states)
-        return hidden_states
-
-    def prepare_attn_mask(self, encoder_attention_mask, encoder_hidden_states, q_seqlen):
-        kv_seqlens = encoder_attention_mask.sum(dim=1).int()
-        mask = torch.zeros([len(kv_seqlens), q_seqlen, max(kv_seqlens)],
-                           dtype=torch.bool,
-                           device=encoder_attention_mask.device)
-        encoder_hidden_states = encoder_hidden_states[:, :max(kv_seqlens)]
-        for i, kv_len in enumerate(kv_seqlens):
-            mask[i, :, :kv_len] = 1
-        return encoder_hidden_states, mask
-
-    @parallel_forward
-    def block_forward(self,
-                      hidden_states,
-                      encoder_hidden_states=None,
-                      timestep=None,
-                      rope_positions=None,
-                      attn_mask=None,
-                      parallel=True,
-                      mask_strategy=None):
-
-        for i, block in enumerate(self.transformer_blocks):
-            hidden_states = block(hidden_states,
-                                  encoder_hidden_states,
-                                  timestep=timestep,
-                                  attn_mask=attn_mask,
-                                  rope_positions=rope_positions,
-                                  mask_strategy=mask_strategy[i])
-
-        return hidden_states
-
-    @torch.inference_mode()
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        encoder_hidden_states_2: Optional[torch.Tensor] = None,
-        timestep: Optional[torch.LongTensor] = None,
-        added_cond_kwargs: Dict[str, torch.Tensor] = None,
-        encoder_attention_mask: Optional[torch.Tensor] = None,
-        fps: torch.Tensor = None,
-        return_dict: bool = True,
-        mask_strategy=None,
-    ):
-        assert hidden_states.ndim == 5
-        "hidden_states's shape should be (bsz, f, ch, h ,w)"
-
-        bsz, frame, _, height, width = hidden_states.shape
-        height, width = height // self.patch_size, width // self.patch_size
-
-        hidden_states = self.patchfy(hidden_states)
-        len_frame = hidden_states.shape[1]
-
-        if self.use_additional_conditions:
-            added_cond_kwargs = {
-                "resolution": torch.tensor([(height, width)] * bsz,
-                                           device=hidden_states.device,
-                                           dtype=hidden_states.dtype),
-                "nframe": torch.tensor([frame] * bsz, device=hidden_states.device, dtype=hidden_states.dtype),
-                "fps": fps
-            }
-        else:
-            added_cond_kwargs = {}
-
-        timestep, embedded_timestep = self.adaln_single(timestep, added_cond_kwargs=added_cond_kwargs)
-
-        encoder_hidden_states = self.caption_projection(self.caption_norm(encoder_hidden_states))
-
-        if encoder_hidden_states_2 is not None and hasattr(self, 'clip_projection'):
-            clip_embedding = self.clip_projection(encoder_hidden_states_2)
-            encoder_hidden_states = torch.cat([clip_embedding, encoder_hidden_states], dim=1)
-
-        hidden_states = rearrange(hidden_states, '(b f) l d->  b (f l) d', b=bsz, f=frame, l=len_frame).contiguous()
-        encoder_hidden_states, attn_mask = self.prepare_attn_mask(encoder_attention_mask,
-                                                                  encoder_hidden_states,
-                                                                  q_seqlen=frame * len_frame)
-
-        hidden_states = self.block_forward(hidden_states,
-                                           encoder_hidden_states,
-                                           timestep=timestep,
-                                           rope_positions=[frame, height, width],
-                                           attn_mask=attn_mask,
-                                           parallel=self.parallel,
-                                           mask_strategy=mask_strategy)
-
-        hidden_states = rearrange(hidden_states, 'b (f l) d -> (b f) l d', b=bsz, f=frame, l=len_frame)
-
-        embedded_timestep = repeat(embedded_timestep, 'b d -> (b f) d', f=frame).contiguous()
-
-        shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1)
-        hidden_states = self.norm_out(hidden_states)
-        # Modulation
-        hidden_states = hidden_states * (1 + scale) + shift
-        hidden_states = self.proj_out(hidden_states)
-
-        # unpatchify
-        hidden_states = hidden_states.reshape(shape=(-1, height, width, self.patch_size, self.patch_size,
-                                                     self.out_channels))
-
-        hidden_states = rearrange(hidden_states, 'n h w p q c -> n c h p w q')
-        output = hidden_states.reshape(shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size))
-
-        output = rearrange(output, '(b f) c h w -> b f c h w', f=frame)
-
-        if return_dict:
-            return {'x': output}
-        return output

fastvideo/models/stepvideo/modules/normalization.py

@@ -1,312 +0,0 @@
-import math
-from typing import Dict, Optional, Tuple
-
-import torch
-import torch.nn as nn
-
-
-class RMSNorm(nn.Module):
-
-    def __init__(
-        self,
-        dim: int,
-        elementwise_affine=True,
-        eps: float = 1e-6,
-        device=None,
-        dtype=None,
-    ):
-        """
-        Initialize the RMSNorm normalization layer.
-
-        Args:
-            dim (int): The dimension of the input tensor.
-            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
-
-        Attributes:
-            eps (float): A small value added to the denominator for numerical stability.
-            weight (nn.Parameter): Learnable scaling parameter.
-
-        """
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.eps = eps
-        if elementwise_affine:
-            self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
-
-    def _norm(self, x):
-        """
-        Apply the RMSNorm normalization to the input tensor.
-
-        Args:
-            x (torch.Tensor): The input tensor.
-
-        Returns:
-            torch.Tensor: The normalized tensor.
-
-        """
-        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
-
-    def forward(self, x):
-        """
-        Forward pass through the RMSNorm layer.
-
-        Args:
-            x (torch.Tensor): The input tensor.
-
-        Returns:
-            torch.Tensor: The output tensor after applying RMSNorm.
-
-        """
-        output = self._norm(x.float()).type_as(x)
-        if hasattr(self, "weight"):
-            output = output * self.weight
-        return output
-
-
-ACTIVATION_FUNCTIONS = {
-    "swish": nn.SiLU(),
-    "silu": nn.SiLU(),
-    "mish": nn.Mish(),
-    "gelu": nn.GELU(),
-    "relu": nn.ReLU(),
-}
-
-
-def get_activation(act_fn: str) -> nn.Module:
-    """Helper function to get activation function from string.
-
-    Args:
-        act_fn (str): Name of activation function.
-
-    Returns:
-        nn.Module: Activation function.
-    """
-
-    act_fn = act_fn.lower()
-    if act_fn in ACTIVATION_FUNCTIONS:
-        return ACTIVATION_FUNCTIONS[act_fn]
-    else:
-        raise ValueError(f"Unsupported activation function: {act_fn}")
-
-
-def get_timestep_embedding(
-    timesteps: torch.Tensor,
-    embedding_dim: int,
-    flip_sin_to_cos: bool = False,
-    downscale_freq_shift: float = 1,
-    scale: float = 1,
-    max_period: int = 10000,
-):
-    """
-    This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
-
-    :param timesteps: a 1-D Tensor of N indices, one per batch element.
-                      These may be fractional.
-    :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the
-    embeddings. :return: an [N x dim] Tensor of positional embeddings.
-    """
-    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
-
-    half_dim = embedding_dim // 2
-    exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
-    exponent = exponent / (half_dim - downscale_freq_shift)
-
-    emb = torch.exp(exponent)
-    emb = timesteps[:, None].float() * emb[None, :]
-
-    # scale embeddings
-    emb = scale * emb
-
-    # concat sine and cosine embeddings
-    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
-
-    # flip sine and cosine embeddings
-    if flip_sin_to_cos:
-        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
-
-    # zero pad
-    if embedding_dim % 2 == 1:
-        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
-    return emb
-
-
-class Timesteps(nn.Module):
-
-    def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float):
-        super().__init__()
-        self.num_channels = num_channels
-        self.flip_sin_to_cos = flip_sin_to_cos
-        self.downscale_freq_shift = downscale_freq_shift
-
-    def forward(self, timesteps):
-        t_emb = get_timestep_embedding(
-            timesteps,
-            self.num_channels,
-            flip_sin_to_cos=self.flip_sin_to_cos,
-            downscale_freq_shift=self.downscale_freq_shift,
-        )
-        return t_emb
-
-
-class TimestepEmbedding(nn.Module):
-
-    def __init__(self,
-                 in_channels: int,
-                 time_embed_dim: int,
-                 act_fn: str = "silu",
-                 out_dim: int = None,
-                 post_act_fn: Optional[str] = None,
-                 cond_proj_dim=None,
-                 sample_proj_bias=True):
-        super().__init__()
-        linear_cls = nn.Linear
-
-        self.linear_1 = linear_cls(
-            in_channels,
-            time_embed_dim,
-            bias=sample_proj_bias,
-        )
-
-        if cond_proj_dim is not None:
-            self.cond_proj = linear_cls(
-                cond_proj_dim,
-                in_channels,
-                bias=False,
-            )
-        else:
-            self.cond_proj = None
-
-        self.act = get_activation(act_fn)
-
-        if out_dim is not None:
-            time_embed_dim_out = out_dim
-        else:
-            time_embed_dim_out = time_embed_dim
-
-        self.linear_2 = linear_cls(
-            time_embed_dim,
-            time_embed_dim_out,
-            bias=sample_proj_bias,
-        )
-
-        if post_act_fn is None:
-            self.post_act = None
-        else:
-            self.post_act = get_activation(post_act_fn)
-
-    def forward(self, sample, condition=None):
-        if condition is not None:
-            sample = sample + self.cond_proj(condition)
-        sample = self.linear_1(sample)
-
-        if self.act is not None:
-            sample = self.act(sample)
-
-        sample = self.linear_2(sample)
-
-        if self.post_act is not None:
-            sample = self.post_act(sample)
-        return sample
-
-
-class PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module):
-
-    def __init__(self, embedding_dim, size_emb_dim, use_additional_conditions: bool = False):
-        super().__init__()
-
-        self.outdim = size_emb_dim
-        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
-        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
-
-        self.use_additional_conditions = use_additional_conditions
-        if self.use_additional_conditions:
-            self.additional_condition_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
-            self.resolution_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=size_emb_dim)
-            self.nframe_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
-            self.fps_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
-
-    def forward(self, timestep, resolution=None, nframe=None, fps=None):
-        hidden_dtype = next(self.timestep_embedder.parameters()).dtype
-
-        timesteps_proj = self.time_proj(timestep)
-        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype))  # (N, D)
-
-        if self.use_additional_conditions:
-            batch_size = timestep.shape[0]
-            resolution_emb = self.additional_condition_proj(resolution.flatten()).to(hidden_dtype)
-            resolution_emb = self.resolution_embedder(resolution_emb).reshape(batch_size, -1)
-            nframe_emb = self.additional_condition_proj(nframe.flatten()).to(hidden_dtype)
-            nframe_emb = self.nframe_embedder(nframe_emb).reshape(batch_size, -1)
-            conditioning = timesteps_emb + resolution_emb + nframe_emb
-
-            if fps is not None:
-                fps_emb = self.additional_condition_proj(fps.flatten()).to(hidden_dtype)
-                fps_emb = self.fps_embedder(fps_emb).reshape(batch_size, -1)
-                conditioning = conditioning + fps_emb
-        else:
-            conditioning = timesteps_emb
-
-        return conditioning
-
-
-class AdaLayerNormSingle(nn.Module):
-    r"""
-        Norm layer adaptive layer norm single (adaLN-single).
-
-        As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3).
-
-        Parameters:
-            embedding_dim (`int`): The size of each embedding vector.
-            use_additional_conditions (`bool`): To use additional conditions for normalization or not.
-    """
-
-    def __init__(self, embedding_dim: int, use_additional_conditions: bool = False, time_step_rescale=1000):
-        super().__init__()
-
-        self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim,
-                                                             size_emb_dim=embedding_dim // 2,
-                                                             use_additional_conditions=use_additional_conditions)
-
-        self.silu = nn.SiLU()
-        self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True)
-
-        self.time_step_rescale = time_step_rescale  ## timestep usually in [0, 1], we rescale it to [0,1000] for stability
-
-    def forward(
-        self,
-        timestep: torch.Tensor,
-        added_cond_kwargs: Dict[str, torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        embedded_timestep = self.emb(timestep * self.time_step_rescale, **added_cond_kwargs)
-
-        out = self.linear(self.silu(embedded_timestep))
-
-        return out, embedded_timestep
-
-
-class PixArtAlphaTextProjection(nn.Module):
-    """
-    Projects caption embeddings. Also handles dropout for classifier-free guidance.
-
-    Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
-    """
-
-    def __init__(self, in_features, hidden_size):
-        super().__init__()
-        self.linear_1 = nn.Linear(
-            in_features,
-            hidden_size,
-            bias=True,
-        )
-        self.act_1 = nn.GELU(approximate="tanh")
-        self.linear_2 = nn.Linear(
-            hidden_size,
-            hidden_size,
-            bias=True,
-        )
-
-    def forward(self, caption):
-        hidden_states = self.linear_1(caption)
-        hidden_states = self.act_1(hidden_states)
-        hidden_states = self.linear_2(hidden_states)
-        return hidden_states

fastvideo/models/stepvideo/modules/rope.py

@@ -1,90 +0,0 @@
-import torch
-
-from fastvideo.utils.parallel_states import nccl_info
-
-
-class RoPE1D:
-
-    def __init__(self, freq=1e4, F0=1.0, scaling_factor=1.0):
-        self.base = freq
-        self.F0 = F0
-        self.scaling_factor = scaling_factor
-        self.cache = {}
-
-    def get_cos_sin(self, D, seq_len, device, dtype):
-        if (D, seq_len, device, dtype) not in self.cache:
-            inv_freq = 1.0 / (self.base**(torch.arange(0, D, 2).float().to(device) / D))
-            t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
-            freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype)
-            freqs = torch.cat((freqs, freqs), dim=-1)
-            cos = freqs.cos()  # (Seq, Dim)
-            sin = freqs.sin()
-            self.cache[D, seq_len, device, dtype] = (cos, sin)
-        return self.cache[D, seq_len, device, dtype]
-
-    @staticmethod
-    def rotate_half(x):
-        x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
-        return torch.cat((-x2, x1), dim=-1)
-
-    def apply_rope1d(self, tokens, pos1d, cos, sin):
-        assert pos1d.ndim == 2
-        cos = torch.nn.functional.embedding(pos1d, cos)[:, :, None, :]
-        sin = torch.nn.functional.embedding(pos1d, sin)[:, :, None, :]
-        return (tokens * cos) + (self.rotate_half(tokens) * sin)
-
-    def __call__(self, tokens, positions):
-        """
-        input:
-            * tokens: batch_size x ntokens x nheads x dim
-            * positions: batch_size x ntokens (t position of each token)
-        output:
-            * tokens after applying RoPE2D (batch_size x ntokens x nheads x dim)
-        """
-        D = tokens.size(3)
-        assert positions.ndim == 2  # Batch, Seq
-        cos, sin = self.get_cos_sin(D, int(positions.max()) + 1, tokens.device, tokens.dtype)
-        tokens = self.apply_rope1d(tokens, positions, cos, sin)
-        return tokens
-
-
-class RoPE3D(RoPE1D):
-
-    def __init__(self, freq=1e4, F0=1.0, scaling_factor=1.0):
-        super(RoPE3D, self).__init__(freq, F0, scaling_factor)
-        self.position_cache = {}
-
-    def get_mesh_3d(self, rope_positions, bsz):
-        f, h, w = rope_positions
-
-        if f"{f}-{h}-{w}" not in self.position_cache:
-            x = torch.arange(f, device='cpu')
-            y = torch.arange(h, device='cpu')
-            z = torch.arange(w, device='cpu')
-            self.position_cache[f"{f}-{h}-{w}"] = torch.cartesian_prod(x, y, z).view(1, f * h * w, 3).expand(bsz, -1, 3)
-        return self.position_cache[f"{f}-{h}-{w}"]
-
-    def __call__(self, tokens, rope_positions, ch_split, parallel=False):
-        """
-        input:
-            * tokens: batch_size x ntokens x nheads x dim
-            * rope_positions: list of (f, h, w)
-        output:
-            * tokens after applying RoPE2D (batch_size x ntokens x nheads x dim)
-        """
-        assert sum(ch_split) == tokens.size(-1)
-
-        mesh_grid = self.get_mesh_3d(rope_positions, bsz=tokens.shape[0])
-        out = []
-        for i, (D, x) in enumerate(zip(ch_split, torch.split(tokens, ch_split, dim=-1))):
-            cos, sin = self.get_cos_sin(D, int(mesh_grid.max()) + 1, tokens.device, tokens.dtype)
-
-            if parallel:
-                mesh = torch.chunk(mesh_grid[:, :, i], nccl_info.sp_size, dim=1)[nccl_info.rank_within_group].clone()
-            else:
-                mesh = mesh_grid[:, :, i].clone()
-            x = self.apply_rope1d(x, mesh.to(tokens.device), cos, sin)
-            out.append(x)
-
-        tokens = torch.cat(out, dim=-1)
-        return tokens

fastvideo/models/stepvideo/parallel.py

@@ -1,21 +0,0 @@
-import torch
-
-from fastvideo.utils.communications import all_gather
-from fastvideo.utils.parallel_states import nccl_info
-
-
-def parallel_forward(fn_):
-
-    def wrapTheFunction(_, hidden_states, *args, **kwargs):
-        if kwargs['parallel']:
-            hidden_states = torch.chunk(hidden_states, nccl_info.sp_size, dim=-2)[nccl_info.rank_within_group]
-            kwargs['attn_mask'] = torch.chunk(kwargs['attn_mask'], nccl_info.sp_size,
-                                              dim=-2)[nccl_info.rank_within_group]
-        output = fn_(_, hidden_states, *args, **kwargs)
-
-        if kwargs['parallel']:
-            output = all_gather(output.contiguous(), dim=-2)
-
-        return output
-
-    return wrapTheFunction

fastvideo/models/stepvideo/text_encoder/__init__.py

@@ -1,12 +0,0 @@
-import os
-
-import torch
-
-from fastvideo.models.stepvideo.config import parse_args
-
-try:
-    args = parse_args()
-    torch.ops.load_library(
-        os.path.join(args.model_dir, 'lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so'))
-except Exception as err:
-    print(err)

fastvideo/models/stepvideo/text_encoder/clip.py

@@ -1,36 +0,0 @@
-import os
-
-import torch
-import torch.nn as nn
-from transformers import BertModel, BertTokenizer
-
-
-class HunyuanClip(nn.Module):
-    """
-        Hunyuan clip code copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py
-        hunyuan's clip used BertModel and BertTokenizer, so we copy it.
-    """
-
-    def __init__(self, model_dir, max_length=77):
-        super(HunyuanClip, self).__init__()
-
-        self.max_length = max_length
-        self.tokenizer = BertTokenizer.from_pretrained(os.path.join(model_dir, 'tokenizer'))
-        self.text_encoder = BertModel.from_pretrained(os.path.join(model_dir, 'clip_text_encoder'))
-
-    @torch.no_grad
-    def forward(self, prompts, with_mask=True):
-        self.device = next(self.text_encoder.parameters()).device
-        text_inputs = self.tokenizer(
-            prompts,
-            padding="max_length",
-            max_length=self.max_length,
-            truncation=True,
-            return_attention_mask=True,
-            return_tensors="pt",
-        )
-        prompt_embeds = self.text_encoder(
-            text_inputs.input_ids.to(self.device),
-            attention_mask=text_inputs.attention_mask.to(self.device) if with_mask else None,
-        )
-        return prompt_embeds.last_hidden_state, prompt_embeds.pooler_output

fastvideo/models/stepvideo/text_encoder/flashattention.py

@@ -1,45 +0,0 @@
-# Copyright 2025 StepFun Inc. All Rights Reserved.
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-# ==============================================================================
-import torch
-
-
-def flash_attn_func(q,
-                    k,
-                    v,
-                    dropout_p=0.0,
-                    softmax_scale=None,
-                    causal=True,
-                    return_attn_probs=False,
-                    tp_group_rank=0,
-                    tp_group_size=1):
-    softmax_scale = q.size(-1)**(-0.5) if softmax_scale is None else softmax_scale
-    return torch.ops.Optimus.fwd(q, k, v, None, dropout_p, softmax_scale, causal, return_attn_probs, None,
-                                 tp_group_rank, tp_group_size)[0]
-
-
-class FlashSelfAttention(torch.nn.Module):
-
-    def __init__(
-        self,
-        attention_dropout=0.0,
-    ):
-        super().__init__()
-        self.dropout_p = attention_dropout
-
-    def forward(self, q, k, v, cu_seqlens=None, max_seq_len=None):
-        if cu_seqlens is None:
-            output = flash_attn_func(q, k, v, dropout_p=self.dropout_p)
-        else:
-            raise ValueError('cu_seqlens is not supported!')
-
-        return output

fastvideo/models/stepvideo/text_encoder/stepllm.py

@@ -1,291 +0,0 @@
-# Copyright 2025 StepFun Inc. All Rights Reserved.
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-# ==============================================================================
-import os
-from typing import Optional
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-from transformers.modeling_utils import PretrainedConfig, PreTrainedModel
-
-from fastvideo.models.stepvideo.modules.normalization import RMSNorm
-from fastvideo.models.stepvideo.text_encoder.flashattention import FlashSelfAttention
-from fastvideo.models.stepvideo.text_encoder.tokenizer import LLaMaEmbedding, Wrapped_StepChatTokenizer
-from fastvideo.models.stepvideo.utils import with_empty_init
-
-
-def safediv(n, d):
-    q, r = divmod(n, d)
-    assert r == 0
-    return q
-
-
-class MultiQueryAttention(nn.Module):
-
-    def __init__(self, cfg, layer_id=None):
-        super().__init__()
-
-        self.head_dim = cfg.hidden_size // cfg.num_attention_heads
-        self.max_seq_len = cfg.seq_length
-        self.use_flash_attention = cfg.use_flash_attn
-        assert self.use_flash_attention, 'FlashAttention is required!'
-
-        self.n_groups = cfg.num_attention_groups
-        self.tp_size = 1
-        self.n_local_heads = cfg.num_attention_heads
-        self.n_local_groups = self.n_groups
-
-        self.wqkv = nn.Linear(
-            cfg.hidden_size,
-            cfg.hidden_size + self.head_dim * 2 * self.n_groups,
-            bias=False,
-        )
-        self.wo = nn.Linear(
-            cfg.hidden_size,
-            cfg.hidden_size,
-            bias=False,
-        )
-
-        assert self.use_flash_attention, 'non-Flash attention not supported yet.'
-        self.core_attention = FlashSelfAttention(attention_dropout=cfg.attention_dropout)
-
-        self.layer_id = layer_id
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        mask: Optional[torch.Tensor],
-        cu_seqlens: Optional[torch.Tensor],
-        max_seq_len: Optional[torch.Tensor],
-    ):
-        seqlen, bsz, dim = x.shape
-        xqkv = self.wqkv(x)
-
-        xq, xkv = torch.split(
-            xqkv,
-            (dim // self.tp_size, self.head_dim * 2 * self.n_groups // self.tp_size),
-            dim=-1,
-        )
-
-        # gather on 1st dimension
-        xq = xq.view(seqlen, bsz, self.n_local_heads, self.head_dim)
-        xkv = xkv.view(seqlen, bsz, self.n_local_groups, 2 * self.head_dim)
-        xk, xv = xkv.chunk(2, -1)
-
-        # rotary embedding + flash attn
-        xq = rearrange(xq, "s b h d -> b s h d")
-        xk = rearrange(xk, "s b h d -> b s h d")
-        xv = rearrange(xv, "s b h d -> b s h d")
-
-        q_per_kv = self.n_local_heads // self.n_local_groups
-        if q_per_kv > 1:
-            b, s, h, d = xk.size()
-            if h == 1:
-                xk = xk.expand(b, s, q_per_kv, d)
-                xv = xv.expand(b, s, q_per_kv, d)
-            else:
-                ''' To cover the cases where h > 1, we have
-                    the following implementation, which is equivalent to:
-                        xk = xk.repeat_interleave(q_per_kv, dim=-2)
-                        xv = xv.repeat_interleave(q_per_kv, dim=-2)
-                    but can avoid calling aten::item() that involves cpu.
-                '''
-                idx = torch.arange(q_per_kv * h, device=xk.device).reshape(q_per_kv, -1).permute(1, 0).flatten()
-                xk = torch.index_select(xk.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous()
-                xv = torch.index_select(xv.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous()
-
-        if self.use_flash_attention:
-            output = self.core_attention(xq, xk, xv, cu_seqlens=cu_seqlens, max_seq_len=max_seq_len)
-            # reduce-scatter only support first dimension now
-            output = rearrange(output, "b s h d -> s b (h d)").contiguous()
-        else:
-            xq, xk, xv = [rearrange(x, "b s ... -> s b ...").contiguous() for x in (xq, xk, xv)]
-            output = self.core_attention(xq, xk, xv, mask)
-        output = self.wo(output)
-        return output
-
-
-class FeedForward(nn.Module):
-
-    def __init__(
-        self,
-        cfg,
-        dim: int,
-        hidden_dim: int,
-        layer_id: int,
-        multiple_of: int = 256,
-    ):
-        super().__init__()
-
-        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
-
-        def swiglu(x):
-            x = torch.chunk(x, 2, dim=-1)
-            return F.silu(x[0]) * x[1]
-
-        self.swiglu = swiglu
-
-        self.w1 = nn.Linear(
-            dim,
-            2 * hidden_dim,
-            bias=False,
-        )
-        self.w2 = nn.Linear(
-            hidden_dim,
-            dim,
-            bias=False,
-        )
-
-    def forward(self, x):
-        x = self.swiglu(self.w1(x))
-        output = self.w2(x)
-        return output
-
-
-class TransformerBlock(nn.Module):
-
-    def __init__(self, cfg, layer_id: int):
-        super().__init__()
-
-        self.n_heads = cfg.num_attention_heads
-        self.dim = cfg.hidden_size
-        self.head_dim = cfg.hidden_size // cfg.num_attention_heads
-        self.attention = MultiQueryAttention(
-            cfg,
-            layer_id=layer_id,
-        )
-
-        self.feed_forward = FeedForward(
-            cfg,
-            dim=cfg.hidden_size,
-            hidden_dim=cfg.ffn_hidden_size,
-            layer_id=layer_id,
-        )
-        self.layer_id = layer_id
-        self.attention_norm = RMSNorm(
-            cfg.hidden_size,
-            eps=cfg.layernorm_epsilon,
-        )
-        self.ffn_norm = RMSNorm(
-            cfg.hidden_size,
-            eps=cfg.layernorm_epsilon,
-        )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        mask: Optional[torch.Tensor],
-        cu_seqlens: Optional[torch.Tensor],
-        max_seq_len: Optional[torch.Tensor],
-    ):
-        residual = self.attention.forward(self.attention_norm(x), mask, cu_seqlens, max_seq_len)
-        h = x + residual
-        ffn_res = self.feed_forward.forward(self.ffn_norm(h))
-        out = h + ffn_res
-        return out
-
-
-class Transformer(nn.Module):
-
-    def __init__(
-        self,
-        config,
-        max_seq_size=8192,
-    ):
-        super().__init__()
-        self.num_layers = config.num_layers
-        self.layers = self._build_layers(config)
-
-    def _build_layers(self, config):
-        layers = torch.nn.ModuleList()
-        for layer_id in range(self.num_layers):
-            layers.append(TransformerBlock(
-                config,
-                layer_id=layer_id + 1,
-            ))
-        return layers
-
-    def forward(
-        self,
-        hidden_states,
-        attention_mask,
-        cu_seqlens=None,
-        max_seq_len=None,
-    ):
-
-        if max_seq_len is not None and not isinstance(max_seq_len, torch.Tensor):
-            max_seq_len = torch.tensor(max_seq_len, dtype=torch.int32, device="cpu")
-
-        for lid, layer in enumerate(self.layers):
-            hidden_states = layer(
-                hidden_states,
-                attention_mask,
-                cu_seqlens,
-                max_seq_len,
-            )
-        return hidden_states
-
-
-class Step1Model(PreTrainedModel):
-    config_class = PretrainedConfig
-
-    @with_empty_init
-    def __init__(
-        self,
-        config,
-    ):
-        super().__init__(config)
-        self.tok_embeddings = LLaMaEmbedding(config)
-        self.transformer = Transformer(config)
-
-    def forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-    ):
-
-        hidden_states = self.tok_embeddings(input_ids)
-
-        hidden_states = self.transformer(
-            hidden_states,
-            attention_mask,
-        )
-        return hidden_states
-
-
-class STEP1TextEncoder(torch.nn.Module):
-
-    def __init__(self, model_dir, max_length=320):
-        super(STEP1TextEncoder, self).__init__()
-        self.max_length = max_length
-        self.text_tokenizer = Wrapped_StepChatTokenizer(os.path.join(model_dir, 'step1_chat_tokenizer.model'))
-        text_encoder = Step1Model.from_pretrained(model_dir)
-        self.text_encoder = text_encoder.eval().to(torch.bfloat16)
-
-    @torch.no_grad
-    def forward(self, prompts, with_mask=True, max_length=None):
-        self.device = next(self.text_encoder.parameters()).device
-        with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
-            if type(prompts) is str:
-                prompts = [prompts]
-
-            txt_tokens = self.text_tokenizer(prompts,
-                                             max_length=max_length or self.max_length,
-                                             padding="max_length",
-                                             truncation=True,
-                                             return_tensors="pt")
-            y = self.text_encoder(txt_tokens.input_ids.to(self.device),
-                                  attention_mask=txt_tokens.attention_mask.to(self.device) if with_mask else None)
-            y_mask = txt_tokens.attention_mask
-        return y.transpose(0, 1), y_mask

fastvideo/models/stepvideo/text_encoder/tokenizer.py

@@ -1,209 +0,0 @@
-# Copyright 2025 StepFun Inc. All Rights Reserved.
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-# ==============================================================================
-from typing import List
-
-import torch
-import torch.nn as nn
-
-
-class LLaMaEmbedding(nn.Module):
-    """Language model embeddings.
-
-    Arguments:
-        hidden_size: hidden size
-        vocab_size: vocabulary size
-        max_sequence_length: maximum size of sequence. This
-                             is used for positional embedding
-        embedding_dropout_prob: dropout probability for embeddings
-        init_method: weight initialization method
-        num_tokentypes: size of the token-type embeddings. 0 value
-                        will ignore this embedding
-    """
-
-    def __init__(
-        self,
-        cfg,
-    ):
-        super().__init__()
-        self.hidden_size = cfg.hidden_size
-        self.params_dtype = cfg.params_dtype
-        self.fp32_residual_connection = cfg.fp32_residual_connection
-        self.embedding_weights_in_fp32 = cfg.embedding_weights_in_fp32
-        self.word_embeddings = torch.nn.Embedding(
-            cfg.padded_vocab_size,
-            self.hidden_size,
-        )
-        self.embedding_dropout = torch.nn.Dropout(cfg.hidden_dropout)
-
-    def forward(self, input_ids):
-        # Embeddings.
-        if self.embedding_weights_in_fp32:
-            self.word_embeddings = self.word_embeddings.to(torch.float32)
-        embeddings = self.word_embeddings(input_ids)
-        if self.embedding_weights_in_fp32:
-            embeddings = embeddings.to(self.params_dtype)
-            self.word_embeddings = self.word_embeddings.to(self.params_dtype)
-
-        # Data format change to avoid explicit transposes : [b s h] --> [s b h].
-        embeddings = embeddings.transpose(0, 1).contiguous()
-
-        # If the input flag for fp32 residual connection is set, convert for float.
-        if self.fp32_residual_connection:
-            embeddings = embeddings.float()
-
-        # Dropout.
-        embeddings = self.embedding_dropout(embeddings)
-
-        return embeddings
-
-
-class StepChatTokenizer:
-    """Step Chat Tokenizer"""
-
-    def __init__(
-            self,
-            model_file,
-            name="StepChatTokenizer",
-            bot_token="<|BOT|>",  # Begin of Turn
-            eot_token="<|EOT|>",  # End of Turn
-            call_start_token="<|CALL_START|>",  # Call Start
-            call_end_token="<|CALL_END|>",  # Call End
-            think_start_token="<|THINK_START|>",  # Think Start
-            think_end_token="<|THINK_END|>",  # Think End
-            mask_start_token="<|MASK_1e69f|>",  # Mask start
-            mask_end_token="<|UNMASK_1e69f|>",  # Mask end
-    ):
-        import sentencepiece
-
-        self._tokenizer = sentencepiece.SentencePieceProcessor(model_file=model_file)
-
-        self._vocab = {}
-        self._inv_vocab = {}
-
-        self._special_tokens = {}
-        self._inv_special_tokens = {}
-
-        self._t5_tokens = []
-
-        for idx in range(self._tokenizer.get_piece_size()):
-            text = self._tokenizer.id_to_piece(idx)
-            self._inv_vocab[idx] = text
-            self._vocab[text] = idx
-
-            if self._tokenizer.is_control(idx) or self._tokenizer.is_unknown(idx):
-                self._special_tokens[text] = idx
-                self._inv_special_tokens[idx] = text
-
-        self._unk_id = self._tokenizer.unk_id()
-        self._bos_id = self._tokenizer.bos_id()
-        self._eos_id = self._tokenizer.eos_id()
-
-        for token in [bot_token, eot_token, call_start_token, call_end_token, think_start_token, think_end_token]:
-            assert token in self._vocab, f"Token '{token}' not found in tokenizer"
-            assert token in self._special_tokens, f"Token '{token}' is not a special token"
-
-        for token in [mask_start_token, mask_end_token]:
-            assert token in self._vocab, f"Token '{token}' not found in tokenizer"
-
-        self._bot_id = self._tokenizer.piece_to_id(bot_token)
-        self._eot_id = self._tokenizer.piece_to_id(eot_token)
-        self._call_start_id = self._tokenizer.piece_to_id(call_start_token)
-        self._call_end_id = self._tokenizer.piece_to_id(call_end_token)
-        self._think_start_id = self._tokenizer.piece_to_id(think_start_token)
-        self._think_end_id = self._tokenizer.piece_to_id(think_end_token)
-        self._mask_start_id = self._tokenizer.piece_to_id(mask_start_token)
-        self._mask_end_id = self._tokenizer.piece_to_id(mask_end_token)
-
-        self._underline_id = self._tokenizer.piece_to_id("\u2581")
-
-    @property
-    def vocab(self):
-        return self._vocab
-
-    @property
-    def inv_vocab(self):
-        return self._inv_vocab
-
-    @property
-    def vocab_size(self):
-        return self._tokenizer.vocab_size()
-
-    def tokenize(self, text: str) -> List[int]:
-        return self._tokenizer.encode_as_ids(text)
-
-    def detokenize(self, token_ids: List[int]) -> str:
-        return self._tokenizer.decode_ids(token_ids)
-
-
-class Tokens:
-
-    def __init__(self, input_ids, cu_input_ids, attention_mask, cu_seqlens, max_seq_len) -> None:
-        self.input_ids = input_ids
-        self.attention_mask = attention_mask
-        self.cu_input_ids = cu_input_ids
-        self.cu_seqlens = cu_seqlens
-        self.max_seq_len = max_seq_len
-
-    def to(self, device):
-        self.input_ids = self.input_ids.to(device)
-        self.attention_mask = self.attention_mask.to(device)
-        self.cu_input_ids = self.cu_input_ids.to(device)
-        self.cu_seqlens = self.cu_seqlens.to(device)
-        return self
-
-
-class Wrapped_StepChatTokenizer(StepChatTokenizer):
-
-    def __call__(self, text, max_length=320, padding="max_length", truncation=True, return_tensors="pt"):
-        # [bos, ..., eos, pad, pad, ..., pad]
-        self.BOS = 1
-        self.EOS = 2
-        self.PAD = 2
-        out_tokens = []
-        attn_mask = []
-        if len(text) == 0:
-            part_tokens = [self.BOS] + [self.EOS]
-            valid_size = len(part_tokens)
-            if len(part_tokens) < max_length:
-                part_tokens += [self.PAD] * (max_length - valid_size)
-            out_tokens.append(part_tokens)
-            attn_mask.append([1] * valid_size + [0] * (max_length - valid_size))
-        else:
-            for part in text:
-                part_tokens = self.tokenize(part)
-                part_tokens = part_tokens[:(max_length - 2)]  # leave 2 space for bos and eos
-                part_tokens = [self.BOS] + part_tokens + [self.EOS]
-                valid_size = len(part_tokens)
-                if len(part_tokens) < max_length:
-                    part_tokens += [self.PAD] * (max_length - valid_size)
-                out_tokens.append(part_tokens)
-                attn_mask.append([1] * valid_size + [0] * (max_length - valid_size))
-
-        out_tokens = torch.tensor(out_tokens, dtype=torch.long)
-        attn_mask = torch.tensor(attn_mask, dtype=torch.long)
-
-        # padding y based on tp size
-        padded_len = 0
-        padded_flag = True if padded_len > 0 else False
-        if padded_flag:
-            pad_tokens = torch.tensor([[self.PAD] * max_length], device=out_tokens.device)
-            pad_attn_mask = torch.tensor([[1] * padded_len + [0] * (max_length - padded_len)], device=attn_mask.device)
-            out_tokens = torch.cat([out_tokens, pad_tokens], dim=0)
-            attn_mask = torch.cat([attn_mask, pad_attn_mask], dim=0)
-
-        # cu_seqlens
-        cu_out_tokens = out_tokens.masked_select(attn_mask != 0).unsqueeze(0)
-        seqlen = attn_mask.sum(dim=1).tolist()
-        cu_seqlens = torch.cumsum(torch.tensor([0] + seqlen), 0).to(device=out_tokens.device, dtype=torch.int32)
-        max_seq_len = max(seqlen)
-        return Tokens(out_tokens, cu_out_tokens, attn_mask, cu_seqlens, max_seq_len)

fastvideo/models/stepvideo/utils/__init__.py

@@ -1,2 +0,0 @@
-from .utils import *
-from .video_process import *

fastvideo/models/stepvideo/utils/quantization.py

@@ -1,117 +0,0 @@
-# from stepvideo.diffusion.video_pipeline import StepVideoPipeline
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-
-
-def get_fp_maxval(bits=8, mantissa_bit=3, sign_bits=1):
-    _bits = torch.tensor(bits)
-    _mantissa_bit = torch.tensor(mantissa_bit)
-    _sign_bits = torch.tensor(sign_bits)
-    M = torch.clamp(torch.round(_mantissa_bit), 1, _bits - _sign_bits)
-    E = _bits - _sign_bits - M
-    bias = 2**(E - 1) - 1
-    mantissa = 1
-    for i in range(mantissa_bit - 1):
-        mantissa += 1 / (2**(i + 1))
-    maxval = mantissa * 2**(2**E - 1 - bias)
-    return maxval
-
-
-def quantize_to_fp8(x, bits=8, mantissa_bit=3, sign_bits=1):
-    """
-    Default is E4M3.
-    """
-    bits = torch.tensor(bits)
-    mantissa_bit = torch.tensor(mantissa_bit)
-    sign_bits = torch.tensor(sign_bits)
-    M = torch.clamp(torch.round(mantissa_bit), 1, bits - sign_bits)
-    E = bits - sign_bits - M
-    bias = 2**(E - 1) - 1
-    mantissa = 1
-    for i in range(mantissa_bit - 1):
-        mantissa += 1 / (2**(i + 1))
-    maxval = mantissa * 2**(2**E - 1 - bias)
-    minval = -maxval
-    minval = -maxval if sign_bits == 1 else torch.zeros_like(maxval)
-    input_clamp = torch.min(torch.max(x, minval), maxval)
-    log_scales = torch.clamp((torch.floor(torch.log2(torch.abs(input_clamp)) + bias)).detach(), 1.0)
-    log_scales = 2.0**(log_scales - M - bias.type(x.dtype))
-    # dequant
-    qdq_out = torch.round(input_clamp / log_scales) * log_scales
-    return qdq_out, log_scales
-
-
-def fp8_tensor_quant(x, scale, bits=8, mantissa_bit=3, sign_bits=1):
-    for i in range(len(x.shape) - 1):
-        scale = scale.unsqueeze(-1)
-    new_x = x / scale
-    quant_dequant_x, log_scales = quantize_to_fp8(new_x, bits=bits, mantissa_bit=mantissa_bit, sign_bits=sign_bits)
-    return quant_dequant_x, scale, log_scales
-
-
-def fp8_activation_dequant(qdq_out, scale, dtype):
-    qdq_out = qdq_out.type(dtype)
-    quant_dequant_x = qdq_out * scale.to(dtype)
-    return quant_dequant_x
-
-
-def fp8_linear_forward(cls, original_dtype, input):
-    weight_dtype = cls.weight.dtype
-    #####
-    if cls.weight.dtype != torch.float8_e4m3fn:
-        assert False
-        maxval = get_fp_maxval()
-        scale = torch.max(torch.abs(cls.weight.flatten())) / maxval
-        linear_weight, scale, log_scales = fp8_tensor_quant(cls.weight, scale)
-        linear_weight = linear_weight.to(torch.float8_e4m3fn)
-        weight_dtype = linear_weight.dtype
-    else:
-        scale = cls.fp8_scale.to(cls.weight.device)
-        linear_weight = cls.weight
-    #####
-
-    if weight_dtype == torch.float8_e4m3fn:
-        if True or len(input.shape) == 3:
-            cls_dequant = fp8_activation_dequant(linear_weight, scale, original_dtype)
-            if cls.bias is not None:
-                print(f"input dtype: {input.dtype}")
-                print(f"cls_dequant dtype: {cls_dequant.dtype}")
-                print(f"cls.bias dtype: {cls.bias.dtype}")
-
-                output = F.linear(input, cls_dequant, cls.bias)
-            else:
-                output = F.linear(input, cls_dequant)
-            return output
-        else:
-            return cls.original_forward(input.to(original_dtype))
-    else:
-        return cls.original_forward(input)
-
-
-def convert_fp8_linear(module, original_dtype, params_to_keep={}):
-    setattr(module, "fp8_matmul_enabled", True)
-    fp8_layers = []
-    scale_dict = {}
-    counter = 0
-    for key, layer in module.named_modules():
-        if isinstance(layer, nn.Linear) and 'transformer_blocks' in key:
-            print(f"Converting {key} to FP8")
-            fp8_layers.append(key)
-            original_forward = layer.forward
-            maxval = get_fp_maxval()
-            scale = torch.max(torch.abs(layer.weight.flatten())) / maxval
-
-            original_weight = layer.weight.data  # Store a reference to the original weights
-            quantized_weight, scale, _ = fp8_tensor_quant(original_weight, scale)
-            scale_dict[key] = scale
-            layer.weight = torch.nn.Parameter(quantized_weight.to(torch.float8_e4m3fn))
-            del original_weight  # Delete the reference to the original weights
-            torch.cuda.empty_cache()
-
-            # print(f"layer weight dtype: {layer.weight.dtype} for layer {key}")
-            setattr(layer, "fp8_scale", scale.to(dtype=original_dtype))
-            setattr(layer, "original_forward", original_forward)
-            setattr(layer, "forward", lambda input, m=layer: fp8_linear_forward(m, original_dtype, input))
-            counter += 1
-    return scale_dict

fastvideo/models/stepvideo/utils/utils.py

@@ -1,60 +0,0 @@
-import random
-from functools import wraps
-
-import numpy as np
-import torch
-import torch.utils._device
-
-
-def setup_seed(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-class EmptyInitOnDevice(torch.overrides.TorchFunctionMode):
-
-    def __init__(self, device=None):
-        self.device = device
-
-    def __torch_function__(self, func, types, args=(), kwargs=None):
-        kwargs = kwargs or {}
-        if getattr(func, '__module__', None) == 'torch.nn.init':
-            if 'tensor' in kwargs:
-                return kwargs['tensor']
-            else:
-                return args[0]
-        if self.device is not None and func in torch.utils._device._device_constructors(
-        ) and kwargs.get('device') is None:
-            kwargs['device'] = self.device
-        return func(*args, **kwargs)
-
-
-def with_empty_init(func):
-
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        with EmptyInitOnDevice('cpu'):
-            return func(*args, **kwargs)
-
-    return wrapper
-
-
-def culens2mask(cu_seqlens=None, cu_seqlens_kv=None, max_seqlen=None, max_seqlen_kv=None, is_causal=False):
-    assert len(cu_seqlens) == len(cu_seqlens_kv)
-    "q k v should have same bsz..."
-    bsz = len(cu_seqlens) - 1
-    seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
-    seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]
-
-    attn_mask = torch.zeros(bsz, max_seqlen, max_seqlen_kv, dtype=torch.bool)
-    for i, (seq_len, seq_len_kv) in enumerate(zip(seqlens, seqlens_kv)):
-        if is_causal:
-            attn_mask[i, :seq_len, :seq_len_kv] = torch.triu(torch.ones(seq_len, seq_len_kv), diagonal=1).bool()
-        else:
-            attn_mask[i, :seq_len, :seq_len_kv] = torch.ones([seq_len, seq_len_kv], dtype=torch.bool)
-
-    return attn_mask

fastvideo/models/stepvideo/utils/video_process.py

@@ -1,51 +0,0 @@
-import os
-
-import imageio
-import numpy as np
-import torch
-
-
-class VideoProcessor:
-
-    def __init__(self, save_path: str = './results', name_suffix: str = ''):
-        self.save_path = save_path
-        os.makedirs(self.save_path, exist_ok=True)
-        self.name_suffix = name_suffix
-
-    def crop2standard540p(self, vid_array):
-        _, height, width, _ = vid_array.shape
-        height_center = height // 2
-        width_center = width // 2
-        if width_center > height_center:  ## horizon mode
-            return vid_array[:, height_center - 270:height_center + 270, width_center - 480:width_center + 480]
-        elif width_center < height_center:  ## portrait mode
-            return vid_array[:, height_center - 480:height_center + 480, width_center - 270:width_center + 270]
-        else:
-            return vid_array
-
-    def save_imageio_video(self, video_array: np.array, output_filename: str, fps=25, codec='libx264'):
-
-        ffmpeg_params = [
-            "-vf",
-            "atadenoise=0a=0.1:0b=0.1:1a=0.1:1b=0.1",  # denoise
-        ]
-
-        with imageio.get_writer(output_filename, fps=fps, codec=codec, ffmpeg_params=ffmpeg_params) as vid_writer:
-            for img_array in video_array:
-                vid_writer.append_data(img_array)
-
-    def postprocess_video(self, video_tensor, output_file_name='', output_type="mp4", crop2standard540p=True):
-        if len(self.name_suffix) == 0:
-            video_path = os.path.join(self.save_path, f"{output_file_name}.{output_type}")
-        else:
-            video_path = os.path.join(self.save_path, f"{output_file_name}-{self.name_suffix}.{output_type}")
-
-        video_tensor = torch.cat([t for t in video_tensor], dim=-2)
-        video_tensor = (video_tensor.cpu().clamp(-1, 1) + 1) * 127.5
-        video_array = video_tensor.clamp(0, 255).to(torch.uint8).numpy().transpose(0, 2, 3, 1)
-
-        if crop2standard540p:
-            video_array = self.crop2standard540p(video_array)
-
-        self.save_imageio_video(video_array, video_path)
-        print(f"Saved the generated video in {video_path}")

fastvideo/models/stepvideo/vae/vae.py

@@ -1,975 +0,0 @@
-# Copyright 2025 StepFun Inc. All Rights Reserved.
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-# ==============================================================================
-import torch
-from einops import rearrange
-from torch import nn
-from torch.nn import functional as F
-
-from fastvideo.models.stepvideo.utils import with_empty_init
-
-
-def base_group_norm(x, norm_layer, act_silu=False, channel_last=False):
-    if hasattr(base_group_norm, 'spatial') and base_group_norm.spatial:
-        assert channel_last
-        x_shape = x.shape
-        x = x.flatten(0, 1)
-        if channel_last:
-            # Permute to NCHW format
-            x = x.permute(0, 3, 1, 2)
-
-        out = F.group_norm(x.contiguous(), norm_layer.num_groups, norm_layer.weight, norm_layer.bias, norm_layer.eps)
-        if act_silu:
-            out = F.silu(out)
-
-        if channel_last:
-            # Permute back to NHWC format
-            out = out.permute(0, 2, 3, 1)
-
-        out = out.view(x_shape)
-    else:
-        if channel_last:
-            # Permute to NCHW format
-            x = x.permute(0, 3, 1, 2)
-        out = F.group_norm(x.contiguous(), norm_layer.num_groups, norm_layer.weight, norm_layer.bias, norm_layer.eps)
-        if act_silu:
-            out = F.silu(out)
-        if channel_last:
-            # Permute back to NHWC format
-            out = out.permute(0, 2, 3, 1)
-    return out
-
-
-def base_conv2d(x, conv_layer, channel_last=False, residual=None):
-    if channel_last:
-        x = x.permute(0, 3, 1, 2)  # NHWC to NCHW
-    out = F.conv2d(x, conv_layer.weight, conv_layer.bias, stride=conv_layer.stride, padding=conv_layer.padding)
-    if residual is not None:
-        if channel_last:
-            residual = residual.permute(0, 3, 1, 2)  # NHWC to NCHW
-        out += residual
-    if channel_last:
-        out = out.permute(0, 2, 3, 1)  # NCHW to NHWC
-    return out
-
-
-def base_conv3d(x, conv_layer, channel_last=False, residual=None, only_return_output=False):
-    if only_return_output:
-        size = cal_outsize(x.shape, conv_layer.weight.shape, conv_layer.stride, conv_layer.padding)
-        return torch.empty(size, device=x.device, dtype=x.dtype)
-    if channel_last:
-        x = x.permute(0, 4, 1, 2, 3)  # NDHWC to NCDHW
-    out = F.conv3d(x, conv_layer.weight, conv_layer.bias, stride=conv_layer.stride, padding=conv_layer.padding)
-    if residual is not None:
-        if channel_last:
-            residual = residual.permute(0, 4, 1, 2, 3)  # NDHWC to NCDHW
-        out += residual
-    if channel_last:
-        out = out.permute(0, 2, 3, 4, 1)  # NCDHW to NDHWC
-    return out
-
-
-def cal_outsize(input_sizes, kernel_sizes, stride, padding):
-    stride_d, stride_h, stride_w = stride
-    padding_d, padding_h, padding_w = padding
-    dilation_d, dilation_h, dilation_w = 1, 1, 1
-
-    in_d = input_sizes[1]
-    in_h = input_sizes[2]
-    in_w = input_sizes[3]
-
-    kernel_d = kernel_sizes[2]
-    kernel_h = kernel_sizes[3]
-    kernel_w = kernel_sizes[4]
-    out_channels = kernel_sizes[0]
-
-    out_d = calc_out_(in_d, padding_d, dilation_d, kernel_d, stride_d)
-    out_h = calc_out_(in_h, padding_h, dilation_h, kernel_h, stride_h)
-    out_w = calc_out_(in_w, padding_w, dilation_w, kernel_w, stride_w)
-    size = [input_sizes[0], out_d, out_h, out_w, out_channels]
-    return size
-
-
-def calc_out_(in_size, padding, dilation, kernel, stride):
-    return (in_size + 2 * padding - dilation * (kernel - 1) - 1) // stride + 1
-
-
-def base_conv3d_channel_last(x, conv_layer, residual=None):
-    in_numel = x.numel()
-    out_numel = int(x.numel() * conv_layer.out_channels / conv_layer.in_channels)
-    if (in_numel >= 2**30) or (out_numel >= 2**30):
-        assert conv_layer.stride[0] == 1, "time split asks time stride = 1"
-
-        B, T, H, W, C = x.shape
-        K = conv_layer.kernel_size[0]
-
-        chunks = 4
-        chunk_size = T // chunks
-
-        if residual is None:
-            out_nhwc = base_conv3d(x, conv_layer, channel_last=True, residual=residual, only_return_output=True)
-        else:
-            out_nhwc = residual
-
-        assert B == 1
-        for i in range(chunks):
-            if i == chunks - 1:
-                xi = x[:1, chunk_size * i:]
-                out_nhwci = out_nhwc[:1, chunk_size * i:]
-            else:
-                xi = x[:1, chunk_size * i:chunk_size * (i + 1) + K - 1]
-                out_nhwci = out_nhwc[:1, chunk_size * i:chunk_size * (i + 1)]
-            if residual is not None:
-                if i == chunks - 1:
-                    ri = residual[:1, chunk_size * i:]
-                else:
-                    ri = residual[:1, chunk_size * i:chunk_size * (i + 1)]
-            else:
-                ri = None
-            out_nhwci.copy_(base_conv3d(xi, conv_layer, channel_last=True, residual=ri))
-    else:
-        out_nhwc = base_conv3d(x, conv_layer, channel_last=True, residual=residual)
-    return out_nhwc
-
-
-class Upsample2D(nn.Module):
-
-    def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_channels=None):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.use_conv_transpose = use_conv_transpose
-
-        if use_conv:
-            self.conv = nn.Conv2d(self.channels, self.out_channels, 3, padding=1)
-        else:
-            assert "Not Supported"
-            self.conv = nn.ConvTranspose2d(channels, self.out_channels, 4, 2, 1)
-
-    def forward(self, x, output_size=None):
-        assert x.shape[-1] == self.channels
-
-        if self.use_conv_transpose:
-            return self.conv(x)
-
-        if output_size is None:
-            x = F.interpolate(x.permute(0, 3, 1, 2).to(memory_format=torch.channels_last),
-                              scale_factor=2.0,
-                              mode='nearest').permute(0, 2, 3, 1).contiguous()
-        else:
-            x = F.interpolate(x.permute(0, 3, 1, 2).to(memory_format=torch.channels_last),
-                              size=output_size,
-                              mode='nearest').permute(0, 2, 3, 1).contiguous()
-
-        # x = self.conv(x)
-        x = base_conv2d(x, self.conv, channel_last=True)
-        return x
-
-
-class Downsample2D(nn.Module):
-
-    def __init__(self, channels, use_conv=False, out_channels=None, padding=1):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.padding = padding
-        stride = 2
-
-        if use_conv:
-            self.conv = nn.Conv2d(self.channels, self.out_channels, 3, stride=stride, padding=padding)
-        else:
-            assert self.channels == self.out_channels
-            self.conv = nn.AvgPool2d(kernel_size=stride, stride=stride)
-
-    def forward(self, x):
-        assert x.shape[-1] == self.channels
-        if self.use_conv and self.padding == 0:
-            pad = (0, 0, 0, 1, 0, 1)
-            x = F.pad(x, pad, mode="constant", value=0)
-
-        assert x.shape[-1] == self.channels
-        # x = self.conv(x)
-        x = base_conv2d(x, self.conv, channel_last=True)
-        return x
-
-
-class CausalConv(nn.Module):
-
-    def __init__(self, chan_in, chan_out, kernel_size, **kwargs):
-        super().__init__()
-
-        if isinstance(kernel_size, int):
-            kernel_size = kernel_size if isinstance(kernel_size, tuple) else ((kernel_size, ) * 3)
-        time_kernel_size, height_kernel_size, width_kernel_size = kernel_size
-
-        self.dilation = kwargs.pop('dilation', 1)
-        self.stride = kwargs.pop('stride', 1)
-        if isinstance(self.stride, int):
-            self.stride = (self.stride, 1, 1)
-        time_pad = self.dilation * (time_kernel_size - 1) + max((1 - self.stride[0]), 0)
-        height_pad = height_kernel_size // 2
-        width_pad = width_kernel_size // 2
-        self.time_causal_padding = (width_pad, width_pad, height_pad, height_pad, time_pad, 0)
-        self.time_uncausal_padding = (width_pad, width_pad, height_pad, height_pad, 0, 0)
-
-        self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=self.stride, dilation=self.dilation, **kwargs)
-        self.is_first_run = True
-
-    def forward(self, x, is_init=True, residual=None):
-        x = nn.functional.pad(x, self.time_causal_padding if is_init else self.time_uncausal_padding)
-
-        x = self.conv(x)
-        if residual is not None:
-            x.add_(residual)
-        return x
-
-
-class ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(nn.Module):
-
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        factor: int,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.factor = factor
-        assert out_channels * factor**3 % in_channels == 0
-        self.repeats = out_channels * factor**3 // in_channels
-
-    def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor:
-        x = x.repeat_interleave(self.repeats, dim=1)
-        x = x.view(x.size(0), self.out_channels, self.factor, self.factor, self.factor, x.size(2), x.size(3), x.size(4))
-        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
-        x = x.view(x.size(0), self.out_channels,
-                   x.size(2) * self.factor,
-                   x.size(4) * self.factor,
-                   x.size(6) * self.factor)
-        x = x[:, :, self.factor - 1:, :, :]
-        return x
-
-
-class ConvPixelShuffleUpSampleLayer3D(nn.Module):
-
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int,
-        factor: int,
-    ):
-        super().__init__()
-        self.factor = factor
-        out_ratio = factor**3
-        self.conv = CausalConv(in_channels, out_channels * out_ratio, kernel_size=kernel_size)
-
-    def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor:
-        x = self.conv(x, is_init)
-        x = self.pixel_shuffle_3d(x, self.factor)
-        return x
-
-    @staticmethod
-    def pixel_shuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor:
-        batch_size, channels, depth, height, width = x.size()
-        new_channels = channels // (factor**3)
-        new_depth = depth * factor
-        new_height = height * factor
-        new_width = width * factor
-
-        x = x.view(batch_size, new_channels, factor, factor, factor, depth, height, width)
-        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
-        x = x.view(batch_size, new_channels, new_depth, new_height, new_width)
-        x = x[:, :, factor - 1:, :, :]
-        return x
-
-
-class ConvPixelUnshuffleDownSampleLayer3D(nn.Module):
-
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int,
-        factor: int,
-    ):
-        super().__init__()
-        self.factor = factor
-        out_ratio = factor**3
-        assert out_channels % out_ratio == 0
-        self.conv = CausalConv(in_channels, out_channels // out_ratio, kernel_size=kernel_size)
-
-    def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor:
-        x = self.conv(x, is_init)
-        x = self.pixel_unshuffle_3d(x, self.factor)
-        return x
-
-    @staticmethod
-    def pixel_unshuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor:
-        pad = (0, 0, 0, 0, factor - 1, 0)  # (left, right, top, bottom, front, back)
-        x = F.pad(x, pad)
-        B, C, D, H, W = x.shape
-        x = x.view(B, C, D // factor, factor, H // factor, factor, W // factor, factor)
-        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
-        x = x.view(B, C * factor**3, D // factor, H // factor, W // factor)
-        return x
-
-
-class PixelUnshuffleChannelAveragingDownSampleLayer3D(nn.Module):
-
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        factor: int,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.factor = factor
-        assert in_channels * factor**3 % out_channels == 0
-        self.group_size = in_channels * factor**3 // out_channels
-
-    def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor:
-        pad = (0, 0, 0, 0, self.factor - 1, 0)  # (left, right, top, bottom, front, back)
-        x = F.pad(x, pad)
-        B, C, D, H, W = x.shape
-        x = x.view(B, C, D // self.factor, self.factor, H // self.factor, self.factor, W // self.factor, self.factor)
-        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
-        x = x.view(B, C * self.factor**3, D // self.factor, H // self.factor, W // self.factor)
-        x = x.view(B, self.out_channels, self.group_size, D // self.factor, H // self.factor, W // self.factor)
-        x = x.mean(dim=2)
-        return x
-
-
-def base_group_norm_with_zero_pad(x, norm_layer, act_silu=True, pad_size=2):
-    out_shape = list(x.shape)
-    out_shape[1] += pad_size
-    out = torch.empty(out_shape, dtype=x.dtype, device=x.device)
-    out[:, pad_size:] = base_group_norm(x, norm_layer, act_silu=act_silu, channel_last=True)
-    out[:, :pad_size] = 0
-    return out
-
-
-class CausalConvChannelLast(CausalConv):
-
-    def __init__(self, chan_in, chan_out, kernel_size, **kwargs):
-        super().__init__(chan_in, chan_out, kernel_size, **kwargs)
-
-        self.time_causal_padding = (0, 0) + self.time_causal_padding
-        self.time_uncausal_padding = (0, 0) + self.time_uncausal_padding
-
-    def forward(self, x, is_init=True, residual=None):
-        if self.is_first_run:
-            self.is_first_run = False
-            # self.conv.weight = nn.Parameter(self.conv.weight.permute(0,2,3,4,1).contiguous())
-
-        x = nn.functional.pad(x, self.time_causal_padding if is_init else self.time_uncausal_padding)
-
-        x = base_conv3d_channel_last(x, self.conv, residual=residual)
-        return x
-
-
-class CausalConvAfterNorm(CausalConv):
-
-    def __init__(self, chan_in, chan_out, kernel_size, **kwargs):
-        super().__init__(chan_in, chan_out, kernel_size, **kwargs)
-
-        if self.time_causal_padding == (1, 1, 1, 1, 2, 0):
-            self.conv = nn.Conv3d(chan_in,
-                                  chan_out,
-                                  kernel_size,
-                                  stride=self.stride,
-                                  dilation=self.dilation,
-                                  padding=(0, 1, 1),
-                                  **kwargs)
-        else:
-            self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=self.stride, dilation=self.dilation, **kwargs)
-        self.is_first_run = True
-
-    def forward(self, x, is_init=True, residual=None):
-        if self.is_first_run:
-            self.is_first_run = False
-
-        if self.time_causal_padding == (1, 1, 1, 1, 2, 0):
-            pass
-        else:
-            x = nn.functional.pad(x, self.time_causal_padding).contiguous()
-
-        x = base_conv3d_channel_last(x, self.conv, residual=residual)
-        return x
-
-
-class AttnBlock(nn.Module):
-
-    def __init__(self, in_channels):
-        super().__init__()
-
-        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels)
-        self.q = CausalConvChannelLast(in_channels, in_channels, kernel_size=1)
-        self.k = CausalConvChannelLast(in_channels, in_channels, kernel_size=1)
-        self.v = CausalConvChannelLast(in_channels, in_channels, kernel_size=1)
-        self.proj_out = CausalConvChannelLast(in_channels, in_channels, kernel_size=1)
-
-    def attention(self, x, is_init=True):
-        x = base_group_norm(x, self.norm, act_silu=False, channel_last=True)
-        q = self.q(x, is_init)
-        k = self.k(x, is_init)
-        v = self.v(x, is_init)
-
-        b, t, h, w, c = q.shape
-        q, k, v = map(lambda x: rearrange(x, "b t h w c -> b 1 (t h w) c"), (q, k, v))
-        x = nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True)
-        x = rearrange(x, "b 1 (t h w) c -> b t h w c", t=t, h=h, w=w)
-
-        return x
-
-    def forward(self, x):
-        x = x.permute(0, 2, 3, 4, 1).contiguous()
-        h = self.attention(x)
-        x = self.proj_out(h, residual=x)
-        x = x.permute(0, 4, 1, 2, 3)
-        return x
-
-
-class Resnet3DBlock(nn.Module):
-
-    def __init__(
-        self,
-        in_channels,
-        out_channels=None,
-        temb_channels=512,
-        conv_shortcut=False,
-    ):
-        super().__init__()
-
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-
-        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels)
-        self.conv1 = CausalConvAfterNorm(in_channels, out_channels, kernel_size=3)
-        if temb_channels > 0:
-            self.temb_proj = nn.Linear(temb_channels, out_channels)
-
-        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels)
-        self.conv2 = CausalConvAfterNorm(out_channels, out_channels, kernel_size=3)
-
-        assert conv_shortcut is False
-        self.use_conv_shortcut = conv_shortcut
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                self.conv_shortcut = CausalConvAfterNorm(in_channels, out_channels, kernel_size=3)
-            else:
-                self.nin_shortcut = CausalConvAfterNorm(in_channels, out_channels, kernel_size=1)
-
-    def forward(self, x, temb=None, is_init=True):
-        x = x.permute(0, 2, 3, 4, 1).contiguous()
-
-        h = base_group_norm_with_zero_pad(x, self.norm1, act_silu=True, pad_size=2)
-        h = self.conv1(h)
-        if temb is not None:
-            h = h + self.temb_proj(nn.functional.silu(temb))[:, :, None, None]
-
-        x = self.nin_shortcut(x) if self.in_channels != self.out_channels else x
-
-        h = base_group_norm_with_zero_pad(h, self.norm2, act_silu=True, pad_size=2)
-        x = self.conv2(h, residual=x)
-
-        x = x.permute(0, 4, 1, 2, 3)
-        return x
-
-
-class Downsample3D(nn.Module):
-
-    def __init__(self, in_channels, with_conv, stride):
-        super().__init__()
-
-        self.with_conv = with_conv
-        if with_conv:
-            self.conv = CausalConv(in_channels, in_channels, kernel_size=3, stride=stride)
-
-    def forward(self, x, is_init=True):
-        if self.with_conv:
-            x = self.conv(x, is_init)
-        else:
-            x = nn.functional.avg_pool3d(x, kernel_size=2, stride=2)
-        return x
-
-
-class VideoEncoder(nn.Module):
-
-    def __init__(
-        self,
-        ch=32,
-        ch_mult=(4, 8, 16, 16),
-        num_res_blocks=2,
-        in_channels=3,
-        z_channels=16,
-        double_z=True,
-        down_sampling_layer=[1, 2],
-        resamp_with_conv=True,
-        version=1,
-    ):
-        super().__init__()
-
-        temb_ch = 0
-
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-
-        # downsampling
-        self.conv_in = CausalConv(in_channels, ch, kernel_size=3)
-        self.down_sampling_layer = down_sampling_layer
-
-        in_ch_mult = (1, ) + tuple(ch_mult)
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch * in_ch_mult[i_level]
-            block_out = ch * ch_mult[i_level]
-            for i_block in range(self.num_res_blocks):
-                block.append(Resnet3DBlock(in_channels=block_in, out_channels=block_out, temb_channels=temb_ch))
-                block_in = block_out
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions - 1:
-                if i_level in self.down_sampling_layer:
-                    down.downsample = Downsample3D(block_in, resamp_with_conv, stride=(2, 2, 2))
-                else:
-                    down.downsample = Downsample2D(block_in, resamp_with_conv, padding=0)  #DIFF
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch)
-
-        # end
-        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in)
-        self.version = version
-        if version == 2:
-            channels = 4 * z_channels * 2**3
-            self.conv_patchify = ConvPixelUnshuffleDownSampleLayer3D(block_in, channels, kernel_size=3, factor=2)
-            self.shortcut_pathify = PixelUnshuffleChannelAveragingDownSampleLayer3D(block_in, channels, 2)
-            self.shortcut_out = PixelUnshuffleChannelAveragingDownSampleLayer3D(
-                channels, 2 * z_channels if double_z else z_channels, 1)
-            self.conv_out = CausalConvChannelLast(channels, 2 * z_channels if double_z else z_channels, kernel_size=3)
-        else:
-            self.conv_out = CausalConvAfterNorm(block_in, 2 * z_channels if double_z else z_channels, kernel_size=3)
-
-    @torch.inference_mode()
-    def forward(self, x, video_frame_num, is_init=True):
-        # timestep embedding
-        temb = None
-
-        t = video_frame_num
-
-        # downsampling
-        h = self.conv_in(x, is_init)
-
-        # make it real channel last, but behave like normal layout
-        h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3)
-
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](h, temb, is_init)
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-
-            if i_level != self.num_resolutions - 1:
-                if isinstance(self.down[i_level].downsample, Downsample2D):
-                    _, _, t, _, _ = h.shape
-                    h = rearrange(h, "b c t h w -> (b t) h w c", t=t)
-                    h = self.down[i_level].downsample(h)
-                    h = rearrange(h, "(b t) h w c -> b c t h w", t=t)
-                else:
-                    h = self.down[i_level].downsample(h, is_init)
-
-        h = self.mid.block_1(h, temb, is_init)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb, is_init)
-
-        h = h.permute(0, 2, 3, 4, 1).contiguous()  # b c l h w -> b l h w c
-        if self.version == 2:
-            h = base_group_norm(h, self.norm_out, act_silu=True, channel_last=True)
-            h = h.permute(0, 4, 1, 2, 3).contiguous()
-            shortcut = self.shortcut_pathify(h, is_init)
-            h = self.conv_patchify(h, is_init)
-            h = h.add_(shortcut)
-            shortcut = self.shortcut_out(h, is_init).permute(0, 2, 3, 4, 1)
-            h = self.conv_out(h.permute(0, 2, 3, 4, 1).contiguous(), is_init)
-            h = h.add_(shortcut)
-        else:
-            h = base_group_norm_with_zero_pad(h, self.norm_out, act_silu=True, pad_size=2)
-            h = self.conv_out(h, is_init)
-        h = h.permute(0, 4, 1, 2, 3)  # b l h w c -> b c l h w
-
-        h = rearrange(h, "b c t h w -> b t c h w")
-        return h
-
-
-class Res3DBlockUpsample(nn.Module):
-
-    def __init__(self, input_filters, num_filters, down_sampling_stride, down_sampling=False):
-        super().__init__()
-
-        self.input_filters = input_filters
-        self.num_filters = num_filters
-
-        self.act_ = nn.SiLU(inplace=True)
-
-        self.conv1 = CausalConvChannelLast(num_filters, num_filters, kernel_size=[3, 3, 3])
-        self.norm1 = nn.GroupNorm(32, num_filters)
-
-        self.conv2 = CausalConvChannelLast(num_filters, num_filters, kernel_size=[3, 3, 3])
-        self.norm2 = nn.GroupNorm(32, num_filters)
-
-        self.down_sampling = down_sampling
-        if down_sampling:
-            self.down_sampling_stride = down_sampling_stride
-        else:
-            self.down_sampling_stride = [1, 1, 1]
-
-        if num_filters != input_filters or down_sampling:
-            self.conv3 = CausalConvChannelLast(input_filters,
-                                               num_filters,
-                                               kernel_size=[1, 1, 1],
-                                               stride=self.down_sampling_stride)
-            self.norm3 = nn.GroupNorm(32, num_filters)
-
-    def forward(self, x, is_init=False):
-        x = x.permute(0, 2, 3, 4, 1).contiguous()
-
-        residual = x
-
-        h = self.conv1(x, is_init)
-        h = base_group_norm(h, self.norm1, act_silu=True, channel_last=True)
-
-        h = self.conv2(h, is_init)
-        h = base_group_norm(h, self.norm2, act_silu=False, channel_last=True)
-
-        if self.down_sampling or self.num_filters != self.input_filters:
-            x = self.conv3(x, is_init)
-            x = base_group_norm(x, self.norm3, act_silu=False, channel_last=True)
-
-        h.add_(x)
-        h = self.act_(h)
-        if residual is not None:
-            h.add_(residual)
-
-        h = h.permute(0, 4, 1, 2, 3)
-        return h
-
-
-class Upsample3D(nn.Module):
-
-    def __init__(self, in_channels, scale_factor=2):
-        super().__init__()
-
-        self.scale_factor = scale_factor
-        self.conv3d = Res3DBlockUpsample(input_filters=in_channels,
-                                         num_filters=in_channels,
-                                         down_sampling_stride=(1, 1, 1),
-                                         down_sampling=False)
-
-    def forward(self, x, is_init=True, is_split=True):
-        b, c, t, h, w = x.shape
-
-        # x = x.permute(0,2,3,4,1).contiguous().permute(0,4,1,2,3).to(memory_format=torch.channels_last_3d)
-        if is_split:
-            split_size = c // 8
-            x_slices = torch.split(x, split_size, dim=1)
-            x = [nn.functional.interpolate(x, scale_factor=self.scale_factor) for x in x_slices]
-            x = torch.cat(x, dim=1)
-        else:
-            x = nn.functional.interpolate(x, scale_factor=self.scale_factor)
-
-        x = self.conv3d(x, is_init)
-        return x
-
-
-class VideoDecoder(nn.Module):
-
-    def __init__(
-        self,
-        ch=128,
-        z_channels=16,
-        out_channels=3,
-        ch_mult=(1, 2, 4, 4),
-        num_res_blocks=2,
-        temporal_up_layers=[2, 3],
-        temporal_downsample=4,
-        resamp_with_conv=True,
-        version=1,
-    ):
-        super().__init__()
-
-        temb_ch = 0
-
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.temporal_downsample = temporal_downsample
-
-        block_in = ch * ch_mult[self.num_resolutions - 1]
-        self.version = version
-        if version == 2:
-            channels = 4 * z_channels * 2**3
-            self.conv_in = CausalConv(z_channels, channels, kernel_size=3)
-            self.shortcut_in = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(z_channels, channels, 1)
-            self.conv_unpatchify = ConvPixelShuffleUpSampleLayer3D(channels, block_in, kernel_size=3, factor=2)
-            self.shortcut_unpathify = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(channels, block_in, 2)
-        else:
-            self.conv_in = CausalConv(z_channels, block_in, kernel_size=3)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch)
-
-        # upsampling
-        self.up_id = len(temporal_up_layers)
-        self.video_frame_num = 1
-        self.cur_video_frame_num = self.video_frame_num // 2**self.up_id + 1
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch * ch_mult[i_level]
-            for i_block in range(self.num_res_blocks + 1):
-                block.append(Resnet3DBlock(in_channels=block_in, out_channels=block_out, temb_channels=temb_ch))
-                block_in = block_out
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                if i_level in temporal_up_layers:
-                    up.upsample = Upsample3D(block_in)
-                    self.cur_video_frame_num = self.cur_video_frame_num * 2
-                else:
-                    up.upsample = Upsample2D(block_in, resamp_with_conv)
-            self.up.insert(0, up)  # prepend to get consistent order
-
-        # end
-        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in)
-        self.conv_out = CausalConvAfterNorm(block_in, out_channels, kernel_size=3)
-
-    @torch.inference_mode()
-    def forward(self, z, is_init=True):
-        z = rearrange(z, "b t c h w -> b c t h w")
-
-        h = self.conv_in(z, is_init=is_init)
-        if self.version == 2:
-            shortcut = self.shortcut_in(z, is_init=is_init)
-            h = h.add_(shortcut)
-            shortcut = self.shortcut_unpathify(h, is_init=is_init)
-            h = self.conv_unpatchify(h, is_init=is_init)
-            h = h.add_(shortcut)
-
-        temb = None
-
-        h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3)
-        h = self.mid.block_1(h, temb, is_init=is_init)
-        h = self.mid.attn_1(h)
-        h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3)
-        h = self.mid.block_2(h, temb, is_init=is_init)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks + 1):
-                h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3)
-                h = self.up[i_level].block[i_block](h, temb, is_init=is_init)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                if isinstance(self.up[i_level].upsample, Upsample2D):
-                    B = h.size(0)
-                    h = h.permute(0, 2, 3, 4, 1).flatten(0, 1)
-                    h = self.up[i_level].upsample(h)
-                    h = h.unflatten(0, (B, -1)).permute(0, 4, 1, 2, 3)
-                else:
-                    h = self.up[i_level].upsample(h, is_init=is_init)
-
-        # end
-        h = h.permute(0, 2, 3, 4, 1)  # b c l h w -> b l h w c
-        h = base_group_norm_with_zero_pad(h, self.norm_out, act_silu=True, pad_size=2)
-        h = self.conv_out(h)
-        h = h.permute(0, 4, 1, 2, 3)
-
-        if is_init:
-            h = h[:, :, (self.temporal_downsample - 1):]
-        return h
-
-
-def rms_norm(input, normalized_shape, eps=1e-6):
-    dtype = input.dtype
-    input = input.to(torch.float32)
-    variance = input.pow(2).flatten(-len(normalized_shape)).mean(-1)[(..., ) + (None, ) * len(normalized_shape)]
-    input = input * torch.rsqrt(variance + eps)
-    return input.to(dtype)
-
-
-class DiagonalGaussianDistribution(object):
-
-    def __init__(self, parameters, deterministic=False, rms_norm_mean=False, only_return_mean=False):
-        self.parameters = parameters
-        self.mean, self.logvar = torch.chunk(parameters, 2, dim=-3)  #N,[X],C,H,W
-        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
-        self.std = torch.exp(0.5 * self.logvar)
-        self.var = torch.exp(self.logvar)
-        self.deterministic = deterministic
-        if self.deterministic:
-            self.var = self.std = torch.zeros_like(self.mean,
-                                                   device=self.parameters.device,
-                                                   dtype=self.parameters.dtype)
-        if rms_norm_mean:
-            self.mean = rms_norm(self.mean, self.mean.size()[1:])
-        self.only_return_mean = only_return_mean
-
-    def sample(self, generator=None):
-        # make sure sample is on the same device
-        # as the parameters and has same dtype
-        sample = torch.randn(self.mean.shape, generator=generator, device=self.parameters.device)
-        sample = sample.to(dtype=self.parameters.dtype)
-        x = self.mean + self.std * sample
-        if self.only_return_mean:
-            return self.mean
-        else:
-            return x
-
-
-class AutoencoderKL(nn.Module):
-
-    @with_empty_init
-    def __init__(
-        self,
-        in_channels=3,
-        out_channels=3,
-        z_channels=16,
-        num_res_blocks=2,
-        model_path=None,
-        weight_dict={},
-        world_size=1,
-        version=1,
-    ):
-        super().__init__()
-
-        self.frame_len = 17
-        self.latent_len = 3 if version == 2 else 5
-
-        base_group_norm.spatial = True if version == 2 else False
-
-        self.encoder = VideoEncoder(
-            in_channels=in_channels,
-            z_channels=z_channels,
-            num_res_blocks=num_res_blocks,
-            version=version,
-        )
-
-        self.decoder = VideoDecoder(
-            z_channels=z_channels,
-            out_channels=out_channels,
-            num_res_blocks=num_res_blocks,
-            version=version,
-        )
-
-        if model_path is not None:
-            weight_dict = self.init_from_ckpt(model_path)
-        if len(weight_dict) != 0:
-            self.load_from_dict(weight_dict)
-        self.convert_channel_last()
-
-        self.world_size = world_size
-
-    def init_from_ckpt(self, model_path):
-        from safetensors import safe_open
-        p = {}
-        with safe_open(model_path, framework="pt", device="cpu") as f:
-            for k in f.keys():
-                tensor = f.get_tensor(k)
-                if k.startswith("decoder.conv_out."):
-                    k = k.replace("decoder.conv_out.", "decoder.conv_out.conv.")
-                p[k] = tensor
-        return p
-
-    def load_from_dict(self, p):
-        self.load_state_dict(p)
-
-    def convert_channel_last(self):
-        #Conv2d NCHW->NHWC
-        pass
-
-    def naive_encode(self, x, is_init_image=True):
-        b, len, c, h, w = x.size()
-        x = rearrange(x, 'b l c h w -> b c l h w').contiguous()
-        z = self.encoder(x, len, True)  # 下采样[1, 4, 8, 16, 16]
-        return z
-
-    @torch.inference_mode()
-    def encode(self, x):
-        # b (nc cf) c h w -> (b nc) cf c h w -> encode -> (b nc) cf c h w -> b (nc cf) c h w
-        chunks = list(x.split(self.frame_len, dim=1))
-        for i in range(len(chunks)):
-            chunks[i] = self.naive_encode(chunks[i], True)
-        z = torch.cat(chunks, dim=1)
-
-        posterior = DiagonalGaussianDistribution(z)
-        return posterior.sample()
-
-    def decode_naive(self, z, is_init=True):
-        z = z.to(next(self.decoder.parameters()).dtype)
-        dec = self.decoder(z, is_init)
-        return dec
-
-    @torch.inference_mode()
-    def decode(self, z):
-        # b (nc cf) c h w -> (b nc) cf c h w -> decode -> (b nc) c cf h w -> b (nc cf) c h w
-        chunks = list(z.split(self.latent_len, dim=1))
-
-        if self.world_size > 1:
-            chunks_total_num = len(chunks)
-            max_num_per_rank = (chunks_total_num + self.world_size - 1) // self.world_size
-            rank = torch.distributed.get_rank()
-            chunks_ = chunks[max_num_per_rank * rank:max_num_per_rank * (rank + 1)]
-            if len(chunks_) < max_num_per_rank:
-                chunks_.extend(chunks[:max_num_per_rank - len(chunks_)])
-            chunks = chunks_
-
-        for i in range(len(chunks)):
-            chunks[i] = self.decode_naive(chunks[i], True).permute(0, 2, 1, 3, 4)
-        x = torch.cat(chunks, dim=1)
-
-        if self.world_size > 1:
-            x_ = torch.empty([x.size(0), (self.world_size * max_num_per_rank) * self.frame_len, *x.shape[2:]],
-                             dtype=x.dtype,
-                             device=x.device)
-            torch.distributed.all_gather_into_tensor(x_, x)
-            x = x_[:, :chunks_total_num * self.frame_len]
-
-        x = self.mix(x)
-        return x
-
-    def mix(self, x):
-        remain_scale = 0.6
-        mix_scale = 1. - remain_scale
-        front = slice(self.frame_len - 1, x.size(1) - 1, self.frame_len)
-        back = slice(self.frame_len, x.size(1), self.frame_len)
-        x[:, back] = x[:, back] * remain_scale + x[:, front] * mix_scale
-        x[:, front] = x[:, front] * remain_scale + x[:, back] * mix_scale
-        return x

fastvideo/models/text_encoder.py

@@ -0,0 +1,255 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from transformers.utils import ModelOutput
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+def use_default(value, default):
+    return value if value is not None else default
+
+@dataclass
+class TextEncoderModelOutput(ModelOutput):
+    """
+    Base class for model's outputs that also contains a pooling of the last hidden states.
+
+    Args:
+        hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
+        hidden_states_list (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        text_outputs (`list`, *optional*, returned when `return_texts=True` is passed):
+            List of decoded texts.
+    """
+
+    hidden_state: torch.FloatTensor = None
+    attention_mask: Optional[torch.LongTensor] = None
+    hidden_states_list: Optional[Tuple[torch.FloatTensor, ...]] = None
+    text_outputs: Optional[list] = None
+
+
+class TextEncoder(nn.Module):
+
+    def __init__(
+        self,
+        text_encoder,
+        tokenizer,
+        max_length: int,
+        text_encoder_precision: Optional[str] = None,
+        text_encoder_path: Optional[str] = None,
+        output_key: Optional[str] = None,
+        use_attention_mask: bool = True,
+        prompt_template: Optional[dict] = None,
+        prompt_template_video: Optional[dict] = None,
+        hidden_state_skip_layer: Optional[int] = None,
+        apply_final_norm: bool = False,
+        reproduce: bool = False,
+        device=None,
+    ):
+        super().__init__()
+        # TODO(will): check if there's a cleaner way to do this
+        self.text_encoder_type = text_encoder.config.architectures[0]
+        self.max_length = max_length
+        self.precision = text_encoder_precision
+        self.model_path = text_encoder_path
+        self.use_attention_mask = use_attention_mask
+        if prompt_template_video is not None:
+            assert (use_attention_mask is True), "Attention mask is True required when training videos."
+        self.prompt_template = prompt_template
+        self.prompt_template_video = prompt_template_video
+        self.hidden_state_skip_layer = hidden_state_skip_layer
+        self.apply_final_norm = apply_final_norm
+        self.reproduce = reproduce
+
+        self.use_template = self.prompt_template is not None
+        if self.use_template:
+            assert (isinstance(self.prompt_template, dict) and "template" in self.prompt_template
+                    ), f"`prompt_template` must be a dictionary with a key 'template', got {self.prompt_template}"
+            assert "{}" in str(self.prompt_template["template"]), (
+                "`prompt_template['template']` must contain a placeholder `{}` for the input text, "
+                f"got {self.prompt_template['template']}")
+
+        self.use_video_template = self.prompt_template_video is not None
+        if self.use_video_template:
+            if self.prompt_template_video is not None:
+                assert (
+                    isinstance(self.prompt_template_video, dict) and "template" in self.prompt_template_video
+                ), f"`prompt_template_video` must be a dictionary with a key 'template', got {self.prompt_template_video}"
+            assert "{}" in str(self.prompt_template_video["template"]), (
+                "`prompt_template_video['template']` must contain a placeholder `{}` for the input text, "
+                f"got {self.prompt_template_video['template']}")
+
+        if "T5" in self.text_encoder_type:
+            self.output_key = output_key or "last_hidden_state"
+        elif "CLIPTextModel" in self.text_encoder_type:
+            self.output_key = output_key or "pooler_output"
+        elif "LlamaModel" in self.text_encoder_type or "glm" in self.text_encoder_type:
+            self.output_key = output_key or "last_hidden_state"
+        else:
+            raise ValueError(f"Unsupported text encoder type: {self.text_encoder_type}")
+        
+        self.model = text_encoder
+        # self.dtype = self.model.dtype
+        self.device = device
+
+        self.tokenizer = tokenizer
+
+    def __repr__(self):
+        return f"{self.text_encoder_type} ({self.precision} - {self.model_path})"
+
+    @staticmethod
+    def apply_text_to_template(text, template, prevent_empty_text=True):
+        """
+        Apply text to template.
+
+        Args:
+            text (str): Input text.
+            template (str or list): Template string or list of chat conversation.
+            prevent_empty_text (bool): If True, we will prevent the user text from being empty
+                by adding a space. Defaults to True.
+        """
+        if isinstance(template, str):
+            # Will send string to tokenizer. Used for llm
+            return template.format(text)
+        else:
+            raise TypeError(f"Unsupported template type: {type(template)}")
+
+    def text2tokens(self, text, data_type="image"):
+        """
+        Tokenize the input text.
+
+        Args:
+            text (str or list): Input text.
+        """
+        tokenize_input_type = "str"
+        if self.use_template:
+            if data_type == "image":
+                prompt_template = self.prompt_template["template"]
+            elif data_type == "video":
+                prompt_template = self.prompt_template_video["template"]
+            else:
+                raise ValueError(f"Unsupported data type: {data_type}")
+            if isinstance(text, (list, tuple)):
+                text = [self.apply_text_to_template(one_text, prompt_template) for one_text in text]
+                if isinstance(text[0], list):
+                    tokenize_input_type = "list"
+            elif isinstance(text, str):
+                text = self.apply_text_to_template(text, prompt_template)
+                if isinstance(text, list):
+                    tokenize_input_type = "list"
+            else:
+                raise TypeError(f"Unsupported text type: {type(text)}")
+
+        kwargs = dict(
+            truncation=True,
+            max_length=self.max_length,
+            padding="max_length",
+            return_tensors="pt",
+        )
+        if tokenize_input_type == "str":
+            return self.tokenizer(
+                text,
+                return_length=False,
+                return_overflowing_tokens=False,
+                return_attention_mask=True,
+                **kwargs,
+            )
+        elif tokenize_input_type == "list":
+            return self.tokenizer.apply_chat_template(
+                text,
+                add_generation_prompt=True,
+                tokenize=True,
+                return_dict=True,
+                **kwargs,
+            )
+        else:
+            raise ValueError(f"Unsupported tokenize_input_type: {tokenize_input_type}")
+
+    def encode(
+        self,
+        batch_encoding,
+        use_attention_mask=None,
+        output_hidden_states=False,
+        do_sample=None,
+        hidden_state_skip_layer=None,
+        return_texts=False,
+        data_type="image",
+        device=None,
+    ):
+        """
+        Args:
+            batch_encoding (dict): Batch encoding from tokenizer.
+            use_attention_mask (bool): Whether to use attention mask. If None, use self.use_attention_mask.
+                Defaults to None.
+            output_hidden_states (bool): Whether to output hidden states. If False, return the value of
+                self.output_key. If True, return the entire output. If set self.hidden_state_skip_layer,
+                output_hidden_states will be set True. Defaults to False.
+            do_sample (bool): Whether to sample from the model. Used for Decoder-Only LLMs. Defaults to None.
+                When self.produce is False, do_sample is set to True by default.
+            hidden_state_skip_layer (int): Number of hidden states to hidden_state_skip_layer. 0 means the last layer.
+                If None, self.output_key will be used. Defaults to None.
+            return_texts (bool): Whether to return the decoded texts. Defaults to False.
+        """
+        device = self.model.device if device is None else device
+        use_attention_mask = use_default(use_attention_mask, self.use_attention_mask)
+        hidden_state_skip_layer = use_default(hidden_state_skip_layer, self.hidden_state_skip_layer)
+        do_sample = use_default(do_sample, not self.reproduce)
+        attention_mask = (batch_encoding["attention_mask"].to(device) if use_attention_mask else None)
+        print(f"batch_encoding: {batch_encoding}")
+        # logger.info(f"attention_mask: {attention_mask}")
+        outputs = self.model(
+            input_ids=batch_encoding["input_ids"].to(device),
+            # attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states or hidden_state_skip_layer is not None,
+        )
+        if hidden_state_skip_layer is not None:
+            last_hidden_state = outputs.hidden_states[-(hidden_state_skip_layer + 1)]
+            # Real last hidden state already has layer norm applied. So here we only apply it
+            # for intermediate layers.
+            if hidden_state_skip_layer > 0 and self.apply_final_norm:
+                last_hidden_state = self.model.final_layer_norm(last_hidden_state)
+        else:
+            last_hidden_state = outputs[self.output_key]
+
+        # Remove hidden states of instruction tokens, only keep prompt tokens.
+        if self.use_template:
+            if data_type == "image":
+                crop_start = self.prompt_template.get("crop_start", -1)
+            elif data_type == "video":
+                crop_start = self.prompt_template_video.get("crop_start", -1)
+            else:
+                raise ValueError(f"Unsupported data type: {data_type}")
+            if crop_start > 0:
+                last_hidden_state = last_hidden_state[:, crop_start:]
+                attention_mask = (attention_mask[:, crop_start:] if use_attention_mask else None)
+
+        if output_hidden_states:
+            return TextEncoderModelOutput(last_hidden_state, attention_mask, outputs.hidden_states)
+        return TextEncoderModelOutput(last_hidden_state, attention_mask)
+
+    def forward(
+        self,
+        text,
+        use_attention_mask=None,
+        output_hidden_states=False,
+        do_sample=False,
+        hidden_state_skip_layer=None,
+        return_texts=False,
+    ):
+        batch_encoding = self.text2tokens(text)
+        return self.encode(
+            batch_encoding,
+            use_attention_mask=use_attention_mask,
+            output_hidden_states=output_hidden_states,
+            do_sample=do_sample,
+            hidden_state_skip_layer=hidden_state_skip_layer,
+            return_texts=return_texts,
+        )

fastvideo/models/utils.py

@@ -5,6 +5,49 @@ from typing import Any, Dict, Optional
 
 import torch
 
+# TODO(PY): move it elsewhere 
+def auto_attributes(init_func):
+    """
+    Decorator that automatically adds all initialization arguments as object attributes.
+    
+    Example:
+        @auto_attributes
+        def __init__(self, a=1, b=2):
+            pass
+        
+        # This will automatically set:
+        # - self.a = 1 and self.b = 2
+        # - self.config.a = 1 and self.config.b = 2
+    """
+    def wrapper(self, *args, **kwargs):
+        # Get the function signature
+        import inspect
+        signature = inspect.signature(init_func)
+        parameters = signature.parameters
+        
+        # Get parameter names (excluding 'self')
+        param_names = list(parameters.keys())[1:]
+        
+        # Bind arguments to parameters
+        bound_args = signature.bind(self, *args, **kwargs)
+        bound_args.apply_defaults()
+        
+        # Create config object if it doesn't exist
+        if not hasattr(self, 'config'):
+            self.config = type('Config', (), {})()
+        
+        # Set attributes on self and self.config
+        for name in param_names:
+            if name in bound_args.arguments:
+                value = bound_args.arguments[name]
+                setattr(self, name, value)
+                setattr(self.config, name, value)
+        
+        # Call the original __init__ function
+        return init_func(self, *args, **kwargs)
+    
+    return wrapper
+
 
 def set_random_seed(seed: int) -> None:
     from fastvideo.platforms import current_platform

fastvideo/models/vaes/common.py

@@ -0,0 +1,469 @@
+import torch
+from typing import Optional, Tuple
+import numpy as np
+import torch.nn as nn
+from abc import abstractmethod, ABC
+import torch.distributed as dist
+from fastvideo.distributed import tensor_model_parallel_all_gather, get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size
+from math import prod
+class ParallelTiledVAE(ABC):
+    def __init__(self, *args, **kwargs):
+        # Check if subclass has defined all required properties
+        required_attributes = [
+            'tile_sample_min_height',
+            'tile_sample_min_width',
+            'tile_sample_min_num_frames',
+            'tile_sample_stride_height',
+            'tile_sample_stride_width',
+            'tile_sample_stride_num_frames',
+            'use_tiling'
+        ]
+        
+        for attr in required_attributes:
+            if not hasattr(self, attr):
+                raise AttributeError(f"Subclasses of ParallelVAE must define '{attr}' property")
+        
+    @abstractmethod
+    def _encode(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def _decode(self, *args, **kwargs):
+        pass
+    
+    
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = x.shape
+
+        if self.use_tiling and num_frames > self.tile_sample_min_num_frames:
+            latents = self.tiled_encode(x)
+        elif self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height):
+            latents = self.spatial_tiled_encode(x)
+        else:
+            latents = self._encode(x)
+        return DiagonalGaussianDistribution(latents)
+
+    
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = z.shape
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+        tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
+        
+        if self.use_tiling and get_sequence_model_parallel_world_size() > 1:
+            return self.parallel_tiled_decode(z)
+        if self.use_tiling and num_frames > tile_latent_min_num_frames:
+            return self.tiled_decode(z)
+
+        if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height):
+            return self.spatial_tiled_decode(z)
+
+        return self._decode(z)
+
+
+    def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (
+                y / blend_extent
+            )
+        return b
+
+    def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (
+                x / blend_extent
+            )
+        return b
+
+    def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-3], b.shape[-3], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * (
+                x / blend_extent
+            )
+        return b
+
+    def spatial_tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Encode a batch of images using a tiled encoder.
+
+        Args:
+            x (`torch.Tensor`): Input batch of videos.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+        batch_size, num_channels, num_frames, height, width = x.shape
+        latent_height = height // self.spatial_compression_ratio
+        latent_width = width // self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        # Split x into overlapping tiles and encode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            row = []
+            for j in range(0, width, self.tile_sample_stride_width):
+                tile = x[:, :, :, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width]
+                tile = self._encode(tile)
+                row.append(tile)
+            rows.append(row)
+
+        return self._merge_spatial_tiles(rows, blend_height, blend_width,tile_latent_stride_height, tile_latent_stride_width)
+    
+    def _parallel_data_generator(self, gathered_results,
+                                 gathered_dim_metadata):
+        global_idx = 0
+        for i, per_rank_metadata in enumerate(gathered_dim_metadata):
+            _start_shape = 0
+            for shape in per_rank_metadata:
+                mul_shape = prod(shape)
+                yield (gathered_results[i, _start_shape:_start_shape +
+                                        mul_shape].reshape(shape), global_idx)
+                _start_shape += mul_shape
+                global_idx += 1
+
+    def parallel_tiled_decode(self, z: torch.FloatTensor) -> torch.FloatTensor:
+        """
+        Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs
+        """
+        world_size, rank = get_sequence_model_parallel_world_size(), get_sequence_model_parallel_rank()
+        B, C, T, H, W = z.shape
+
+        # Calculate parameters
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+        tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio
+        
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+        blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
+
+        # Calculate tile dimensions
+        num_t_tiles = (T + tile_latent_stride_num_frames - 1) // tile_latent_stride_num_frames
+        num_h_tiles = (H + tile_latent_stride_height - 1) // tile_latent_stride_height
+        num_w_tiles = (W + tile_latent_stride_width - 1) // tile_latent_stride_width
+        total_spatial_tiles = num_h_tiles * num_w_tiles
+        total_tiles = num_t_tiles * total_spatial_tiles
+
+        # Calculate tiles per rank and padding
+        tiles_per_rank = (total_tiles + world_size - 1) // world_size
+        start_tile_idx = rank * tiles_per_rank
+        end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles)
+
+        local_results = []
+        local_dim_metadata = []
+        # Process assigned tiles
+        for local_idx, global_idx in enumerate(range(start_tile_idx, end_tile_idx)):
+            t_idx = global_idx // total_spatial_tiles
+            spatial_idx = global_idx % total_spatial_tiles
+            h_idx = spatial_idx // num_w_tiles
+            w_idx = spatial_idx % num_w_tiles
+
+            # Calculate positions
+            t_start = t_idx * tile_latent_stride_num_frames
+            h_start = h_idx * tile_latent_stride_height
+            w_start = w_idx * tile_latent_stride_width
+
+            # Extract and process tile
+            tile = z[:, :, t_start:t_start + tile_latent_min_num_frames + 1,
+                     h_start:h_start + tile_latent_min_height,
+                     w_start:w_start + tile_latent_min_width]
+
+            # Process tile
+            tile = self._decode(tile)
+
+            if t_start > 0:
+                tile = tile[:, :, 1:, :, :]
+
+            # Store metadata
+            shape = tile.shape
+            # Store decoded data (flattened)
+            decoded_flat = tile.reshape(-1)
+            local_results.append(decoded_flat)
+            local_dim_metadata.append(shape)
+
+        results = torch.cat(local_results, dim=0).contiguous()
+        del local_results
+        torch.cuda.empty_cache()
+        # first gather size to pad the results
+        local_size = torch.tensor([results.size(0)],
+                                  device=results.device,
+                                  dtype=torch.int64)
+        all_sizes = [
+            torch.zeros(1, device=results.device, dtype=torch.int64)
+            for _ in range(world_size)
+        ]
+        dist.all_gather(all_sizes, local_size)
+        max_size = max(size.item() for size in all_sizes)
+        padded_results = torch.zeros(max_size, device=results.device)
+        padded_results[:results.size(0)] = results
+        del results
+        torch.cuda.empty_cache()
+        # Gather all results
+        gathered_dim_metadata = [None] * world_size
+        gathered_results = torch.zeros_like(padded_results).repeat(
+            world_size, *[1] * len(padded_results.shape)
+        ).contiguous()  # use contiguous to make sure it won't copy data in the following operations
+        # TODO (PY): use fastvideo distributed methods
+        dist.all_gather_into_tensor(gathered_results, padded_results)
+        dist.all_gather_object(gathered_dim_metadata, local_dim_metadata)
+        # Process gathered results
+        data = [[[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)]
+                for _ in range(num_t_tiles)]
+        for current_data, global_idx in self._parallel_data_generator(
+                gathered_results, gathered_dim_metadata):
+            t_idx = global_idx // total_spatial_tiles
+            spatial_idx = global_idx % total_spatial_tiles
+            h_idx = spatial_idx // num_w_tiles
+            w_idx = spatial_idx % num_w_tiles
+            data[t_idx][h_idx][w_idx] = current_data
+        # Merge results
+        result_slices = []
+        last_slice_data = None
+        for i, tem_data in enumerate(data):
+            slice_data = self._merge_spatial_tiles(tem_data, blend_height, blend_width, 
+                                                  self.tile_sample_stride_height, self.tile_sample_stride_width)
+            if i > 0:
+                slice_data = self.blend_t(last_slice_data, slice_data, blend_num_frames)
+                result_slices.append(slice_data[:, :, :self.tile_sample_stride_num_frames, :, :])
+            else:
+                result_slices.append(slice_data[:, :, :self.tile_sample_stride_num_frames + 1, :, :])
+            last_slice_data = slice_data
+        dec = torch.cat(result_slices, dim=2)
+
+        return dec
+    
+    def _merge_spatial_tiles(self, tiles, blend_height, blend_width, stride_height, stride_width):
+        """Helper function to merge spatial tiles with blending"""
+        result_rows = []
+        for i, row in enumerate(tiles):
+            result_row = []
+            for j, tile in enumerate(row):
+                if i > 0:
+                    tile = self.blend_v(tiles[i - 1][j], tile,
+                                        blend_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(tile[:, :, :, :stride_height, :stride_width])
+            result_rows.append(torch.cat(result_row, dim=-1))
+        return torch.cat(result_rows, dim=-2)
+    
+    def spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+
+        Returns:
+            `torch.Tensor`:
+                The decoded images.
+        """
+
+        batch_size, num_channels, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, tile_latent_stride_height):
+            row = []
+            for j in range(0, width, tile_latent_stride_width):
+                tile = z[:, :, :, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                decoded = self._decode(tile)
+                row.append(decoded)
+            rows.append(row)
+        return self._merge_spatial_tiles(rows, blend_height, blend_width,self.tile_sample_stride_height, self.tile_sample_stride_width)
+
+    def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = x.shape
+        latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1
+
+        tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
+        tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio
+        blend_num_frames = tile_latent_min_num_frames - tile_latent_stride_num_frames
+
+        row = []
+        for i in range(0, num_frames, self.tile_sample_stride_num_frames):
+            tile = x[:, :, i : i + self.tile_sample_min_num_frames + 1, :, :]
+            if self.use_tiling and (height > self.tile_sample_min_height or width > self.tile_sample_min_width):
+                tile = self.spatial_tiled_encode(tile)
+            else:
+                tile = self._encode(tile)
+            if i > 0:
+                tile = tile[:, :, 1:, :, :]
+            row.append(tile)
+        result_row = []
+        for i, tile in enumerate(row):
+            if i > 0:
+                tile = self.blend_t(row[i - 1], tile, blend_num_frames)
+                result_row.append(tile[:, :, :tile_latent_stride_num_frames, :, :])
+            else:
+                result_row.append(tile[:, :, : tile_latent_stride_num_frames + 1, :, :])
+
+        enc = torch.cat(result_row, dim=2)[:, :, :latent_num_frames]
+        return enc
+
+    def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = z.shape
+        num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
+        tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio
+        blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
+
+        row = []
+        for i in range(0, num_frames, tile_latent_stride_num_frames):
+            tile = z[:, :, i : i + tile_latent_min_num_frames + 1, :, :]
+            if self.use_tiling and (tile.shape[-1] > tile_latent_min_width or tile.shape[-2] > tile_latent_min_height):
+                decoded = self.spatial_tiled_decode(tile)
+            else:
+                decoded = self._decode(tile)
+            if i > 0:
+                decoded = decoded[:, :, 1:, :, :]
+            row.append(decoded)
+        result_row = []
+        result_row = []
+        for i, tile in enumerate(row):
+            if i > 0:
+                tile = self.blend_t(row[i - 1], tile, blend_num_frames)
+                result_row.append(tile[:, :, : self.tile_sample_stride_num_frames, :, :])
+            else:
+                result_row.append(tile[:, :, : self.tile_sample_stride_num_frames + 1, :, :])
+
+        dec = torch.cat(result_row, dim=2)[:, :, :num_sample_frames]
+        return dec
+
+
+
+    def enable_tiling(
+        self,
+        tile_sample_min_height: Optional[int] = None,
+        tile_sample_min_width: Optional[int] = None,
+        tile_sample_min_num_frames: Optional[int] = None,
+        tile_sample_stride_height: Optional[float] = None,
+        tile_sample_stride_width: Optional[float] = None,
+        tile_sample_stride_num_frames: Optional[float] = None,
+    ) -> None:
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+
+        Args:
+            tile_sample_min_height (`int`, *optional*):
+                The minimum height required for a sample to be separated into tiles across the height dimension.
+            tile_sample_min_width (`int`, *optional*):
+                The minimum width required for a sample to be separated into tiles across the width dimension.
+            tile_sample_min_num_frames (`int`, *optional*):
+                The minimum number of frames required for a sample to be separated into tiles across the frame
+                dimension.
+            tile_sample_stride_height (`int`, *optional*):
+                The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are
+                no tiling artifacts produced across the height dimension.
+            tile_sample_stride_width (`int`, *optional*):
+                The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling
+                artifacts produced across the width dimension.
+            tile_sample_stride_num_frames (`int`, *optional*):
+                The stride between two consecutive frame tiles. This is to ensure that there are no tiling artifacts
+                produced across the frame dimension.
+        """
+        self.use_tiling = True
+        self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height
+        self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
+        self.tile_sample_min_num_frames = tile_sample_min_num_frames or self.tile_sample_min_num_frames
+        self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height
+        self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width
+        self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames
+        
+    
+        
+    def disable_tiling(self) -> None:
+        r"""
+        Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.use_tiling = False
+        
+        
+
+# adapted from https://github.com/huggingface/diffusers/blob/e7ffeae0a191f710881d1fbde00cd6ff025e81f2/src/diffusers/models/autoencoders/vae.py#L691
+class DiagonalGaussianDistribution(object):
+    def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(
+                self.mean, device=self.parameters.device, dtype=self.parameters.dtype
+            )
+
+    def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor:
+        # make sure sample is on the same device as the parameters and has same dtype
+        sample = torch.randn(
+            self.mean.shape,
+            generator=generator,
+            device=self.parameters.device,
+            dtype=self.parameters.dtype,
+        )
+        x = self.mean + self.std * sample
+        return x
+
+    def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        else:
+            if other is None:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+                    dim=[1, 2, 3],
+                )
+            else:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean - other.mean, 2) / other.var
+                    + self.var / other.var
+                    - 1.0
+                    - self.logvar
+                    + other.logvar,
+                    dim=[1, 2, 3],
+                )
+
+    def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims,
+        )
+
+    def mode(self) -> torch.Tensor:
+        return self.mean

fastvideo/models/vaes/hunyuanvae.py

@@ -0,0 +1,799 @@
+# Copyright 2024 The Hunyuan Team, The HuggingFace Team and The FastVideo Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+
+from fastvideo.models.vaes.common import DiagonalGaussianDistribution, ParallelTiledVAE
+from fastvideo.layers.activation import get_act_fn
+
+from fastvideo.models.utils import auto_attributes
+
+def prepare_causal_attention_mask(
+    num_frames: int, height_width: int, dtype: torch.dtype, device: torch.device, batch_size: int = None
+) -> torch.Tensor:
+    indices = torch.arange(1, num_frames + 1, dtype=torch.int32, device=device)
+    indices_blocks = indices.repeat_interleave(height_width)
+    x, y = torch.meshgrid(indices_blocks, indices_blocks, indexing="xy")
+    mask = torch.where(x <= y, 0, -float("inf")).to(dtype=dtype)
+
+    if batch_size is not None:
+        mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
+    return mask
+
+class HunyuanVAEAttention(nn.Module):
+    def __init__(self, in_channels, heads, dim_head, eps, norm_num_groups, bias):
+        super().__init__()
+        self.in_channels = in_channels
+        self.heads = heads
+        self.dim_head = dim_head
+        self.eps = eps
+        self.norm_num_groups = norm_num_groups
+        self.bias = bias
+        
+        inner_dim = heads * dim_head
+        
+        # Define the projection layers
+        self.to_q = nn.Linear(in_channels, inner_dim, bias=bias)
+        self.to_k = nn.Linear(in_channels, inner_dim, bias=bias)
+        self.to_v = nn.Linear(in_channels, inner_dim, bias=bias)
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, in_channels, bias=bias)
+        )
+        
+        # Optional normalization layers
+        self.group_norm = nn.GroupNorm(norm_num_groups, in_channels, eps=eps, affine=True)
+        
+    def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        residual = hidden_states
+        
+        batch_size, sequence_length, _ = hidden_states.shape
+        
+        hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+        
+        # Project to query, key, value
+        query = self.to_q(hidden_states)
+        key = self.to_k(hidden_states)
+        value = self.to_v(hidden_states)
+        
+        # Reshape for multi-head attention
+        head_dim = self.dim_head
+        
+        query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        
+        # Perform scaled dot-product attention
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+        
+        # Reshape back
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+        
+        # Linear projection
+        hidden_states = self.to_out(hidden_states)
+        
+        # Residual connection and rescale
+        hidden_states = hidden_states + residual
+        
+        return hidden_states
+
+class HunyuanVideoCausalConv3d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int, int, int]] = 3,
+        stride: Union[int, Tuple[int, int, int]] = 1,
+        padding: Union[int, Tuple[int, int, int]] = 0,
+        dilation: Union[int, Tuple[int, int, int]] = 1,
+        bias: bool = True,
+        pad_mode: str = "replicate",
+    ) -> None:
+        super().__init__()
+
+        kernel_size = (kernel_size, kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
+
+        self.pad_mode = pad_mode
+        self.time_causal_padding = (
+            kernel_size[0] // 2,
+            kernel_size[0] // 2,
+            kernel_size[1] // 2,
+            kernel_size[1] // 2,
+            kernel_size[2] - 1,
+            0,
+        )
+
+        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = F.pad(hidden_states, self.time_causal_padding, mode=self.pad_mode)
+        return self.conv(hidden_states)
+
+
+class HunyuanVideoUpsampleCausal3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        kernel_size: int = 3,
+        stride: int = 1,
+        bias: bool = True,
+        upsample_factor: Tuple[float, float, float] = (2, 2, 2),
+    ) -> None:
+        super().__init__()
+
+        out_channels = out_channels or in_channels
+        self.upsample_factor = upsample_factor
+
+        self.conv = HunyuanVideoCausalConv3d(in_channels, out_channels, kernel_size, stride, bias=bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        num_frames = hidden_states.size(2)
+
+        first_frame, other_frames = hidden_states.split((1, num_frames - 1), dim=2)
+        first_frame = F.interpolate(
+            first_frame.squeeze(2), scale_factor=self.upsample_factor[1:], mode="nearest"
+        ).unsqueeze(2)
+
+        if num_frames > 1:
+            # See: https://github.com/pytorch/pytorch/issues/81665
+            # Unless you have a version of pytorch where non-contiguous implementation of F.interpolate
+            # is fixed, this will raise either a runtime error, or fail silently with bad outputs.
+            # If you are encountering an error here, make sure to try running encoding/decoding with
+            # `vae.enable_tiling()` first. If that doesn't work, open an issue at:
+            # https://github.com/huggingface/diffusers/issues
+            other_frames = other_frames.contiguous()
+            other_frames = F.interpolate(other_frames, scale_factor=self.upsample_factor, mode="nearest")
+            hidden_states = torch.cat((first_frame, other_frames), dim=2)
+        else:
+            hidden_states = first_frame
+
+        hidden_states = self.conv(hidden_states)
+        return hidden_states
+
+
+class HunyuanVideoDownsampleCausal3D(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        padding: int = 1,
+        kernel_size: int = 3,
+        bias: bool = True,
+        stride=2,
+    ) -> None:
+        super().__init__()
+        out_channels = out_channels or channels
+
+        self.conv = HunyuanVideoCausalConv3d(channels, out_channels, kernel_size, stride, padding, bias=bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv(hidden_states)
+        return hidden_states
+
+
+class HunyuanVideoResnetBlockCausal3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        dropout: float = 0.0,
+        groups: int = 32,
+        eps: float = 1e-6,
+        non_linearity: str = "silu",
+    ) -> None:
+        super().__init__()
+        out_channels = out_channels or in_channels
+
+        self.nonlinearity = get_act_fn(non_linearity)
+
+        self.norm1 = nn.GroupNorm(groups, in_channels, eps=eps, affine=True)
+        self.conv1 = HunyuanVideoCausalConv3d(in_channels, out_channels, 3, 1, 0)
+
+        self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps, affine=True)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = HunyuanVideoCausalConv3d(out_channels, out_channels, 3, 1, 0)
+
+        self.conv_shortcut = None
+        if in_channels != out_channels:
+            self.conv_shortcut = HunyuanVideoCausalConv3d(in_channels, out_channels, 1, 1, 0)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.contiguous()
+        residual = hidden_states
+
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.conv1(hidden_states)
+
+        hidden_states = self.norm2(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        if self.conv_shortcut is not None:
+            residual = self.conv_shortcut(residual)
+
+        hidden_states = hidden_states + residual
+        return hidden_states
+
+
+class HunyuanVideoMidBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_act_fn: str = "silu",
+        resnet_groups: int = 32,
+        add_attention: bool = True,
+        attention_head_dim: int = 1,
+    ) -> None:
+        super().__init__()
+        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+        self.add_attention = add_attention
+
+        # There is always at least one resnet
+        resnets = [
+            HunyuanVideoResnetBlockCausal3D(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                eps=resnet_eps,
+                groups=resnet_groups,
+                dropout=dropout,
+                non_linearity=resnet_act_fn,
+            )
+        ]
+        attentions = []
+
+        for _ in range(num_layers):
+            if self.add_attention:
+                attentions.append(
+                    HunyuanVAEAttention(
+                        in_channels,
+                        heads=in_channels // attention_head_dim,
+                        dim_head=attention_head_dim,
+                        eps=resnet_eps,
+                        norm_num_groups=resnet_groups,
+                        bias=True,
+                    )
+                )
+            else:
+                attentions.append(None)
+
+            resnets.append(
+                HunyuanVideoResnetBlockCausal3D(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    non_linearity=resnet_act_fn,
+                )
+            )
+
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            hidden_states = self._gradient_checkpointing_func(self.resnets[0], hidden_states)
+
+            for attn, resnet in zip(self.attentions, self.resnets[1:]):
+                if attn is not None:
+                    batch_size, num_channels, num_frames, height, width = hidden_states.shape
+                    hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3)
+                    attention_mask = prepare_causal_attention_mask(
+                        num_frames, height * width, hidden_states.dtype, hidden_states.device, batch_size=batch_size
+                    )
+                    hidden_states = attn(hidden_states, attention_mask=attention_mask)
+                    hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(0, 4, 1, 2, 3)
+
+                hidden_states = self._gradient_checkpointing_func(resnet, hidden_states)
+
+        else:
+            hidden_states = self.resnets[0](hidden_states)
+
+            for attn, resnet in zip(self.attentions, self.resnets[1:]):
+                if attn is not None:
+                    batch_size, num_channels, num_frames, height, width = hidden_states.shape
+                    hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3)
+                    attention_mask = prepare_causal_attention_mask(
+                        num_frames, height * width, hidden_states.dtype, hidden_states.device, batch_size=batch_size
+                    )
+                    hidden_states = attn(hidden_states, attention_mask=attention_mask)
+                    hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(0, 4, 1, 2, 3)
+
+                hidden_states = resnet(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideoDownBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_act_fn: str = "silu",
+        resnet_groups: int = 32,
+        add_downsample: bool = True,
+        downsample_stride: int = 2,
+        downsample_padding: int = 1,
+    ) -> None:
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                HunyuanVideoResnetBlockCausal3D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    non_linearity=resnet_act_fn,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    HunyuanVideoDownsampleCausal3D(
+                        out_channels,
+                        out_channels=out_channels,
+                        padding=downsample_padding,
+                        stride=downsample_stride,
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for resnet in self.resnets:
+                hidden_states = self._gradient_checkpointing_func(resnet, hidden_states)
+        else:
+            for resnet in self.resnets:
+                hidden_states = resnet(hidden_states)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideoUpBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_act_fn: str = "silu",
+        resnet_groups: int = 32,
+        add_upsample: bool = True,
+        upsample_scale_factor: Tuple[int, int, int] = (2, 2, 2),
+    ) -> None:
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            input_channels = in_channels if i == 0 else out_channels
+
+            resnets.append(
+                HunyuanVideoResnetBlockCausal3D(
+                    in_channels=input_channels,
+                    out_channels=out_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    non_linearity=resnet_act_fn,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList(
+                [
+                    HunyuanVideoUpsampleCausal3D(
+                        out_channels,
+                        out_channels=out_channels,
+                        upsample_factor=upsample_scale_factor,
+                    )
+                ]
+            )
+        else:
+            self.upsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for resnet in self.resnets:
+                hidden_states = self._gradient_checkpointing_func(resnet, hidden_states)
+
+        else:
+            for resnet in self.resnets:
+                hidden_states = resnet(hidden_states)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideoEncoder3D(nn.Module):
+    r"""
+    Causal encoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603).
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        down_block_types: Tuple[str, ...] = (
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+        ),
+        block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        double_z: bool = True,
+        mid_block_add_attention=True,
+        temporal_compression_ratio: int = 4,
+        spatial_compression_ratio: int = 8,
+    ) -> None:
+        super().__init__()
+
+        self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1)
+        self.mid_block = None
+        self.down_blocks = nn.ModuleList([])
+
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            if down_block_type != "HunyuanVideoDownBlock3D":
+                raise ValueError(f"Unsupported down_block_type: {down_block_type}")
+
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio))
+            num_time_downsample_layers = int(np.log2(temporal_compression_ratio))
+
+            if temporal_compression_ratio == 4:
+                add_spatial_downsample = bool(i < num_spatial_downsample_layers)
+                add_time_downsample = bool(
+                    i >= (len(block_out_channels) - 1 - num_time_downsample_layers) and not is_final_block
+                )
+            elif temporal_compression_ratio == 8:
+                add_spatial_downsample = bool(i < num_spatial_downsample_layers)
+                add_time_downsample = bool(i < num_time_downsample_layers)
+            else:
+                raise ValueError(f"Unsupported time_compression_ratio: {temporal_compression_ratio}")
+
+            downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1)
+            downsample_stride_T = (2,) if add_time_downsample else (1,)
+            downsample_stride = tuple(downsample_stride_T + downsample_stride_HW)
+
+            down_block = HunyuanVideoDownBlock3D(
+                num_layers=layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                add_downsample=bool(add_spatial_downsample or add_time_downsample),
+                resnet_eps=1e-6,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                downsample_stride=downsample_stride,
+                downsample_padding=0,
+            )
+
+            self.down_blocks.append(down_block)
+
+        self.mid_block = HunyuanVideoMidBlock3D(
+            in_channels=block_out_channels[-1],
+            resnet_eps=1e-6,
+            resnet_act_fn=act_fn,
+            attention_head_dim=block_out_channels[-1],
+            resnet_groups=norm_num_groups,
+            add_attention=mid_block_add_attention,
+        )
+
+        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
+        self.conv_act = nn.SiLU()
+
+        conv_out_channels = 2 * out_channels if double_z else out_channels
+        self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv_in(hidden_states)
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for down_block in self.down_blocks:
+                hidden_states = self._gradient_checkpointing_func(down_block, hidden_states)
+
+            hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states)
+        else:
+            for down_block in self.down_blocks:
+                hidden_states = down_block(hidden_states)
+
+            hidden_states = self.mid_block(hidden_states)
+
+        hidden_states = self.conv_norm_out(hidden_states)
+        hidden_states = self.conv_act(hidden_states)
+        hidden_states = self.conv_out(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideoDecoder3D(nn.Module):
+    r"""
+    Causal decoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603).
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        up_block_types: Tuple[str, ...] = (
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+        ),
+        block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        mid_block_add_attention=True,
+        time_compression_ratio: int = 4,
+        spatial_compression_ratio: int = 8,
+    ):
+        super().__init__()
+        self.layers_per_block = layers_per_block
+
+        self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1)
+        self.up_blocks = nn.ModuleList([])
+
+        # mid
+        self.mid_block = HunyuanVideoMidBlock3D(
+            in_channels=block_out_channels[-1],
+            resnet_eps=1e-6,
+            resnet_act_fn=act_fn,
+            attention_head_dim=block_out_channels[-1],
+            resnet_groups=norm_num_groups,
+            add_attention=mid_block_add_attention,
+        )
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            if up_block_type != "HunyuanVideoUpBlock3D":
+                raise ValueError(f"Unsupported up_block_type: {up_block_type}")
+
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio))
+            num_time_upsample_layers = int(np.log2(time_compression_ratio))
+
+            if time_compression_ratio == 4:
+                add_spatial_upsample = bool(i < num_spatial_upsample_layers)
+                add_time_upsample = bool(
+                    i >= len(block_out_channels) - 1 - num_time_upsample_layers and not is_final_block
+                )
+            else:
+                raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}")
+
+            upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1)
+            upsample_scale_factor_T = (2,) if add_time_upsample else (1,)
+            upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW)
+
+            up_block = HunyuanVideoUpBlock3D(
+                num_layers=self.layers_per_block + 1,
+                in_channels=prev_output_channel,
+                out_channels=output_channel,
+                add_upsample=bool(add_spatial_upsample or add_time_upsample),
+                upsample_scale_factor=upsample_scale_factor,
+                resnet_eps=1e-6,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+            )
+
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6)
+        self.conv_act = nn.SiLU()
+        self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[0], out_channels, kernel_size=3)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv_in(hidden_states)
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states)
+
+            for up_block in self.up_blocks:
+                hidden_states = self._gradient_checkpointing_func(up_block, hidden_states)
+        else:
+            hidden_states = self.mid_block(hidden_states)
+
+            for up_block in self.up_blocks:
+                hidden_states = up_block(hidden_states)
+
+        # post-process
+        hidden_states = self.conv_norm_out(hidden_states)
+        hidden_states = self.conv_act(hidden_states)
+        hidden_states = self.conv_out(hidden_states)
+
+        return hidden_states
+
+
+class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE):
+    r"""
+    A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos.
+    Introduced in [HunyuanVideo](https://huggingface.co/papers/2412.03603).
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+    """
+
+    _supports_gradient_checkpointing = True
+    @auto_attributes
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        latent_channels: int = 16,
+        down_block_types: Tuple[str, ...] = (
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+        ),
+        up_block_types: Tuple[str, ...] = (
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+        ),
+        block_out_channels: Tuple[int] = (128, 256, 512, 512),
+        layers_per_block: int = 2,
+        act_fn: str = "silu",
+        norm_num_groups: int = 32,
+        scaling_factor: float = 0.476986,
+        spatial_compression_ratio: int = 8,
+        temporal_compression_ratio: int = 4,
+        mid_block_add_attention: bool = True,
+        load_encoder: bool = True,
+        load_decoder: bool = True,
+    ) -> None:
+        super().__init__()
+
+        # TODO(will): only pass in config. We do this by manually defining a
+        # config for hunyuan vae
+        self.block_out_channels = block_out_channels
+
+        if load_encoder:
+            self.encoder = HunyuanVideoEncoder3D(
+                in_channels=in_channels,
+                out_channels=latent_channels,
+                down_block_types=down_block_types,
+                block_out_channels=block_out_channels,
+                layers_per_block=layers_per_block,
+            norm_num_groups=norm_num_groups,
+            act_fn=act_fn,
+            double_z=True,
+            mid_block_add_attention=mid_block_add_attention,
+            temporal_compression_ratio=temporal_compression_ratio,
+                spatial_compression_ratio=spatial_compression_ratio,
+            )
+            self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1)
+            
+        if load_decoder:
+            self.decoder = HunyuanVideoDecoder3D(
+                in_channels=latent_channels,
+                out_channels=out_channels,
+                up_block_types=up_block_types,
+                block_out_channels=block_out_channels,
+                layers_per_block=layers_per_block,
+                norm_num_groups=norm_num_groups,
+                act_fn=act_fn,
+                time_compression_ratio=temporal_compression_ratio,
+                spatial_compression_ratio=spatial_compression_ratio,
+                mid_block_add_attention=mid_block_add_attention,
+            )
+            self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1)
+
+
+
+        # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
+        # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
+        # intermediate tiles together, the memory requirement can be lowered.
+        self.use_tiling = True
+
+        # The minimal tile height and width for spatial tiling to be used
+        self.tile_sample_min_height = 256
+        self.tile_sample_min_width = 256
+        self.tile_sample_min_num_frames = 16
+
+        # The minimal distance between two spatial tiles
+        self.tile_sample_stride_height = 192
+        self.tile_sample_stride_width = 192
+        self.tile_sample_stride_num_frames = 12
+
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.encoder(x)
+        enc = self.quant_conv(x)
+        return enc
+    
+    def _decode(self, z: torch.Tensor) -> torch.Tensor:
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+        return dec
+    
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        generator: Optional[torch.Generator] = None,
+    ) -> torch.Tensor:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z)
+        return dec
+    

fastvideo/pipelines/__init__.py

@@ -3,12 +3,25 @@ Diffusion pipelines for FastVideo.
 
 This package contains diffusion pipelines for generating videos and images.
 """
+import os
+import json
+from copy import deepcopy
 
 from typing import Dict, Optional, Type, Any
 
-# First, import the registry
-from fastvideo.pipelines.pipeline_registry import PipelineRegistry, register_pipeline
+from fastvideo.pipelines.pipeline_registry import PipelineRegistry
 from fastvideo.inference_args import InferenceArgs
+from fastvideo.logger import init_logger
+from huggingface_hub import snapshot_download
+from transformers import PretrainedConfig
+from fastvideo.models.hf_transformer_utils import get_hf_config, get_diffusers_config
+from fastvideo.models import get_scheduler
+from fastvideo.loader.loader import get_model_loader
+import glob
+from fastvideo.loader.fsdp_load import load_fsdp_model
+from fastvideo.models.loader import PipelineComponentLoader
+
+logger = init_logger(__name__)
 
 # Then import the base classes
 from fastvideo.pipelines.composed.composed_pipeline_base import (
@@ -20,30 +33,199 @@ def get_pipeline_type(inference_args: InferenceArgs) -> str:
     # hardcode for now
     return "hunyuan_video"
 
+def maybe_download_model(model_path: str) -> str:
+    """
+    Check if the model path is a Hugging Face Hub model ID and download it if needed.
+    
+    Args:
+        model_path: Local path or Hugging Face Hub model ID
+        
+    Returns:
+        Local path to the model
+    """
+    
+    # If the path exists locally, return it
+    if os.path.exists(model_path):
+        logger.info(f"Model already exists locally at {model_path}")
+        return model_path
+    
+    # Otherwise, assume it's a HF Hub model ID and try to download it
+    try:
+        logger.info(f"Downloading model snapshot from HF Hub for {model_path}...")
+        local_path = snapshot_download(
+            repo_id=model_path,
+            # allow_patterns=["*.json", "*.bin", "*.safetensors", "*.pt", "*.pth", "*.ckpt"],
+            ignore_patterns=["*.onnx", "*.msgpack"],
+        )
+        logger.info(f"Downloaded model to {local_path}")
+        return local_path
+    except Exception as e:
+        raise ValueError(f"Could not find model at {model_path} and failed to download from HF Hub: {e}")
+
+def verify_model_config_and_directory(model_path: str) -> dict:
+    """
+    Verify that the model directory contains a valid diffusers configuration.
+    
+    Args:
+        model_path: Path to the model directory
+        
+    Returns:
+        The loaded model configuration as a dictionary
+    """
+    
+    # Check for model_index.json which is required for diffusers models
+    config_path = os.path.join(model_path, "model_index.json")
+    if not os.path.exists(config_path):
+        raise ValueError(
+            f"Model directory {model_path} does not contain model_index.json. "
+            "Only Hugging Face diffusers format is supported."
+        )
+    
+    # Check for transformer and vae directories
+    transformer_dir = os.path.join(model_path, "transformer")
+    vae_dir = os.path.join(model_path, "vae")
+    
+    if not os.path.exists(transformer_dir):
+        raise ValueError(f"Model directory {model_path} does not contain a transformer/ directory.")
+    
+    if not os.path.exists(vae_dir):
+        raise ValueError(f"Model directory {model_path} does not contain a vae/ directory.")
+    
+    # Load the config
+    try:
+        with open(config_path, "r") as f:
+            config = json.load(f)
+    except Exception as e:
+        raise ValueError(f"Failed to load model configuration from {config_path}: {e}")
+
+    # Verify diffusers version exists
+    if "_diffusers_version" not in config:
+        raise ValueError(f"model_index.json does not contain _diffusers_version")
+    
+    logger.info(f"Diffusers version: {config['_diffusers_version']}")
+    return config
+
+def load_pipeline_module(module_name: str, component_model_path: str, transformers_or_diffusers: str, architecture: str, inference_args: InferenceArgs) -> Any:
+    """
+    Load a pipeline module using the appropriate loader.
+    
+    Args:
+        module_name: Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler")
+        component_model_path: Path to the component model
+        transformers_or_diffusers: Whether the module is from transformers or diffusers
+        architecture: Architecture of the component model
+        inference_args: Inference arguments
+        
+    Returns:
+        The loaded module
+    """
+    return PipelineComponentLoader.load_module(
+        module_name=module_name,
+        component_model_path=component_model_path,
+        transformers_or_diffusers=transformers_or_diffusers,
+        architecture=architecture,
+        inference_args=inference_args
+    )
+
+def load_pipeline_modules(model_path: str, config: Dict, inference_args: InferenceArgs) -> dict[str, Any]:
+    """
+    Load the pipeline modules from the config.
+    
+    Args:
+        config: The model_index.json config
+        inference_args: Inference arguments
+        
+    Returns:
+        Dictionary mapping module names to loaded modules
+    """
+    logger.info(f"Loading pipeline modules from config: {config}")
+    modules_config = deepcopy(config)
+    
+    # remove keys that are not pipeline modules
+    modules_config.pop("_class_name")
+    modules_config.pop("_diffusers_version")
+    
+    # some sanity checks
+    assert len(modules_config) > 1, "model_index.json must contain at least one pipeline module"
+    
+    required_modules = ["vae", "text_encoder", "transformer", "scheduler", "tokenizer"]
+    for module_name in required_modules:
+        if module_name not in modules_config:
+            raise ValueError(f"model_index.json must contain a {module_name} module")
+    logger.info(f"Diffusers config passed sanity checks")
+    
+
+    # all the component models used by the pipeline
+    pipeline_modules = {}
+    for module_name, (transformers_or_diffusers, architecture) in modules_config.items():
+        component_model_path = os.path.join(model_path, module_name)
+        module = load_pipeline_module(
+            module_name,
+            component_model_path,
+            transformers_or_diffusers,
+            architecture,
+            inference_args,
+        )
+
+        pipeline_modules[module_name] = module
+
+
+    # Check if all required modules were loaded
+    for module_name in required_modules:
+        if module_name not in pipeline_modules or pipeline_modules[module_name] is None:
+            logger.warning(f"Required module {module_name} was not loaded properly")
+    
+    return pipeline_modules
+
 def build_pipeline(inference_args: InferenceArgs) -> ComposedPipelineBase:
     """
     Only works with valid hf diffusers configs. (model_index.json)
-    We want to build a pipeline based on the inference args.
-    1. load the correct hf config from disk or download from hub
-    2. based on the config, determine the pipeline class and pipeline loader class
-    3. parse the config to get the model component (vae, text_encoders, etc...)
-    names and paths
-    4. the pipeline loader class will use the model component names and paths to load
-    the models. 
-    5. the pipeline class will be composed of the models returned by the pipeline loader
+    We want to build a pipeline based on the inference args mode_path:
+    1. download the model from the hub if it's not already downloaded
+    2. verify the model config and directory
+    3. based on the config, determine the pipeline class 
+    4. parse the config to get the model components (vae, text_encoders, etc...)
+    5. the pipeline loader class will use the model component names and paths to load
+    6. the pipeline class will be composed of the models returned by the pipeline loader
     """
     # Get pipeline type
-    pipeline_type = get_pipeline_type(inference_args)
+    model_path = inference_args.model_path
+    model_path = maybe_download_model(model_path)
+    # inference_args.downloaded_model_path = model_path
+    logger.info(f"Model path: {model_path}")
+    config = verify_model_config_and_directory(model_path)
+
+    pipeline_architecture = config.get("_class_name")
+    if pipeline_architecture is None:
+        raise ValueError("Model config does not contain a _class_name attribute. "
+                         "Only diffusers format is supported.")
     
-    if not hasattr(inference_args, 'model_type') or not inference_args.model_type:
-        inference_args.model_type = pipeline_type
+    pipeline_cls, pipeline_architecture = PipelineRegistry.resolve_pipeline_cls(pipeline_architecture)
+
+    # instantiate the pipeline
+    pipeline = pipeline_cls()
     
-    # TODO(will): Get pipeline loader (hardcoded for now)
-    from fastvideo.pipelines.loader import get_pipeline_loader
-    pipeline_loader = get_pipeline_loader(inference_args)
+    try:
+        pipeline_modules = load_pipeline_modules(model_path, config, inference_args)
+    except Exception as e:
+        logger.error(f"Failed to load pipeline modules: {e}")
+        raise ValueError(f"Failed to load pipeline modules: {e}")
+
+    logger.info(f"Initializing encoders")
+    pipeline.initialize_encoders(pipeline_modules, inference_args)
+
+    logger.info(f"Registering modules")
+    pipeline.register_modules(pipeline_modules)
     
-    # Load and return pipeline
-    return pipeline_loader.load_pipeline_v2(inference_args)
+    logger.info(f"Setting up pipeline")
+    pipeline.setup_pipeline(inference_args)
+
+    logger.info(f"Initializing pipeline")
+    pipeline.initialize_pipeline(inference_args)
+    
+    # pipeline is now initialized and ready to use
+    return pipeline
+
 
 
 def list_available_pipelines() -> Dict[str, Type[Any]]:
@@ -55,16 +237,9 @@ def list_available_pipelines() -> Dict[str, Type[Any]]:
     """
     return PipelineRegistry.list()
 
-
-# Import all pipeline implementations to register them
-# These imports should be at the end to avoid circular imports
-from fastvideo.pipelines.implementations.hunyuan import HunyuanVideoPipeline
-
 __all__ = [
-    "create_pipeline",
+    "build_pipeline",
     "list_available_pipelines",
     "ComposedPipelineBase",
     "DiffusionPipelineOutput",
-    "register_pipeline",
-    "HunyuanVideoPipeline",
 ] 

fastvideo/pipelines/composed/composed_pipeline_base.py

@@ -63,6 +63,14 @@ class ComposedPipelineBase(ABC):
         """
         ...
     
+    @abstractmethod
+    def initialize_encoders(self, modules: Dict[str, Any], inference_args: InferenceArgs):
+        """
+        Initialize the encoders. Will remove the encoders/tokenizers modules from the
+        modules. Will add the TextEncoder or ImageEncoder to the modules.
+        """
+        ...
+    
     def register_modules(self, modules: Dict[str, Any]):
         """
         Register modules with the pipeline and its stages.

fastvideo/pipelines/composed/text_to_video.py

@@ -1,113 +0,0 @@
-"""
-Text-to-video pipeline implementation.
-
-This module contains a standard text-to-video pipeline implementation that can be
-used as a base for more specific implementations.
-"""
-
-from typing import Any, Dict, List, Optional, Union, Tuple
-import torch
-
-from fastvideo.pipelines.composed.composed_pipeline_base import ComposedPipelineBase
-from fastvideo.pipelines.stages import (
-    InputValidationStage,
-    PromptEncodingStage,
-    TimestepPreparationStage,
-    LatentPreparationStage,
-    ConditioningStage,
-    DenoisingStage,
-    DecodingStage,
-    PostProcessingStage,
-)
-# Import from the main pipelines package to avoid circular imports
-from fastvideo.pipelines import register_pipeline
-
-
-@register_pipeline("text-to-video")
-class TextToVideoPipeline(ComposedPipelineBase):
-    """
-    Standard text-to-video pipeline.
-    
-    This pipeline implements a standard text-to-video diffusion process with the
-    following stages:
-    1. Input validation
-    2. Prompt encoding
-    3. Timestep preparation
-    4. Latent preparation
-    5. Conditioning (classifier-free guidance)
-    6. Denoising
-    7. Decoding
-    8. Post-processing
-    
-    This pipeline supports models with either a single text encoder or dual text encoders.
-    Different model architectures may use secondary encoders for various purposes such as
-    additional conditioning information, temporal control, stylistic aspects, or other
-    model-specific features.
-    """
-    
-    is_video_pipeline = True
-    
-    def __init__(
-        self,
-        input_validation_stage: InputValidationStage,
-        prompt_encoding_stage: PromptEncodingStage,
-        timestep_preparation_stage: TimestepPreparationStage,
-        latent_preparation_stage: LatentPreparationStage,
-        conditioning_stage: ConditioningStage,
-        denoising_stage: DenoisingStage,
-        decoding_stage: DecodingStage,
-        post_processing_stage: Optional[PostProcessingStage] = None,
-        secondary_prompt_encoding_stage: Optional[PromptEncodingStage] = None,
-    ):
-        """
-        Initialize the text-to-video pipeline.
-        
-        Args:
-            input_validation_stage: The input validation stage.
-            prompt_encoding_stage: The primary prompt encoding stage.
-            timestep_preparation_stage: The timestep preparation stage.
-            latent_preparation_stage: The latent preparation stage.
-            conditioning_stage: The conditioning stage.
-            denoising_stage: The denoising stage.
-            decoding_stage: The decoding stage.
-            post_processing_stage: The post-processing stage (optional).
-            secondary_prompt_encoding_stage: Optional secondary prompt encoding stage for models
-                with dual text encoders. The purpose of this encoder varies by model architecture
-                and may handle additional conditioning aspects such as temporal control, style,
-                or other model-specific features.
-        """
-        super().__init__()
-        
-        # Add the stages in order
-        self.add_stage(input_validation_stage)
-        self.add_stage(prompt_encoding_stage)
-        
-        # Add optional secondary prompt encoding stage if provided
-        if secondary_prompt_encoding_stage is not None:
-            self.add_stage(secondary_prompt_encoding_stage)
-            
-        self.add_stage(timestep_preparation_stage)
-        self.add_stage(latent_preparation_stage)
-        self.add_stage(conditioning_stage)
-        self.add_stage(denoising_stage)
-        self.add_stage(decoding_stage)
-        
-        if post_processing_stage is not None:
-            self.add_stage(post_processing_stage)
-    
-    @property
-    def components(self) -> Dict[str, Any]:
-        """Get all components used by this pipeline."""
-        components = {}
-        
-        # Add components from each stage
-        for stage in self._stages:
-            for attr_name in dir(stage):
-                if attr_name.startswith("__") or attr_name.startswith("_"):
-                    continue
-                
-                attr = getattr(stage, attr_name)
-                if isinstance(attr, torch.nn.Module):
-                    components[attr_name] = attr
-        
-        return components 

fastvideo/pipelines/hunyuan/hunyuan_loader.py

@@ -221,7 +221,6 @@ class HunyuanPipelineLoader(PipelineLoader):
             hidden_state_skip_layer=inference_args.hidden_state_skip_layer,
             apply_final_norm=inference_args.apply_final_norm,
             reproduce=inference_args.reproduce,
-            logger=logger,
             device=self.device if not inference_args.use_cpu_offload else "cpu",
         )
         text_encoder_2 = None
@@ -232,7 +231,6 @@ class HunyuanPipelineLoader(PipelineLoader):
                 text_encoder_precision=inference_args.text_encoder_precision_2,
                 tokenizer_type=inference_args.tokenizer_2,
                 reproduce=inference_args.reproduce,
-                logger=logger,
                 device=self.device if not inference_args.use_cpu_offload else "cpu",
             )
             

fastvideo/pipelines/implementations/__init__.py

@@ -1,8 +0,0 @@
-"""
-Concrete implementations of diffusion pipelines.
-
-This package contains concrete implementations of diffusion pipelines for
-specific models or architectures.
-"""
-
-__all__ = [] 

fastvideo/pipelines/implementations/hunyuan/__init__.py

@@ -5,9 +5,9 @@ This package contains implementations of diffusion pipelines for HunYuan models.
 """
 
 from fastvideo.pipelines.implementations.hunyuan.hunyuan_pipeline import (
-    HunyuanVideoPipeline,
+    EntryClass,
 )
 
 __all__ = [
-    "HunyuanVideoPipeline",
+    "EntryClass",
 ] 

fastvideo/pipelines/implementations/hunyuan/hunyuan_loader.py

@@ -8,14 +8,14 @@ from fastvideo.inference_args import InferenceArgs
 from fastvideo.pipelines.implementations.hunyuan.constants import PRECISION_TO_TYPE, PROMPT_TEMPLATE
 from fastvideo.logger import init_logger
 from typing import Tuple, Optional
-
+from fastvideo.models.dits import HunyuanVideoTransformer3DModel
 import os
 import glob
 import json
 from fastvideo.loader.fsdp_load import load_fsdp_model
-
 from fastvideo.models.hunyuan.text_encoder import TextEncoder
-
+from fastvideo.models.vaes import AutoencoderKLHunyuanVideo
+from safetensors.torch import load_file as safetensors_load_file
 from fastvideo.platforms import current_platform
 
 logger = init_logger(__name__)
@@ -113,7 +113,7 @@ class HunyuanPipelineLoader(PipelineLoader):
         """Custom transformer loading for Hunyuan"""
         
         # Path to model files
-        # TODO(PY): remove this hardcode
+        # TODO(PY/Will): remove this hardcode
         path = "data/hunyuanvideo_community/transformer"
         path = Path(path)
         
@@ -130,7 +130,7 @@ class HunyuanPipelineLoader(PipelineLoader):
             config.pop("_class_name")
         if "_diffusers_version" in config:
             config.pop("_diffusers_version")
-        
+            
         # Find all safetensors files
         safetensors_list = glob.glob(os.path.join(str(path), "*.safetensors"))
         if not safetensors_list:
@@ -139,8 +139,9 @@ class HunyuanPipelineLoader(PipelineLoader):
         logger.info(f"Loading model from {len(safetensors_list)} safetensors files in {path}")
         
         # Load the model using FSDP loader
+        # TOOD(Will) new registry for DiT
         model = load_fsdp_model(
-            model_name="HunyuanVideoTransformer3DModel",
+            model_cls=HunyuanVideoTransformer3DModel,
             init_params=config,
             weight_dir_list=safetensors_list,
             device=self.device,
@@ -181,6 +182,49 @@ class HunyuanPipelineLoader(PipelineLoader):
         return model
     
     def load_vae(self, inference_args: InferenceArgs):
+        """Custom VAE loading for Hunyuan"""
+        # TODO(will): replace this with abstracted model
+        use_v1_loader = True
+        if use_v1_loader:
+            return self.load_vae_v1(inference_args)
+        else:
+            return self.load_vae_v0(inference_args)
+        
+    def load_vae_v1(self, inference_args: InferenceArgs):
+        """Custom VAE loading for Hunyuan"""
+        path = "data/hunyuanvideo_community/vae"
+        path = Path(path)
+        # TODO(PY/Will): remove this hardcode
+        # Load config file
+        config_path = path / "config.json"
+        if not config_path.exists():
+            raise ValueError(f"Config file not found at {config_path}")
+
+        with open(config_path, "r") as f:
+            config = json.load(f)
+        
+        # Clean up config
+
+        class_name = config.pop("_class_name")
+        if "_diffusers_version" in config:
+            config.pop("_diffusers_version")
+        vae = AutoencoderKLHunyuanVideo(**config).to(self.device)
+        
+        # Find all safetensors files
+        safetensors_list = glob.glob(os.path.join(str(path), "*.safetensors"))
+        # TODO(PY)
+        assert len(safetensors_list) == 1, f"Found {len(safetensors_list)} safetensors files in {path}"
+        loaded = safetensors_load_file(safetensors_list[0])
+        vae.load_state_dict(loaded)
+        dtype = PRECISION_TO_TYPE[inference_args.vae_precision]
+        vae = vae.eval().to(dtype)
+        
+        return vae, {"s_ratio": config["spatial_compression_ratio"], "t_ratio": config["temporal_compression_ratio"]}
+        
+        
+        
+        
+    def load_vae_v0(self, inference_args: InferenceArgs):
         """Custom VAE loading for Hunyuan"""
         # TODO(will): replace this with abstracted model
         from fastvideo.models.hunyuan.vae import load_vae
@@ -188,7 +232,7 @@ class HunyuanPipelineLoader(PipelineLoader):
             inference_args.vae,
             inference_args.vae_precision,
             logger=logger,
-            device=self.device if not inference_args.use_cpu_offload else "cpu",
+            device=self.device
         )
         vae_kwargs = {"s_ratio": s_ratio, "t_ratio": t_ratio}
         return vae, vae_kwargs

fastvideo/pipelines/implementations/hunyuan/hunyuan_pipeline.py

@@ -5,7 +5,7 @@ This module contains an implementation of the HunYuan video diffusion pipeline
 using the modular pipeline architecture.
 """
 
-from typing import Any, Dict, List, Optional, Union, Tuple
+from typing import Union, Any, Dict
 import torch
 from diffusers.image_processor import VaeImageProcessor
 
@@ -20,14 +20,16 @@ from fastvideo.pipelines.stages import (
     DecodingStage,
     PostProcessingStage,
 )
-from fastvideo.pipelines import register_pipeline
 from fastvideo.pipelines.stages.prompt_encoding import PromptEncodingStage
 # from fastvideo.pipelines.stages.timestep_preparation import FlowMatchingTimestepPreparationStage
 from fastvideo.inference_args import InferenceArgs
 # from fastvideo.pipelines.composed.composed_pipeline_base import DiffusionPipelineOutput
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.models.text_encoder import TextEncoder
+# TODO(will): move PRECISION_TO_TYPE to better place
+from .constants import PROMPT_TEMPLATE, PRECISION_TO_TYPE
+
 from diffusers.utils import BaseOutput
-from typing import Union, Optional, Tuple, List
 import numpy as np
 from dataclasses import dataclass
 
@@ -49,11 +51,75 @@ class DiffusionPipelineOutput(BaseOutput):
     videos: Union[torch.Tensor, np.ndarray]
 
 
-@register_pipeline("hunyuan-video")
 class HunyuanVideoPipeline(ComposedPipelineBase):
 
+    def initialize_encoders(self, modules: Dict[str, Any], inference_args: InferenceArgs):
+        """
+        Initialize the encoders. Will remove the encoders/tokenizers modules from the
+        modules. Will add the TextEncoder or ImageEncoder to the modules.
+        """
+        if inference_args.prompt_template_video is not None:
+            crop_start = PROMPT_TEMPLATE[inference_args.prompt_template_video].get("crop_start", 0)
+        elif inference_args.prompt_template is not None:
+            crop_start = PROMPT_TEMPLATE[inference_args.prompt_template].get("crop_start", 0)
+        else:
+            crop_start = 0
+
+        max_length = inference_args.text_len + crop_start
+
+        # prompt_template
+        prompt_template = (PROMPT_TEMPLATE[inference_args.prompt_template] if inference_args.prompt_template is not None else None)
+
+        # prompt_template_video
+        prompt_template_video = (PROMPT_TEMPLATE[inference_args.prompt_template_video]
+                                 if inference_args.prompt_template_video is not None else None)
+
+        encoder_1 = modules.pop("text_encoder")
+        assert encoder_1 is not None, "Text encoder is not found"
+        encoder_1.to(inference_args.device)
+        encoder_1.to(dtype=PRECISION_TO_TYPE[inference_args.text_encoder_precision])
+        encoder_1.requires_grad_(False)
+
+        tokenizer_1 = modules.pop("tokenizer")
+        assert tokenizer_1 is not None, "Tokenizer is not found"
+
+        text_encoder = TextEncoder(
+            text_encoder=encoder_1,
+            tokenizer=tokenizer_1,
+            # text_encoder_type="text_encoder"
+            max_length=max_length,
+            # text_encoder_precision=inference_args.text_encoder_precision,
+            prompt_template=prompt_template,
+            prompt_template_video=prompt_template_video,
+            hidden_state_skip_layer=inference_args.hidden_state_skip_layer,
+            apply_final_norm=inference_args.apply_final_norm,
+            reproduce=inference_args.reproduce,
+            device=inference_args.device if not inference_args.use_cpu_offload else "cpu",
+        )
+
+        encoder_2 = modules.pop("text_encoder_2")
+        assert encoder_2 is not None, "Text encoder 2 is not found"
+        encoder_2.to(inference_args.device)
+        encoder_2.to(dtype=PRECISION_TO_TYPE[inference_args.text_encoder_precision])
+        encoder_2.requires_grad_(False)
+
+        tokenizer_2 = modules.pop("tokenizer_2")
+        assert tokenizer_2 is not None, "Tokenizer 2 is not found"
+
+        text_encoder_2 = TextEncoder(
+            text_encoder=encoder_2,
+            tokenizer=tokenizer_2,
+            # text_encoder_type="text_encoder_2",
+            max_length=inference_args.text_len_2,
+            # text_encoder_precision=inference_args.text_encoder_precision,
+            reproduce=inference_args.reproduce,
+            device=inference_args.device if not inference_args.use_cpu_offload else "cpu",
+        )
+        modules["text_encoder"] = text_encoder
+        modules["text_encoder_2"] = text_encoder_2
+
+
     def setup_pipeline(self, inference_args: InferenceArgs):
-        self._stages = []
         self.add_stage("input_validation_stage", 
                         InputValidationStage())
         self.add_stage("prompt_encoding_stage_primary", 
@@ -76,7 +142,7 @@ class HunyuanVideoPipeline(ComposedPipelineBase):
         assert len(self._modules) > 0, "Pipeline modules are not set"
 
 
-        vae_scale_factor = 2**(len(self.vae.config.block_out_channels) - 1)
+        vae_scale_factor = 2**(len(self.vae.block_out_channels) - 1)
         inference_args.vae_scale_factor = vae_scale_factor
 
         self.image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
@@ -123,3 +189,5 @@ class HunyuanVideoPipeline(ComposedPipelineBase):
         batch = self.decoding_stage(batch, inference_args)
 
         return DiffusionPipelineOutput(videos=batch.videos) 
+
+EntryClass = HunyuanVideoPipeline

fastvideo/pipelines/loader.py

@@ -246,7 +246,15 @@ class PipelineLoader(ABC):
                 solver=inference_args.flow_solver,
             )
         return None
-    
+
+    @staticmethod
+    def load_pipeline_from_config(config: dict):
+        """Load a pipeline from a config."""
+        pipeline_cls = resolve_pipeline_cls(inference_args)
+        assert issubclass(pipeline_cls, ComposedPipelineBase)
+        logger.info(f"Pipeline class: {pipeline_cls}")
+        
+        
 
 def get_pipeline_loader(inference_args: InferenceArgs) -> PipelineLoader:
     """Get a pipeline loader based on the inference args."""

fastvideo/pipelines/pipeline_registry.py

@@ -1,69 +1,91 @@
-"""
-Registry for pipeline implementations.
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/models/registry.py
+# and https://github.com/sgl-project/sglang/blob/v0.4.3/python/sglang/srt/models/registry.py
 
-This module provides a registry for pipeline implementations, allowing them to be
-looked up by name.
-"""
+import importlib
+import pkgutil
+from dataclasses import dataclass, field
+from functools import lru_cache
+from typing import AbstractSet, Dict, List, Optional, Tuple, Type, Union
 
-from typing import Dict, Type, Any, Optional, Callable, TypeVar, cast
+from fastvideo.logger import init_logger
 
-# Define a type variable for the pipeline class
-T = TypeVar('T')
+import torch.nn as nn
 
-class PipelineRegistry:
-    """Registry for pipeline implementations."""
-    
-    _registry: Dict[str, Type[Any]] = {}
-    
-    @classmethod
-    def register(cls, name: str) -> Callable[[Type[T]], Type[T]]:
-        """
-        Register a pipeline implementation.
-        
-        Args:
-            name: The name to register the pipeline under.
-            
-        Returns:
-            A decorator that registers the pipeline.
-        """
-        def decorator(pipeline_cls: Type[T]) -> Type[T]:
-            cls._registry[name] = pipeline_cls
-            return pipeline_cls
-        
-        return decorator
-    
-    @classmethod
-    def get(cls, name: str) -> Optional[Type[Any]]:
-        """
-        Get a pipeline implementation by name.
-        
-        Args:
-            name: The name of the pipeline.
-            
-        Returns:
-            The pipeline class, or None if not found.
-        """
-        return cls._registry.get(name)
-    
-    @classmethod
-    def list(cls) -> Dict[str, Type[Any]]:
-        """
-        List all registered pipelines.
-        
-        Returns:
-            A dictionary of pipeline names to pipeline classes.
-        """
-        return cls._registry.copy()
+logger = init_logger(__name__)
 
 
-def register_pipeline(name: str) -> Callable[[Type[T]], Type[T]]:
-    """
-    Register a pipeline implementation.
-    
-    Args:
-        name: The name to register the pipeline under.
-        
-    Returns:
-        A decorator that registers the pipeline.
-    """
-    return PipelineRegistry.register(name) 
+@dataclass
+class _PipelineRegistry:
+    # Keyed by pipeline_arch
+    pipelines: Dict[str, Union[Type[nn.Module], str]] = field(default_factory=dict)
+
+    def get_supported_archs(self) -> AbstractSet[str]:
+        return self.pipelines.keys()
+
+    def _raise_for_unsupported(self, architectures: List[str]):
+        all_supported_archs = self.get_supported_archs()
+
+        if any(arch in all_supported_archs for arch in architectures):
+            raise ValueError(
+                f"Pipeline architectures {architectures} failed "
+                "to be inspected. Please check the logs for more details."
+            )
+
+        raise ValueError(
+            f"Pipeline architectures {architectures} are not supported for now. "
+            f"Supported architectures: {all_supported_archs}"
+        )
+
+    def _try_load_pipeline_cls(self, pipeline_arch: str) -> Optional[Type[nn.Module]]:
+        if pipeline_arch not in self.pipelines:
+            return None
+
+        return self.pipelines[pipeline_arch]
+
+    def resolve_pipeline_cls(
+        self,
+        architecture: str,
+    ) -> Tuple[Type[nn.Module], str]:
+        if not architecture:
+            logger.warning("No pipeline architecture is specified")
+
+        pipeline_cls = self._try_load_pipeline_cls(architecture)
+        if pipeline_cls is not None:
+            return (pipeline_cls, architecture)
+
+        return self._raise_for_unsupported(architecture)
+
+
+@lru_cache()
+def import_pipeline_classes():
+    pipeline_arch_name_to_cls = {}
+    package_name = "fastvideo.pipelines.implementations"
+    package = importlib.import_module(package_name)
+    for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."):
+        if ispkg:
+            try:
+                module = importlib.import_module(name)
+            except Exception as e:
+                logger.warning(f"Ignore import error when loading {name}. " f"{e}")
+                continue
+            if hasattr(module, "EntryClass"):
+                entry = module.EntryClass
+                print(entry)
+                print(entry.__name__)
+                if isinstance(
+                    entry, list
+                ):  # To support multiple pipeline classes in one module
+                    for tmp in entry:
+                        assert (
+                            tmp.__name__ not in pipeline_arch_name_to_cls
+                        ), f"Duplicated pipeline implementation for {tmp.__name__}"
+                        pipeline_arch_name_to_cls[tmp.__name__] = tmp
+                else:
+                    assert (
+                        entry.__name__ not in pipeline_arch_name_to_cls
+                    ), f"Duplicated pipeline implementation for {entry.__name__}"
+                    pipeline_arch_name_to_cls[entry.__name__] = entry
+    return pipeline_arch_name_to_cls
+
+
+PipelineRegistry = _PipelineRegistry(import_pipeline_classes())

fastvideo/pipelines/stages/decoding.py

@@ -3,17 +3,11 @@ Decoding stage for diffusion pipelines.
 """
 
 import torch
-from typing import Optional, Union, List, Tuple
-import numpy as np
-
 from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.inference_args import InferenceArgs
 from fastvideo.models.hunyuan.constants import PRECISION_TO_TYPE
-from fastvideo.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D
 from fastvideo.logger import init_logger
-from diffusers.utils import BaseOutput
-
 logger = init_logger(__name__)
 
 
@@ -50,18 +44,6 @@ class DecodingStage(PipelineStage):
             vae_dtype = PRECISION_TO_TYPE[inference_args.vae_precision]
             vae_autocast_enabled = (vae_dtype != torch.float32) and not inference_args.disable_autocast
             
-            # Handle different latent shapes
-            expand_temporal_dim = False
-            if len(latents.shape) == 4:
-                if isinstance(self.vae, AutoencoderKLCausal3D):
-                    latents = latents.unsqueeze(2)
-                    expand_temporal_dim = True
-            elif len(latents.shape) == 5:
-                pass
-            else:
-                raise ValueError(
-                    f"Only support latents with shape (b, c, h, w) or (b, c, f, h, w), but got {latents.shape}."
-                )
 
             # Apply scaling/shifting if needed
             if (hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor):
@@ -73,11 +55,9 @@ class DecodingStage(PipelineStage):
             with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
                 if inference_args.vae_tiling:
                     self.vae.enable_tiling()
-                image = self.vae.decode(latents, return_dict=False, generator=batch.generator)[0]
+                image = self.vae.decode(latents)
+
 
-            # Handle temporal dimension if needed
-            if expand_temporal_dim or image.shape[2] == 1:
-                image = image.squeeze(2)
 
         # Normalize image to [0, 1] range
         image = (image / 2 + 0.5).clamp(0, 1)

fastvideo/pipelines/stages/prompt_encoding.py

@@ -10,7 +10,7 @@ import torch
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.inference_args import InferenceArgs
 from fastvideo.pipelines.stages import PipelineStage
-from fastvideo.models.hunyuan.text_encoder import TextEncoder
+from fastvideo.models.text_encoder import TextEncoder
 from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)

fastvideo/sample/v1_fastvideo_inference.py

@@ -14,38 +14,31 @@ from fastvideo.inference_args import InferenceArgs
 from fastvideo.utils.utils import FlexibleArgumentParser
 from fastvideo.distributed import init_distributed_environment, initialize_model_parallel
 from fastvideo.logger import init_logger
-import torch.distributed as dist
 from fastvideo.utils.parallel_states import initialize_sequence_parallel_state
 logger = init_logger(__name__)
 
-def initialize_distributed():
-    local_rank = int(os.getenv("RANK", 0))
-    world_size = int(os.getenv("WORLD_SIZE", 1))
-    print("world_size", world_size)
-    torch.cuda.set_device(local_rank)
-    dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
-    initialize_sequence_parallel_state(world_size)
-
 def main(inference_args: InferenceArgs):
     # initialize_distributed()
     # print(nccl_info.sp_size)
-    # local_rank = int(os.environ.get("LOCAL_RANK", 0))
-    # rank = int(os.environ.get("RANK", 0))
-    # world_size = int(os.environ.get("WORLD_SIZE", 1))
-
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    rank = int(os.environ.get("RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    torch.cuda.set_device(local_rank)
     # logger.info(f"Initializing process: rank={rank}, local_rank={local_rank}, world_size={world_size}")
-    # init_distributed_environment(
-    #     world_size=world_size,
-    #     rank=rank,
-    #     local_rank=local_rank
-    # )
-    # print(inference_args.sp_size)
+    init_distributed_environment(
+        world_size=world_size,
+        rank=rank,
+        local_rank=local_rank
+    )
+    print(inference_args.sp_size)
 
     # Initialize tensor model parallel groups
-    # initialize_model_parallel(
-    #     sequence_model_parallel_size=inference_args.sp_size
-    # )
-    initialize_distributed()
+    initialize_model_parallel(
+        sequence_model_parallel_size=inference_args.sp_size,
+        tensor_model_parallel_size=inference_args.sp_size
+    )
+    # initialize_sequence_parallel_state(world_size)
+    # initialize_distributed()
 
 
     print('Creating engine')

fastvideo/sample/v2_fastvideo_inference.py

@@ -1,102 +0,0 @@
-import argparse
-import os
-from pathlib import Path
-
-import imageio
-import numpy as np
-import torch
-import torchvision
-from einops import rearrange
-
-# Fix the import path
-from fastvideo.inference_engine import InferenceEngine
-from fastvideo.inference_args import InferenceArgs
-from fastvideo.utils.utils import FlexibleArgumentParser
-from fastvideo.distributed import init_distributed_environment, initialize_model_parallel
-from fastvideo.logger import init_logger
-import torch.distributed as dist
-from fastvideo.utils.parallel_states import initialize_sequence_parallel_state
-logger = init_logger(__name__)
-
-def initialize_distributed():
-    local_rank = int(os.getenv("RANK", 0))
-    world_size = int(os.getenv("WORLD_SIZE", 1))
-    print("world_size", world_size)
-    torch.cuda.set_device(local_rank)
-    dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
-    initialize_sequence_parallel_state(world_size)
-
-def main(inference_args: InferenceArgs):
-    # initialize_distributed()
-    # print(nccl_info.sp_size)
-    local_rank = int(os.environ.get("LOCAL_RANK", 0))
-    rank = int(os.environ.get("RANK", 0))
-    world_size = int(os.environ.get("WORLD_SIZE", 1))
-    torch.cuda.set_device(local_rank)
-    # logger.info(f"Initializing process: rank={rank}, local_rank={local_rank}, world_size={world_size}")
-    init_distributed_environment(
-        world_size=world_size,
-        rank=rank,
-        local_rank=local_rank
-    )
-    print(inference_args.sp_size)
-
-    # Initialize tensor model parallel groups
-    initialize_model_parallel(
-        sequence_model_parallel_size=inference_args.sp_size
-    )
-    initialize_sequence_parallel_state(world_size)
-    # initialize_distributed()
-
-
-    print('Creating engine')
-    # Create inference object using the updated API
-    engine = InferenceEngine.create_engine(
-        inference_args,
-    )
-    print('Engine created')
-    # return
-
-    # Load prompts
-    if inference_args.prompt.endswith('.txt'):
-        with open(inference_args.prompt) as f:
-            prompts = [line.strip() for line in f.readlines()]
-    else:
-        prompts = [inference_args.prompt]
-
-    # Process each prompt
-    for prompt in prompts:
-        outputs = engine.run(
-            prompt=prompt,
-            inference_args=inference_args,
-        )
-        
-        # Process outputs
-        videos = rearrange(outputs["samples"], "b c t h w -> t b c h w")
-        frames = []
-        for x in videos:
-            x = torchvision.utils.make_grid(x, nrow=6)
-            x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-            frames.append((x * 255).numpy().astype(np.uint8))
-            
-        # Save video
-        os.makedirs(os.path.dirname(inference_args.output_path), exist_ok=True)
-        imageio.mimsave(
-            os.path.join(inference_args.output_path, f"{prompt[:100]}.mp4"), 
-            frames, 
-            fps=inference_args.fps
-        )
-
-
-if __name__ == "__main__":
-    parser = FlexibleArgumentParser()
-    InferenceArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    inference_args = InferenceArgs.from_cli_args(args)
-    inference_args.check_inference_args()
-
-    # Validate arguments
-    if inference_args.vae_sp and not inference_args.vae_tiling:
-        raise ValueError("Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True.")
-        
-    main(inference_args)

fastvideo/v1/test_hunyuan_vae.py

@@ -0,0 +1,161 @@
+from fastvideo.models.vaes import AutoencoderKLHunyuanVideo as MyHunyuanVAE 
+from diffusers import AutoencoderKLHunyuanVideo as DiffusersHunyuanVAE 
+
+import os
+import torch
+import torch.nn as nn
+import argparse
+import numpy as np
+import json
+from fastvideo.logger import init_logger
+from safetensors.torch import load_file
+logger = init_logger(__name__)
+
+def initialize_identical_weights(model1, model2, seed=42):
+    """Initialize both models with identical weights using a fixed seed for reproducibility."""
+    # Get all parameters from both models
+    params1 = dict(model1.named_parameters())
+    params2 = dict(model2.named_parameters())
+    
+    # Initialize each layer with identical values
+    with torch.no_grad():
+        # Initialize weights
+        for name1, param1 in params1.items():
+            if 'weight' in name1:
+                # Set seed before each weight initialization
+                torch.manual_seed(seed)
+                nn.init.normal_(param1, mean=0.0, std=0.05)
+                
+        for name2, param2 in params2.items():
+            if 'weight' in name2:
+                # Reset seed to get same initialization
+                torch.manual_seed(seed)
+                nn.init.normal_(param2, mean=0.0, std=0.05)
+                
+        # Initialize biases
+        for name1, param1 in params1.items():
+            if 'bias' in name1:
+                torch.manual_seed(seed)
+                nn.init.normal_(param1, mean=0.0, std=0.05)
+                param1.data = param1.data.to(torch.bfloat16)
+                
+        for name2, param2 in params2.items():
+            if 'bias' in name2:
+                torch.manual_seed(seed)
+                nn.init.normal_(param2, mean=0.0, std=0.05)
+                param2.data = param2.data.to(torch.bfloat16)
+    
+    logger.info("Both models initialized with identical weights in bfloat16")
+    return model1, model2
+
+def setup_args():
+    parser = argparse.ArgumentParser(description='HunyuanVAE Test')
+    parser.add_argument('--in-channels', type=int, default=4,
+                        help='Number of input channels')
+    parser.add_argument('--out-channels', type=int, default=4,
+                        help='Number of output channels')
+    parser.add_argument('--latent-channels', type=int, default=4,
+                        help='Number of latent channels')
+    return parser.parse_args()
+
+def test_hunyuan_vae():
+    args = setup_args()
+    
+    # Set fixed random seed for reproducibility
+    torch.manual_seed(42)
+    np.random.seed(42)
+    
+    # Model parameters
+    in_channels = args.in_channels
+    out_channels = args.out_channels
+    latent_channels = args.latent_channels
+    
+    device = torch.device("cuda:0")
+    print
+    # Initialize the two model implementations
+    path = "/mbz/users/hao.zhang/peiyuan/FastVideo/data/hunyuanvideo_community/vae"
+    config_path = os.path.join(path, "config.json")
+    config = json.load(open(config_path))
+    config.pop("_class_name")
+    config.pop("_diffusers_version")
+    model1 = MyHunyuanVAE(
+        **config
+    ).to(torch.bfloat16)
+    
+    model2 = DiffusersHunyuanVAE(
+        **config
+    ).to(torch.bfloat16)
+    
+    loaded = load_file(os.path.join(path, "diffusion_pytorch_model.safetensors"))
+    model1.load_state_dict(loaded)
+    model2.load_state_dict(loaded)
+
+    
+    # Set both models to eval mode
+    model1.eval()
+    model2.eval()
+    
+    # Move to GPU
+    model1 = model1.to(device)
+    model2 = model2.to(device)
+    
+
+    model1.enable_tiling(
+        tile_sample_min_height=32,
+        tile_sample_min_width=32,
+        tile_sample_min_num_frames=8,
+        tile_sample_stride_height=16,
+        tile_sample_stride_width=16,
+        tile_sample_stride_num_frames=4
+    )
+    model2.enable_tiling(
+        tile_sample_min_height=32,
+        tile_sample_min_width=32,
+        tile_sample_min_num_frames=8,
+        tile_sample_stride_height=16,
+        tile_sample_stride_width=16,
+        tile_sample_stride_num_frames=4
+    )
+    
+    # Create identical inputs for both models
+    batch_size = 1
+    
+    # Video input [B, C, T, H, W]
+    input_tensor = torch.randn(batch_size, 3, 21, 64, 64, device=device, dtype=torch.bfloat16)
+    
+    # Disable gradients for inference
+    with torch.no_grad():
+        # Test encoding
+        logger.info("Testing encoding...")
+        latent1 = model1.encode(input_tensor).mean
+        print("--------------------------------")
+        latent2 = model2.encode(input_tensor).latent_dist.mean
+        # Check if latents have the same shape
+        assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}"
+        assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}"
+        # Check if latents are similar
+        max_diff_encode = torch.max(torch.abs(latent1 - latent2))
+        mean_diff_encode = torch.mean(torch.abs(latent1 - latent2))
+        logger.info(f"Maximum difference between encoded latents: {max_diff_encode.item()}")
+        logger.info(f"Mean difference between encoded latents: {mean_diff_encode.item()}")
+        assert max_diff_encode < 1e-4, f"Encoded latents differ significantly: max diff = {max_diff_encode.item()}"
+        # Test decoding
+        logger.info("Testing decoding...")
+        output1 = model1.decode(latent1)
+        output2 = model2.decode(latent2).sample
+        # Check if outputs have the same shape
+        assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
+        
+        # Check if outputs are similar
+        max_diff_decode = torch.max(torch.abs(output1 - output2))
+        mean_diff_decode = torch.mean(torch.abs(output1 - output2))
+        logger.info(f"Maximum difference between decoded outputs: {max_diff_decode.item()}")
+        logger.info(f"Mean difference between decoded outputs: {mean_diff_decode.item()}")
+        assert max_diff_decode < 1e-4, f"Decoded outputs differ significantly: max diff = {max_diff_decode.item()}"
+
+    
+    logger.info("Test passed! Both VAE implementations produce similar outputs.")
+    logger.info("Test completed successfully")
+
+if __name__ == "__main__":
+    test_hunyuan_vae()

fastvideo/v1/test_hunyuanvideo.py

@@ -20,7 +20,7 @@ from fastvideo.utils.parallel_states import initialize_sequence_parallel_state
 from torch.distributed._composable.fsdp import CPUOffloadPolicy, fully_shard
 from torch.distributed.device_mesh import init_device_mesh
 from fastvideo.loader.fsdp_load import shard_model
-from fastvideo.models.dits.hunyuanvideo import HunyuanVideoDiT
+from fastvideo.models.dits.hunyuanvideo import HunyuanVideoTransformer3DModel as HunyuanVideoDit 
 from fastvideo.models.hunyuan.modules.models import HYVideoDiffusionTransformer
 from fastvideo.models.hunyuan_hf.modeling_hunyuan import HunyuanVideoTransformer3DModel
 logger = init_logger(__name__)
@@ -116,7 +116,7 @@ def test_hunyuanvideo_distributed():
     patch_size = [1, 2, 2]
     torch.cuda.set_device(f"cuda:{local_rank}")
     # Initialize the two model implementations
-    model1 = HunyuanVideoDiT(
+    model1 = HunyuanVideoDit(
         patch_size=2,
         patch_size_t=1,
         in_channels=4,

fastvideo/v1/test_hunyuanvideo_load.py

@@ -15,7 +15,7 @@ from fastvideo.distributed.parallel_state import (
     cleanup_dist_env_and_memory
 )
 import json
-from fastvideo.models.dits.hunyuanvideo import HunyuanVideoDiT
+from fastvideo.models.dits.hunyuanvideo import HunyuanVideoTransformer3DModel as HunyuanVideoDit 
 from fastvideo.models.hunyuan.modules.models import HUNYUAN_VIDEO_CONFIG
 from fastvideo.models.hunyuan.modules.models import HYVideoDiffusionTransformer
 from fastvideo.loader.fsdp_load import load_fsdp_model

scripts/inference/v1_inference_hunyuan.sh

@@ -1,7 +1,8 @@
 #!/bin/bash
 
 num_gpus=4
-export MODEL_BASE=data/FastHunyuan
+export MODEL_BASE=FastVideo/FastHunyuan-diffusers
+# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
     fastvideo/sample/v1_fastvideo_inference.py \
     --sp_size 4 \
@@ -15,7 +16,6 @@ torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
     --flow-reverse \
     --prompt ./assets/prompt.txt \
     --seed 1024 \
-    --output_path outputs_video/hunyuan/vae_sp_v1/ \
+    --output_path outputs_video/hunyuan/vae_sp_v2/ \
     --model_path $MODEL_BASE \
-    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
     --vae-sp

scripts/inference/v2_inference_hunyuan.sh

@@ -1,21 +0,0 @@
-#!/bin/bash
-
-num_gpus=4
-export MODEL_BASE=data/FastHunyuan
-torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
-    fastvideo/sample/v2_fastvideo_inference.py \
-    --sp_size 4 \
-    --height 720 \
-    --width 1280 \
-    --num_frames 125 \
-    --num_inference_steps 6 \
-    --guidance_scale 1 \
-    --embedded_cfg_scale 6 \
-    --flow_shift 17 \
-    --flow-reverse \
-    --prompt ./assets/prompt.txt \
-    --seed 1024 \
-    --output_path outputs_video/hunyuan/vae_sp_v2/ \
-    --model_path $MODEL_BASE \
-    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
-    --vae-sp