Adapt to new registry format

examples/inference/basic/basic_hy15.py

@@ -39,4 +39,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()

examples/inference/basic/basic_hy15_1080p.py

@@ -0,0 +1,43 @@
+from fastvideo import VideoGenerator
+import json
+# from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples_hy15_1080p"
+def main():
+    # FastVideo will automatically use the optimal default arguments for the
+    # model.
+    # If a local path is provided, FastVideo will make a best effort
+    # attempt to identify the optimal arguments.
+    generator = VideoGenerator.from_pretrained(
+        "weizhou03/HunyuanVideo-1.5-Diffusers-1080p-2SR", # 480p -> 720p -> 1080p
+        # or "weizhou03/HunyuanVideo-1.5-Diffusers-1080p" # 720p -> 1080p
+        # FastVideo will automatically handle distributed setup
+        num_gpus=1,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
+        dit_cpu_offload=True,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
+        # image_encoder_cpu_offload=False,
+    )
+
+    prompt = (
+        "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."
+    )
+
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="")
+
+    prompt2 = (
+        "A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
+
+    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, negative_prompt="")
+
+
+if __name__ == "__main__":
+    main()

fastvideo/configs/models/__init__.py

@@ -5,6 +5,7 @@ from fastvideo.configs.models.vaes.base import VAEConfig
 from fastvideo.configs.models.audio import (LTX2AudioDecoderConfig,
                                             LTX2AudioEncoderConfig,
                                             LTX2VocoderConfig)
+from fastvideo.configs.models.upsamplers.base import UpsamplerConfig
 
 __all__ = [
     "ModelConfig",
@@ -14,4 +15,5 @@ __all__ = [
     "LTX2AudioEncoderConfig",
     "LTX2AudioDecoderConfig",
     "LTX2VocoderConfig",
+    "UpsamplerConfig",
 ]

fastvideo/configs/models/dits/hunyuanvideo15.py

@@ -55,6 +55,8 @@ class HunyuanVideo15ArchConfig(DiTArchConfig):
             r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",
             r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear\.(.*)$":
             r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.self_attn_qkv\.(.*)$":
+            r"txt_in.refiner_blocks.\1.self_attn_qkv.\2",
 
             # 2. txt_in_2 mapping:
             r"^context_embedder_2\.(.*)$":

fastvideo/configs/models/upsamplers/__init__.py

@@ -0,0 +1,6 @@
+from fastvideo.configs.models.upsamplers.hunyuan15 import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig
+from fastvideo.configs.models.upsamplers.base import UpsamplerConfig
+
+__all__ = [
+    "SRTo720pUpsamplerConfig", "SRTo1080pUpsamplerConfig", "UpsamplerConfig"
+]

fastvideo/configs/models/upsamplers/base.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+from fastvideo.configs.models.base import ModelConfig
+
+
+@dataclass
+class UpsamplerConfig(ModelConfig):
+    pass

fastvideo/configs/models/upsamplers/hunyuan15.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from fastvideo.configs.models.upsamplers.base import UpsamplerConfig
+
+
+@dataclass
+class SRTo720pUpsamplerConfig(UpsamplerConfig):
+    in_channels: int = 0
+    out_channels: int = 0
+    hidden_channels: int = 64
+    num_blocks: int = 6
+    global_residual: bool = False
+
+
+@dataclass
+class SRTo1080pUpsamplerConfig(UpsamplerConfig):
+    z_channels: int = 0
+    out_channels: int = 0
+    block_out_channels: tuple[int, ...] = (0, 0)
+    num_res_blocks: int = 2
+    is_residual: bool = False

fastvideo/configs/pipelines/base.py

@@ -8,7 +8,7 @@ from typing import Any, cast
 import torch
 
 from fastvideo.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
-                                      VAEConfig)
+                                      VAEConfig, UpsamplerConfig)
 from fastvideo.configs.models.encoders import BaseEncoderOutput
 from fastvideo.configs.utils import update_config_from_args
 from fastvideo.logger import init_logger
@@ -44,12 +44,15 @@ class PipelineConfig:
     # Video generation parameters
     embedded_cfg_scale: float = 6.0
     flow_shift: float | None = None
+    flow_shift_sr: float | None = None
     disable_autocast: bool = False
     is_causal: bool = False
 
     # Model configuration
     dit_config: DiTConfig = field(default_factory=DiTConfig)
     dit_precision: str = "bf16"
+    upsampler_config: UpsamplerConfig = field(default_factory=UpsamplerConfig)
+    upsampler_precision: str = "fp32"
 
     # VAE configuration
     vae_config: VAEConfig = field(default_factory=VAEConfig)

fastvideo/configs/pipelines/hunyuan15.py

@@ -11,7 +11,8 @@ from fastvideo.configs.models.dits import HunyuanVideo15Config
 from fastvideo.configs.models.encoders import (BaseEncoderOutput,
                                                Qwen2_5_VLConfig, T5Config)
 from fastvideo.configs.models.vaes import Hunyuan15VAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.models.upsamplers import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig
+from fastvideo.configs.pipelines.base import PipelineConfig, UpsamplerConfig
 
 PROMPT_TEMPLATE_TOKEN_LENGTH = 108
 
@@ -131,9 +132,35 @@ class Hunyuan15T2V480PConfig(PipelineConfig):
         self.vae_config.load_decoder = True
 
 
+@dataclass
+class Hunyuan15I2V480PStepDistilledConfig(Hunyuan15T2V480PConfig):
+    flow_shift: int = 7
+
+
 @dataclass
 class Hunyuan15T2V720PConfig(Hunyuan15T2V480PConfig):
     """Base configuration for HunYuan pipeline architecture."""
 
     # HunyuanConfig-specific parameters with defaults
     flow_shift: int = 9
+
+
+@dataclass
+class Hunyuan15I2V720PConfig(Hunyuan15T2V720PConfig):
+    """Base configuration for HunYuan pipeline architecture."""
+
+    # HunyuanConfig-specific parameters with defaults
+    flow_shift: int = 7
+
+
+@dataclass
+class Hunyuan15SR1080PConfig(Hunyuan15T2V720PConfig):
+    """Base configuration for HunYuan pipeline architecture."""
+
+    # HunyuanConfig-specific parameters with defaults
+    flow_shift: int = 7
+    flow_shift_sr: int = 2
+    upsampler_config: tuple[UpsamplerConfig, ...] = field(
+        default_factory=lambda:
+        (SRTo720pUpsamplerConfig(), SRTo1080pUpsamplerConfig()))
+    upsampler_precision: str = "fp32"

fastvideo/configs/sample/base.py

@@ -58,10 +58,13 @@ class SamplingParam:
     num_frames_round_down: bool = False  # Whether to round down num_frames if it's not divisible by num_gpus
     height: int = 720
     width: int = 1280
+    height_sr: int = 1072
+    width_sr: int = 1920
     fps: int = 24
 
     # Denoising parameters
     num_inference_steps: int = 50
+    num_inference_steps_sr: int = 50
     guidance_scale: float = 1.0
     guidance_rescale: float = 0.0
     boundary_ratio: float | None = None

fastvideo/configs/sample/hunyuan15.py

@@ -28,7 +28,33 @@ class Hunyuan15_480P_SamplingParam(SamplingParam):
             np.linspace(1.0, 0.0, self.num_inference_steps + 1)[:-1])
 
 
+@dataclass
+class Hunyuan15_480P_StepDistilled_I2V_SamplingParam(
+        Hunyuan15_480P_SamplingParam):
+    num_inference_steps: int = 12
+
+    height: int = 720
+    width: int = 1280
+    guidance_scale: float = 1.0
+
+
 @dataclass
 class Hunyuan15_720P_SamplingParam(Hunyuan15_480P_SamplingParam):
     height: int = 720
     width: int = 1280
+
+
+@dataclass
+class Hunyuan15_720P_Distilled_I2V_SamplingParam(Hunyuan15_720P_SamplingParam):
+    guidance_scale: float = 1.0
+
+
+@dataclass
+class Hunyuan15_SR_1080P_SamplingParam(Hunyuan15_480P_SamplingParam):
+    height_sr: int = 1072
+    width_sr: int = 1920
+
+    num_inference_steps: int = 12
+    num_inference_steps_sr: int = 8
+
+    guidance_scale: float = 1.0

fastvideo/fastvideo_args.py

@@ -179,6 +179,7 @@ class FastVideoArgs:
     model_loaded: dict[str, bool] = field(default_factory=lambda: {
         "transformer": True,
         "vae": True,
+        "upsampler": True,
     })
 
     override_text_encoder_safetensors: str | None = None  # path to safetensors file for text encoder override

fastvideo/models/loader/component_loader.py

@@ -79,6 +79,7 @@ class ComponentLoader(ABC):
             "scheduler": (SchedulerLoader, "diffusers"),
             "transformer": (TransformerLoader, "diffusers"),
             "transformer_2": (TransformerLoader, "diffusers"),
+            "transformer_3": (TransformerLoader, "diffusers"),
             "vae": (VAELoader, "diffusers"),
             "audio_vae": (AudioDecoderLoader, "diffusers"),
             "audio_decoder": (AudioDecoderLoader, "diffusers"),
@@ -90,6 +91,8 @@ class ComponentLoader(ABC):
             "image_processor": (ImageProcessorLoader, "transformers"),
             "feature_extractor": (ImageProcessorLoader, "transformers"),
             "image_encoder": (ImageEncoderLoader, "transformers"),
+            "upsampler": (UpsamplerLoader, "diffusers"),
+            "upsampler_2": (UpsamplerLoader, "diffusers"),
         }
 
         if module_type in module_loaders:
@@ -558,7 +561,8 @@ class VAELoader(ComponentLoader):
     def load(self, model_path: str, fastvideo_args: FastVideoArgs):
         """Load the VAE based on the model path, and inference args."""
         config = get_diffusers_config(model=model_path)
-        class_name = config.get("_class_name")
+        class_name = config.pop("_class_name")
+        config.pop("_name_or_path", None)
         assert class_name is not None, (
             "Model config does not contain a _class_name attribute. Only diffusers format is supported."
         )
@@ -731,6 +735,7 @@ class TransformerLoader(ComponentLoader):
         config = get_diffusers_config(model=model_path)
         hf_config = deepcopy(config)
         cls_name = config.pop("_class_name")
+        config.pop("_name_or_path", None)
         if cls_name is None:
             raise ValueError(
                 "Model config does not contain a _class_name attribute. "
@@ -745,7 +750,7 @@ class TransformerLoader(ComponentLoader):
         fastvideo_args.model_paths["transformer"] = model_path
 
         # Config from Diffusers supersedes fastvideo's model config
-        dit_config = fastvideo_args.pipeline_config.dit_config
+        dit_config = deepcopy(fastvideo_args.pipeline_config.dit_config)
         dit_config.update_model_arch(config)
 
         model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
@@ -877,6 +882,50 @@ class SchedulerLoader(ComponentLoader):
         return scheduler
 
 
+class UpsamplerLoader(ComponentLoader):
+    """Loader for upsamplers."""
+
+    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
+        """Load the upsampler based on the model path, and inference args."""
+        config_dict = get_diffusers_config(model=model_path)
+        class_name = config_dict.pop("_class_name", None)
+
+        if class_name is None:
+            raise ValueError(
+                "Model config does not contain a _class_name attribute. "
+                "Only diffusers format is supported."
+            )
+
+        try:
+            upsampler_cfg = deepcopy(fastvideo_args.pipeline_config.upsampler_config[0])
+            upsampler_cfg.update_model_config(config_dict)
+        except Exception as e:
+            upsampler_cfg = deepcopy(fastvideo_args.pipeline_config.upsampler_config[1])
+            upsampler_cfg.update_model_config(config_dict)
+
+        model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+        model = model_cls(upsampler_cfg)
+
+        target_device = get_local_torch_device()
+        model = model.to(target_device, dtype=PRECISION_TO_TYPE[fastvideo_args.pipeline_config.upsampler_precision])
+
+        # 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}")
+        
+        if len(safetensors_list) == 1:
+            loaded = safetensors_load_file(safetensors_list[0])
+        else:
+            loaded = {}
+            for sf_file in safetensors_list:
+                loaded.update(safetensors_load_file(sf_file))
+        
+        model.load_state_dict(loaded, strict=True)
+
+        return model.eval()
+
 class GenericComponentLoader(ComponentLoader):
     """Generic loader for components that don't have a specific loader."""
 

fastvideo/models/registry.py

@@ -99,6 +99,11 @@ _SCHEDULERS = {
     ("schedulers", "scheduling_rcm", "RCMScheduler"),
 }
 
+_UPSAMPLERS = {
+    "SRTo720pUpsampler": ("upsamplers", "hunyuan15", "SRTo720pUpsampler"),
+    "SRTo1080pUpsampler": ("upsamplers", "hunyuan15", "SRTo1080pUpsampler"),
+}
+
 _LEGACY_FAST_VIDEO_MODELS = {
     **_TEXT_TO_VIDEO_DIT_MODELS,
     **_IMAGE_TO_VIDEO_DIT_MODELS,
@@ -107,6 +112,7 @@ _LEGACY_FAST_VIDEO_MODELS = {
     **_VAE_MODELS,
     **_AUDIO_MODELS,
     **_SCHEDULERS,
+    **_UPSAMPLERS,
 }
 
 MODELS_PATH = os.path.dirname(__file__)

fastvideo/models/upsamplers/hunyuan15.py

@@ -0,0 +1,170 @@
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE
+#
+# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any
+# output and results therefrom are provided "AS IS" without any express or implied warranties of
+# any kind including any warranties of title, merchantability, noninfringement, course of dealing,
+# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the
+# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of
+# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a
+# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise
+# of rights and permissions under this agreement.
+# See the License for the specific language governing permissions and limitations under the License.
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import Enum
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Tensor
+
+from fastvideo.models.vaes.hunyuan15vae import (
+    HunyuanVideo15CausalConv3d,
+    HunyuanVideo15RMS_norm,
+)
+
+from fastvideo.layers.activation import get_act_fn
+from fastvideo.configs.models.upsamplers import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig
+
+class HunyuanVideo15ResnetBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int | None = None,
+        non_linearity: str = "swish",
+    ) -> None:
+        super().__init__()
+        out_channels = out_channels or in_channels
+
+        self.nonlinearity = get_act_fn(non_linearity)
+
+        self.norm1 = HunyuanVideo15RMS_norm(in_channels, images=False)
+        self.conv1 = HunyuanVideo15CausalConv3d(in_channels, out_channels, kernel_size=3)
+
+        self.norm2 = HunyuanVideo15RMS_norm(out_channels, images=False)
+        self.conv2 = HunyuanVideo15CausalConv3d(out_channels, out_channels, kernel_size=3)
+
+        self.nin_shortcut = None
+        if in_channels != out_channels:
+            self.nin_shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        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.conv2(hidden_states)
+
+        if self.nin_shortcut is not None:
+            residual = self.nin_shortcut(residual)
+
+        return hidden_states + residual
+
+class SRResidualCausalBlock3D(nn.Module):
+    def __init__(self, channels: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            HunyuanVideo15CausalConv3d(channels, channels, kernel_size=3),
+            nn.SiLU(inplace=True),
+            HunyuanVideo15CausalConv3d(channels, channels, kernel_size=3),
+            nn.SiLU(inplace=True),
+            HunyuanVideo15CausalConv3d(channels, channels, kernel_size=3),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x + self.block(x)
+
+
+class SRTo720pUpsampler(nn.Module):
+
+    def __init__(
+        self,
+        config: SRTo720pUpsamplerConfig,
+    ):
+        super().__init__()
+        self.in_conv = HunyuanVideo15CausalConv3d(config.in_channels, config.hidden_channels, kernel_size=3)
+        self.blocks = nn.ModuleList([SRResidualCausalBlock3D(config.hidden_channels) for _ in range(config.num_blocks)])
+        self.out_conv = HunyuanVideo15CausalConv3d(config.hidden_channels, config.out_channels, kernel_size=3)
+        self.global_residual = bool(config.global_residual)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        y = self.in_conv(x)
+        for blk in self.blocks:
+            y = blk(y)
+        y = self.out_conv(y)
+        if self.global_residual and (y.shape == residual.shape):
+            y = y + residual
+        return y
+
+
+class SRTo1080pUpsampler(nn.Module):
+
+    def __init__(
+        self,
+        config: SRTo1080pUpsamplerConfig,
+    ):
+        super().__init__()
+        self.num_res_blocks = config.num_res_blocks
+        self.block_out_channels = config.block_out_channels
+        self.z_channels = config.z_channels
+
+        block_in = config.block_out_channels[0]
+        self.conv_in = HunyuanVideo15CausalConv3d(config.z_channels, block_in, kernel_size=3)
+
+        self.up = nn.ModuleList()
+        for i_level, ch in enumerate(config.block_out_channels):
+            block = nn.ModuleList()
+            block_out = ch
+            for _ in range(self.num_res_blocks + 1):
+                block.append(HunyuanVideo15ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+
+            self.up.append(up)
+
+        self.norm_out = HunyuanVideo15RMS_norm(block_in, images=False)
+        self.conv_out = HunyuanVideo15CausalConv3d(block_in, config.out_channels, kernel_size=3)
+
+        self.gradient_checkpointing = False
+        self.is_residual = config.is_residual
+
+    def forward(self, z: Tensor, target_shape: Sequence[int] = None) -> Tensor:
+        """
+        Args:
+            z: (B, C, T, H, W)
+            target_shape: (H, W)
+        """
+        if target_shape is not None and z.shape[-2:] != target_shape:
+            bsz = z.shape[0]
+            z = rearrange(z, "b c f h w -> (b f) c h w")
+            z = F.interpolate(z, size=target_shape, mode="bilinear", align_corners=False)
+            z = rearrange(z, "(b f) c h w -> b c f h w", b=bsz)
+
+        # z to block_in
+        repeats = self.block_out_channels[0] // (self.z_channels)
+        h = self.conv_in(z) + z.repeat_interleave(repeats=repeats, dim=1)
+
+        # upsampling
+        for i_level in range(len(self.block_out_channels)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+            if hasattr(self.up[i_level], "upsample"):
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = get_act_fn("swish")(h)
+        h = self.conv_out(h)
+        return h

fastvideo/pipelines/basic/hunyuan15/hunyuan15_2sr_pipeline.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Hunyuan video diffusion pipeline implementation.
+
+This module contains an implementation of the Hunyuan video diffusion pipeline
+using the modular pipeline architecture.
+"""
+import torch
+import time
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import (
+    ConditioningStage, DecodingStage, DenoisingStage, InputValidationStage,
+    LatentPreparationStage, TextEncodingStage, TimestepPreparationStage,
+    Hy15ImageEncodingStage, SRDenoisingStage)
+from fastvideo.distributed.parallel_state import get_local_torch_device
+
+# TODO(will): move PRECISION_TO_TYPE to better place
+
+logger = init_logger(__name__)
+
+
+class HunyuanVideo152SRPipeline(ComposedPipelineBase):
+
+    _required_config_modules = [
+        "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae",
+        "transformer", "transformer_2", "transformer_3", "scheduler",
+        "upsampler", "upsampler_2"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(stage_name="prompt_encoding_stage_primary",
+                       stage=TextEncodingStage(
+                           text_encoders=[
+                               self.get_module("text_encoder"),
+                               self.get_module("text_encoder_2")
+                           ],
+                           tokenizers=[
+                               self.get_module("tokenizer"),
+                               self.get_module("tokenizer_2")
+                           ],
+                       ))
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer")))
+
+        self.add_stage(stage_name="image_encoding_stage",
+                       stage=Hy15ImageEncodingStage(image_encoder=None,
+                                                    image_processor=None))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="sr_720p_latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer_2")))
+
+        self.add_stage(stage_name="sr_720p_denoising_stage",
+                       stage=SRDenoisingStage(
+                           transformer=self.get_module("transformer_2"),
+                           scheduler=self.get_module("scheduler"),
+                           upsampler=self.get_module("upsampler")))
+
+        self.add_stage(stage_name="sr_1080p_latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer_3")))
+
+        self.add_stage(stage_name="sr_1080p_denoising_stage",
+                       stage=SRDenoisingStage(
+                           transformer=self.get_module("transformer_3"),
+                           scheduler=self.get_module("scheduler"),
+                           upsampler=self.get_module("upsampler_2")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+    @torch.no_grad()
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Generate a video or image using the pipeline.
+        
+        Args:
+            batch: The batch to generate from.
+            fastvideo_args: The inference arguments.
+        Returns:
+            ForwardBatch: The batch with the generated video or image.
+        """
+        if not self.post_init_called:
+            self.post_init()
+
+        self.get_module("transformer").to(get_local_torch_device())
+        # Execute each stage
+        logger.info("Running pipeline stages: %s",
+                    self._stage_name_mapping.keys())
+        # logger.info("Batch: %s", batch)
+        batch = self.input_validation_stage(batch, fastvideo_args)
+        batch = self.prompt_encoding_stage_primary(batch, fastvideo_args)
+        batch = self.conditioning_stage(batch, fastvideo_args)
+        batch = self.timestep_preparation_stage(batch, fastvideo_args)
+        batch = self.latent_preparation_stage(batch, fastvideo_args)
+        batch = self.image_encoding_stage(batch, fastvideo_args)
+        batch = self.denoising_stage(batch, fastvideo_args)
+        self.get_module("transformer").to("cpu")
+
+        # 720p SR
+        self.get_module("transformer_2").to(get_local_torch_device())
+        batch.lq_latents = batch.latents
+        batch.latents = None
+        batch.height = 720
+        batch.width = 1280
+        batch.num_inference_steps_sr = 6
+        batch = self.sr_720p_latent_preparation_stage(batch, fastvideo_args)
+        batch = self.image_encoding_stage(batch, fastvideo_args)
+        batch = self.sr_720p_denoising_stage(batch, fastvideo_args)
+        self.get_module("transformer_2").to("cpu")
+
+        # 1080p SR
+        self.get_module("transformer_3").to(get_local_torch_device())
+        batch.lq_latents = batch.latents
+        batch.latents = None
+        batch.height = 1072
+        batch.width = 1920
+        batch.num_inference_steps_sr = 8
+        batch = self.sr_1080p_latent_preparation_stage(batch, fastvideo_args)
+        batch = self.image_encoding_stage(batch, fastvideo_args)
+        batch = self.sr_1080p_denoising_stage(batch, fastvideo_args)
+        self.get_module("transformer_3").to("cpu")
+
+        start_time = time.time()
+        batch = self.decoding_stage(batch, fastvideo_args)
+        end_time = time.time()
+        logger.info("Decoding time: %s seconds", end_time - start_time)
+
+        # Return the output
+        return batch
+
+
+EntryClass = HunyuanVideo152SRPipeline

fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Hunyuan video diffusion pipeline implementation.
+
+This module contains an implementation of the Hunyuan video diffusion pipeline
+using the modular pipeline architecture.
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
+                                        DenoisingStage, InputValidationStage,
+                                        LatentPreparationStage,
+                                        TextEncodingStage,
+                                        TimestepPreparationStage,
+                                        Hy15ImageEncodingStage)
+
+# TODO(will): move PRECISION_TO_TYPE to better place
+
+logger = init_logger(__name__)
+
+
+class HunyuanVideo15ImageToVideoPipeline(ComposedPipelineBase):
+
+    _required_config_modules = [
+        "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae",
+        "transformer", "scheduler"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(stage_name="prompt_encoding_stage_primary",
+                       stage=TextEncodingStage(
+                           text_encoders=[
+                               self.get_module("text_encoder"),
+                               self.get_module("text_encoder_2")
+                           ],
+                           tokenizers=[
+                               self.get_module("tokenizer"),
+                               self.get_module("tokenizer_2")
+                           ],
+                       ))
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer")))
+
+        self.add_stage(stage_name="image_encoding_stage",
+                       stage=Hy15ImageEncodingStage(image_encoder=None,
+                                                    image_processor=None))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+
+EntryClass = HunyuanVideo15ImageToVideoPipeline

fastvideo/pipelines/basic/hunyuan15/hunyuan15_sr_pipeline.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Hunyuan video diffusion pipeline implementation.
+
+This module contains an implementation of the Hunyuan video diffusion pipeline
+using the modular pipeline architecture.
+"""
+import torch
+import time
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import (
+    ConditioningStage, DecodingStage, DenoisingStage, InputValidationStage,
+    LatentPreparationStage, TextEncodingStage, TimestepPreparationStage,
+    Hy15ImageEncodingStage, SRDenoisingStage)
+from fastvideo.distributed.parallel_state import get_local_torch_device
+
+# TODO(will): move PRECISION_TO_TYPE to better place
+
+logger = init_logger(__name__)
+
+
+class HunyuanVideo15SRPipeline(ComposedPipelineBase):
+
+    _required_config_modules = [
+        "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae",
+        "transformer", "transformer_2", "scheduler", "upsampler"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(stage_name="prompt_encoding_stage_primary",
+                       stage=TextEncodingStage(
+                           text_encoders=[
+                               self.get_module("text_encoder"),
+                               self.get_module("text_encoder_2")
+                           ],
+                           tokenizers=[
+                               self.get_module("tokenizer"),
+                               self.get_module("tokenizer_2")
+                           ],
+                       ))
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer")))
+
+        self.add_stage(stage_name="image_encoding_stage",
+                       stage=Hy15ImageEncodingStage(image_encoder=None,
+                                                    image_processor=None))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="sr_latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer_2")))
+
+        self.add_stage(stage_name="sr_denoising_stage",
+                       stage=SRDenoisingStage(
+                           transformer=self.get_module("transformer_2"),
+                           scheduler=self.get_module("scheduler"),
+                           upsampler=self.get_module("upsampler")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+    @torch.no_grad()
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Generate a video or image using the pipeline.
+        
+        Args:
+            batch: The batch to generate from.
+            fastvideo_args: The inference arguments.
+        Returns:
+            ForwardBatch: The batch with the generated video or image.
+        """
+        if not self.post_init_called:
+            self.post_init()
+
+        # Execute each stage
+        logger.info("Running pipeline stages: %s",
+                    self._stage_name_mapping.keys())
+        # logger.info("Batch: %s", batch)
+        self.get_module("transformer").to(get_local_torch_device())
+        batch = self.input_validation_stage(batch, fastvideo_args)
+        batch = self.prompt_encoding_stage_primary(batch, fastvideo_args)
+        batch = self.conditioning_stage(batch, fastvideo_args)
+        batch = self.timestep_preparation_stage(batch, fastvideo_args)
+        batch = self.latent_preparation_stage(batch, fastvideo_args)
+        batch = self.image_encoding_stage(batch, fastvideo_args)
+        batch = self.denoising_stage(batch, fastvideo_args)
+        self.get_module("transformer").to("cpu")
+
+        self.get_module("transformer_2").to(get_local_torch_device())
+        batch.lq_latents = batch.latents
+        batch.latents = None
+        batch.height = batch.height_sr
+        batch.width = batch.width_sr
+        batch = self.sr_latent_preparation_stage(batch, fastvideo_args)
+        batch = self.image_encoding_stage(batch, fastvideo_args)
+        batch = self.sr_denoising_stage(batch, fastvideo_args)
+        self.get_module("transformer_2").to("cpu")
+
+        start_time = time.time()
+        batch = self.decoding_stage(batch, fastvideo_args)
+        end_time = time.time()
+        logger.info("Decoding time: %s seconds", end_time - start_time)
+
+        # Return the output
+        return batch
+
+
+EntryClass = HunyuanVideo15SRPipeline

fastvideo/pipelines/composed_pipeline_base.py

@@ -293,6 +293,7 @@ class ComposedPipelineBase(ABC):
         # remove keys that are not pipeline modules
         model_index.pop("_class_name")
         model_index.pop("_diffusers_version")
+        model_index.pop("_name_or_path", None)
         model_index.pop("workload_type", None)
         if "boundary_ratio" in model_index and model_index[
                 "boundary_ratio"] is not None:

fastvideo/pipelines/pipeline_batch_info.py

@@ -123,6 +123,7 @@ class ForwardBatch:
 
     # Latent tensors
     latents: torch.Tensor | None = None
+    lq_latents: torch.Tensor | None = None
     raw_latent_shape: tuple[int, ...] | None = None
     noise_pred: torch.Tensor | None = None
     image_latent: torch.Tensor | None = None
@@ -144,6 +145,8 @@ class ForwardBatch:
     # Original dimensions (before VAE scaling)
     height: list[int] | int | None = None
     width: list[int] | int | None = None
+    height_sr: list[int] | int | None = None
+    width_sr: list[int] | int | None = None
     fps: list[int] | int | None = None
 
     # Timesteps
@@ -154,6 +157,7 @@ class ForwardBatch:
 
     # Scheduler parameters
     num_inference_steps: int = 50
+    num_inference_steps_sr: int = 50
     guidance_scale: float = 1.0
     guidance_scale_2: float | None = None
     guidance_rescale: float = 0.0

fastvideo/pipelines/stages/__init__.py

@@ -14,6 +14,7 @@ from fastvideo.pipelines.stages.denoising import (
     Cosmos25AutoDenoisingStage, Cosmos25DenoisingStage,
     Cosmos25V2WDenoisingStage, Cosmos25T2WDenoisingStage, CosmosDenoisingStage,
     DenoisingStage, DmdDenoisingStage)
+from fastvideo.pipelines.stages.sr_denoising import SRDenoisingStage
 from fastvideo.pipelines.stages.encoding import EncodingStage
 from fastvideo.pipelines.stages.image_encoding import (
     ImageEncodingStage, MatrixGameImageEncodingStage, RefImageEncodingStage,
@@ -70,6 +71,7 @@ __all__ = [
     "Cosmos25AutoDenoisingStage",
     "LTX2DenoisingStage",
     "LTX2TextEncodingStage",
+    "SRDenoisingStage",
     "EncodingStage",
     "DecodingStage",
     "ImageEncodingStage",

fastvideo/pipelines/stages/denoising.py

@@ -315,6 +315,26 @@ class DenoisingStage(PipelineStage):
                 else:
                     t_expand = t.repeat(latent_model_input.shape[0])
 
+                use_meanflow = getattr(self.transformer.config, "use_meanflow",
+                                       False)
+                if use_meanflow:
+                    if i == len(timesteps) - 1:
+                        timesteps_r = torch.tensor(
+                            [0.0], device=get_local_torch_device())
+                    else:
+                        timesteps_r = timesteps[i + 1]
+                    timesteps_r = timesteps_r.repeat(
+                        latent_model_input.shape[0])
+                else:
+                    timesteps_r = None
+
+                timesteps_r_kwarg = self.prepare_extra_func_kwargs(
+                    self.transformer.forward,
+                    {
+                        "timestep_r": timesteps_r,
+                    },
+                )
+
                 latent_model_input = self.scheduler.scale_model_input(
                     latent_model_input, t)
 
@@ -406,6 +426,7 @@ class DenoisingStage(PipelineStage):
                             **image_kwargs,
                             **pos_cond_kwargs,
                             **action_kwargs,
+                            **timesteps_r_kwarg,
                         )
 
                     if batch.do_classifier_free_guidance:
@@ -423,6 +444,7 @@ class DenoisingStage(PipelineStage):
                                 **image_kwargs,
                                 **neg_cond_kwargs,
                                 **action_kwargs,
+                                **timesteps_r_kwarg,
                             )
 
                         noise_pred_text = noise_pred

fastvideo/pipelines/stages/sr_denoising.py

@@ -0,0 +1,431 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Denoising stage for diffusion pipelines.
+"""
+
+import inspect
+import weakref
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+import numpy as np
+from tqdm.auto import tqdm
+
+from fastvideo.attention import get_attn_backend
+from fastvideo.configs.pipelines.base import STA_Mode
+from fastvideo.distributed import (get_local_torch_device, get_world_group)
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import TransformerLoader, UpsamplerLoader
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.validators import StageValidators as V
+from fastvideo.pipelines.stages.validators import VerificationResult
+from fastvideo.platforms import AttentionBackendEnum
+
+try:
+    from fastvideo.attention.backends.sliding_tile_attn import (
+        SlidingTileAttentionBackend)
+    st_attn_available = True
+except ImportError:
+    st_attn_available = False
+
+try:
+    from fastvideo.attention.backends.vmoba import VMOBAAttentionBackend
+    from fastvideo.utils import is_vmoba_available
+    vmoba_attn_available = is_vmoba_available()
+except ImportError:
+    vmoba_attn_available = False
+
+try:
+    from fastvideo.attention.backends.video_sparse_attn import (
+        VideoSparseAttentionBackend)
+    vsa_available = True
+except ImportError:
+    vsa_available = False
+
+logger = init_logger(__name__)
+
+
+class SRDenoisingStage(PipelineStage):
+    """
+    Stage for running the denoising loop in SR diffusion pipelines. Used by Hunyuan15 SR pipeline.
+    
+    This stage handles the iterative denoising process that transforms
+    the initial noise into the final output.
+    """
+
+    def __init__(self,
+                 transformer,
+                 scheduler,
+                 upsampler,
+                 pipeline=None,
+                 vae=None) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.scheduler = scheduler
+        self.vae = vae
+        self.upsampler = upsampler
+        self.pipeline = weakref.ref(pipeline) if pipeline else None
+        attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads
+        self.attn_backend = get_attn_backend(
+            head_size=attn_head_size,
+            dtype=torch.float16,  # TODO(will): hack
+            supported_attention_backends=(
+                AttentionBackendEnum.SLIDING_TILE_ATTN,
+                AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+                AttentionBackendEnum.VMOBA_ATTN,
+                AttentionBackendEnum.FLASH_ATTN,
+                AttentionBackendEnum.TORCH_SDPA,
+                AttentionBackendEnum.SAGE_ATTN_THREE)  # hack
+        )
+
+    def add_noise_to_lq(self,
+                        lq_latents: torch.Tensor,
+                        strength: float = 0.7) -> torch.Tensor:
+
+        def expand_dims(tensor: torch.Tensor, ndim: int):
+            shape = tensor.shape + (1, ) * (ndim - tensor.ndim)
+            return tensor.reshape(shape)
+
+        noise = torch.randn_like(lq_latents)
+        timestep = torch.tensor([1000.0],
+                                device=get_local_torch_device()) * strength
+        t = expand_dims(timestep, lq_latents.ndim)
+        return (1 - t / 1000.0) * lq_latents + (t / 1000.0) * noise
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Run the denoising loop.
+        
+        Args:
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
+            
+        Returns:
+            The batch with denoised latents.
+        """
+        pipeline = self.pipeline() if self.pipeline else None
+        if not fastvideo_args.model_loaded["transformer"]:
+            loader = TransformerLoader()
+            self.transformer = loader.load(
+                fastvideo_args.model_paths["transformer"], fastvideo_args)
+            if pipeline:
+                pipeline.add_module("transformer", self.transformer)
+            fastvideo_args.model_loaded["transformer"] = True
+
+        if not fastvideo_args.model_loaded["upsampler"]:
+            loader = UpsamplerLoader()
+            self.upsampler = loader.load(
+                fastvideo_args.model_paths["upsampler"], fastvideo_args)
+            if pipeline:
+                pipeline.add_module("upsampler", self.upsampler)
+            fastvideo_args.model_loaded["upsampler"] = True
+
+        # Setup precision and autocast settings
+        # TODO(will): make the precision configurable for inference
+        # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
+        target_dtype = torch.bfloat16
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        self.scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift_sr)
+        sigmas = np.linspace(1.0, 0.0, batch.num_inference_steps_sr + 1)[:-1]
+        self.scheduler.set_timesteps(sigmas=sigmas,
+                                     device=get_local_torch_device())
+        timesteps = self.scheduler.timesteps
+        logger.info("timesteps: %s", timesteps)
+        num_inference_steps = len(timesteps)
+        num_warmup_steps = len(
+            timesteps) - num_inference_steps * self.scheduler.order
+        assert num_inference_steps == batch.num_inference_steps_sr, "num_inference_steps_sr must match the number of timesteps"
+
+        pos_cond_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "encoder_hidden_states_2": batch.clip_embedding_pos,
+                "encoder_attention_mask": batch.prompt_attention_mask,
+            },
+        )
+
+        image_embeds = batch.image_embeds
+        if len(image_embeds) > 0:
+            assert not torch.isnan(
+                image_embeds[0]).any(), "image_embeds contains nan"
+            image_embeds = [
+                image_embed.to(target_dtype) for image_embed in image_embeds
+            ]
+
+        image_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "encoder_hidden_states_image": image_embeds,
+            },
+        )
+
+        # Prepare STA parameters
+        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
+            self.prepare_sta_param(batch, fastvideo_args)
+
+        # Get latents and embeddings
+        prompt_embeds = batch.prompt_embeds
+        assert not torch.isnan(
+            prompt_embeds[0]).any(), "prompt_embeds contains nan"
+
+        latents = batch.latents
+        lq_latents = batch.lq_latents
+        logger.info("lq_latents: %s", lq_latents.shape)
+        logger.info("latents: %s", latents.shape)
+        tgt_shape = latents.shape[-2:]  # (h w)
+        bsz = lq_latents.shape[0]
+        lq_latents = rearrange(lq_latents, "b c f h w -> (b f) c h w")
+        lq_latents = F.interpolate(lq_latents,
+                                   size=tgt_shape,
+                                   mode="bilinear",
+                                   align_corners=False)
+        lq_latents = rearrange(lq_latents, "(b f) c h w -> b c f h w", b=bsz)
+        lq_latents = self.upsampler(
+            lq_latents.to(dtype=torch.float32, device=get_local_torch_device()))
+        lq_latents = lq_latents.to(dtype=latents.dtype)
+        lq_latents = self.add_noise_to_lq(lq_latents, 0.7)
+        b, c, f, h, w = lq_latents.shape
+        mask_ones = torch.ones(b, 1, f, h, w).to(lq_latents.device)
+        lq_cond_latents = torch.concat([lq_latents, mask_ones],
+                                       dim=1).to(target_dtype)
+        cond_latents = torch.cat(
+            [batch.video_latent, torch.zeros_like(latents)],
+            dim=1).to(target_dtype)
+        condition = torch.concat([cond_latents, lq_cond_latents], dim=1)
+        zero_lq_condition = condition.clone()
+        zero_lq_condition[:, c + 1:2 * c + 1] = torch.zeros_like(lq_latents)
+        zero_lq_condition[:, 2 * c + 1] = 0
+
+        latent_model_input = latents.to(target_dtype)
+        assert latent_model_input.shape[0] == 1, "only support batch size 1"
+
+        # Run denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # Skip if interrupted
+                if hasattr(self, 'interrupt') and self.interrupt:
+                    continue
+
+                latent_model_input = latents.to(target_dtype)
+                if t < 1000 * 0.7:
+                    condition = zero_lq_condition
+
+                latent_model_input = torch.concat([latents, condition], dim=1)
+
+                assert not torch.isnan(
+                    latent_model_input).any(), "latent_model_input contains nan"
+                t_expand = t.repeat(latent_model_input.shape[0])
+
+                if i == len(timesteps) - 1:
+                    timesteps_r = torch.tensor([0.0],
+                                               device=get_local_torch_device())
+                else:
+                    timesteps_r = timesteps[i + 1]
+                timesteps_r = timesteps_r.repeat(latent_model_input.shape[0])
+
+                latent_model_input = self.scheduler.scale_model_input(
+                    latent_model_input, t)
+
+                # Prepare inputs for transformer
+                guidance_expand = (
+                    torch.tensor(
+                        [fastvideo_args.pipeline_config.embedded_cfg_scale] *
+                        latent_model_input.shape[0],
+                        dtype=torch.float32,
+                        device=get_local_torch_device(),
+                    ).to(target_dtype) *
+                    1000.0 if fastvideo_args.pipeline_config.embedded_cfg_scale
+                    is not None else None)
+
+                # Predict noise residual
+                with torch.autocast(device_type="cuda",
+                                    dtype=target_dtype,
+                                    enabled=autocast_enabled):
+                    if (st_attn_available
+                            and self.attn_backend == SlidingTileAttentionBackend
+                        ) or (vsa_available and self.attn_backend
+                              == VideoSparseAttentionBackend):
+                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
+                        )
+
+                        if self.attn_metadata_builder_cls is not None:
+                            self.attn_metadata_builder = self.attn_metadata_builder_cls(
+                            )
+                            # TODO(will): clean this up
+                            attn_metadata = self.attn_metadata_builder.build(  # type: ignore
+                                current_timestep=i,  # type: ignore
+                                raw_latent_shape=batch.
+                                raw_latent_shape[2:5],  # type: ignore
+                                patch_size=fastvideo_args.
+                                pipeline_config.  # type: ignore
+                                dit_config.patch_size,  # type: ignore
+                                STA_param=batch.STA_param,  # type: ignore
+                                VSA_sparsity=fastvideo_args.
+                                VSA_sparsity,  # type: ignore
+                                device=get_local_torch_device(),
+                            )
+                            assert attn_metadata is not None, "attn_metadata cannot be None"
+                        else:
+                            attn_metadata = None
+                    elif (vmoba_attn_available
+                          and self.attn_backend == VMOBAAttentionBackend):
+                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
+                        )
+                        if self.attn_metadata_builder_cls is not None:
+                            self.attn_metadata_builder = self.attn_metadata_builder_cls(
+                            )
+                            # Prepare V-MoBA parameters from config
+                            moba_params = fastvideo_args.moba_config.copy()
+                            moba_params.update({
+                                "current_timestep":
+                                i,
+                                "raw_latent_shape":
+                                batch.raw_latent_shape[2:5],
+                                "patch_size":
+                                fastvideo_args.pipeline_config.dit_config.
+                                patch_size,
+                                "device":
+                                get_local_torch_device(),
+                            })
+                            attn_metadata = self.attn_metadata_builder.build(
+                                **moba_params)
+                            assert attn_metadata is not None, "attn_metadata cannot be None"
+                        else:
+                            attn_metadata = None
+                    else:
+                        attn_metadata = None
+                    # TODO(will): finalize the interface. vLLM uses this to
+                    # support torch dynamo compilation. They pass in
+                    # attn_metadata, vllm_config, and num_tokens. We can pass in
+                    # fastvideo_args or training_args, and attn_metadata.
+                    batch.is_cfg_negative = False
+                    with set_forward_context(
+                            current_timestep=i,
+                            attn_metadata=attn_metadata,
+                            forward_batch=batch,
+                            # fastvideo_args=fastvideo_args
+                    ):
+                        # Run transformer
+                        noise_pred = self.transformer(latent_model_input,
+                                                      prompt_embeds,
+                                                      t_expand,
+                                                      guidance=guidance_expand,
+                                                      timestep_r=timesteps_r,
+                                                      **pos_cond_kwargs,
+                                                      **image_kwargs)
+
+                    # Compute the previous noisy sample
+                    latents = self.scheduler.step(noise_pred,
+                                                  t,
+                                                  latents,
+                                                  return_dict=False)[0]
+
+                # Update progress bar
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and
+                    (i + 1) % self.scheduler.order == 0
+                        and progress_bar is not None):
+                    progress_bar.update()
+
+        # Update batch with final latents
+        batch.latents = latents
+
+        # Save STA mask search results if needed
+        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.STA_mode == STA_Mode.STA_SEARCHING:
+            self.save_sta_search_results(batch)
+
+        # deallocate transformer if on mps
+        if torch.backends.mps.is_available():
+            logger.info("Memory before deallocating transformer: %s",
+                        torch.mps.current_allocated_memory())
+            del self.transformer
+            if pipeline is not None and "transformer" in pipeline.modules:
+                del pipeline.modules["transformer"]
+            fastvideo_args.model_loaded["transformer"] = False
+            logger.info("Memory after deallocating transformer: %s",
+                        torch.mps.current_allocated_memory())
+
+        return batch
+
+    def prepare_extra_func_kwargs(self, func, kwargs) -> dict[str, Any]:
+        """
+        Prepare extra kwargs for the scheduler step / denoise step.
+        
+        Args:
+            func: The function to prepare kwargs for.
+            kwargs: The kwargs to prepare.
+            
+        Returns:
+            The prepared kwargs.
+        """
+        extra_step_kwargs = {}
+        for k, v in kwargs.items():
+            accepts = k in set(inspect.signature(func).parameters.keys())
+            if accepts:
+                extra_step_kwargs[k] = v
+        return extra_step_kwargs
+
+    def progress_bar(self,
+                     iterable: Iterable | None = None,
+                     total: int | None = None) -> tqdm:
+        """
+        Create a progress bar for the denoising process.
+        
+        Args:
+            iterable: The iterable to iterate over.
+            total: The total number of items.
+            
+        Returns:
+            A tqdm progress bar.
+        """
+        local_rank = get_world_group().local_rank
+        if local_rank == 0:
+            return tqdm(iterable=iterable, total=total)
+        else:
+            return tqdm(iterable=iterable, total=total, disable=True)
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify denoising stage inputs."""
+        result = VerificationResult()
+        result.add_check("timesteps", batch.timesteps,
+                         [V.is_tensor, V.min_dims(1)])
+        result.add_check("latents", batch.latents,
+                         [V.is_tensor, V.with_dims(5)])
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
+        result.add_check("image_embeds", batch.image_embeds, V.is_list)
+        result.add_check("image_latent", batch.image_latent,
+                         V.none_or_tensor_with_dims(5))
+        result.add_check("num_inference_steps", batch.num_inference_steps,
+                         V.positive_int)
+        result.add_check("guidance_scale", batch.guidance_scale,
+                         V.positive_float)
+        result.add_check("eta", batch.eta, V.non_negative_float)
+        result.add_check("generator", batch.generator,
+                         V.generator_or_list_generators)
+        result.add_check("do_classifier_free_guidance",
+                         batch.do_classifier_free_guidance, V.bool_value)
+        result.add_check(
+            "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x:
+            not batch.do_classifier_free_guidance or V.list_not_empty(x))
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify denoising stage outputs."""
+        result = VerificationResult()
+        result.add_check("latents", batch.latents,
+                         [V.is_tensor, V.with_dims(5)])
+        return result

fastvideo/registry.py

@@ -18,8 +18,9 @@ from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.configs.pipelines.cosmos import CosmosConfig
 from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
-from fastvideo.configs.pipelines.hunyuan15 import (Hunyuan15T2V480PConfig,
-                                                   Hunyuan15T2V720PConfig)
+from fastvideo.configs.pipelines.hunyuan15 import (
+    Hunyuan15T2V480PConfig, Hunyuan15I2V480PStepDistilledConfig,
+    Hunyuan15T2V720PConfig, Hunyuan15I2V720PConfig, Hunyuan15SR1080PConfig)
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
 from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
 from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig
@@ -50,8 +51,11 @@ from fastvideo.configs.sample.cosmos import (
 from fastvideo.configs.sample.cosmos2_5 import Cosmos25SamplingParamBase
 from fastvideo.configs.sample.hunyuan import (FastHunyuanSamplingParam,
                                               HunyuanSamplingParam)
-from fastvideo.configs.sample.hunyuan15 import (Hunyuan15_480P_SamplingParam,
-                                                Hunyuan15_720P_SamplingParam)
+from fastvideo.configs.sample.hunyuan15 import (
+    Hunyuan15_480P_SamplingParam,
+    Hunyuan15_480P_StepDistilled_I2V_SamplingParam,
+    Hunyuan15_720P_SamplingParam, Hunyuan15_720P_Distilled_I2V_SamplingParam,
+    Hunyuan15_SR_1080P_SamplingParam)
 from fastvideo.configs.sample.hyworld import HYWorld_SamplingParam
 from fastvideo.configs.sample.ltx2 import LTX2SamplingParam
 from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam
@@ -265,6 +269,13 @@ def _register_configs() -> None:
             )),
         ],
     )
+    register_configs(
+        sampling_param_cls=Hunyuan15_480P_StepDistilled_I2V_SamplingParam,
+        pipeline_config_cls=Hunyuan15I2V480PStepDistilledConfig,
+        hf_model_paths=[
+            "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled",
+        ],
+    )
     register_configs(
         sampling_param_cls=Hunyuan15_720P_SamplingParam,
         pipeline_config_cls=Hunyuan15T2V720PConfig,
@@ -272,6 +283,21 @@ def _register_configs() -> None:
             "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v",
         ],
     )
+    register_configs(
+        sampling_param_cls=Hunyuan15_720P_Distilled_I2V_SamplingParam,
+        pipeline_config_cls=Hunyuan15I2V720PConfig,
+        hf_model_paths=[
+            "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v_distilled",
+        ],
+    )
+    register_configs(
+        sampling_param_cls=Hunyuan15_SR_1080P_SamplingParam,
+        pipeline_config_cls=Hunyuan15SR1080PConfig,
+        hf_model_paths=[
+            "weizhou03/HunyuanVideo-1.5-Diffusers-1080p",
+            "weizhou03/HunyuanVideo-1.5-Diffusers-1080p-2SR"
+        ],
+    )
 
     # Hunyuan
     register_configs(