From 98308db7e0a992cc5101fe5e0f59f989643e609d Mon Sep 17 00:00:00 2001 From: Wei Zhou <69577934+JerryZhou54@users.noreply.github.com> Date: Thu, 5 Feb 2026 19:39:20 -0500 Subject: [PATCH] [Feature] [Hy1.5] Support HY1.5 super-resolution pipeline for 1080p videos (#1046) --- examples/inference/basic/basic_hy15.py | 2 +- examples/inference/basic/basic_hy15_1080p.py | 43 ++ fastvideo/configs/models/__init__.py | 2 + .../configs/models/dits/hunyuanvideo15.py | 2 + .../configs/models/upsamplers/__init__.py | 6 + fastvideo/configs/models/upsamplers/base.py | 7 + .../configs/models/upsamplers/hunyuan15.py | 20 + fastvideo/configs/pipelines/base.py | 5 +- fastvideo/configs/pipelines/hunyuan15.py | 29 +- fastvideo/configs/sample/base.py | 3 + fastvideo/configs/sample/hunyuan15.py | 26 ++ fastvideo/fastvideo_args.py | 1 + fastvideo/models/loader/component_loader.py | 53 ++- fastvideo/models/registry.py | 6 + fastvideo/models/upsamplers/hunyuan15.py | 170 +++++++ .../basic/hunyuan15/hunyuan15_2sr_pipeline.py | 162 +++++++ .../basic/hunyuan15/hunyuan15_i2v_pipeline.py | 74 +++ .../basic/hunyuan15/hunyuan15_sr_pipeline.py | 136 ++++++ fastvideo/pipelines/composed_pipeline_base.py | 1 + fastvideo/pipelines/pipeline_batch_info.py | 4 + fastvideo/pipelines/stages/__init__.py | 2 + fastvideo/pipelines/stages/denoising.py | 22 + fastvideo/pipelines/stages/sr_denoising.py | 431 ++++++++++++++++++ fastvideo/registry.py | 34 +- 24 files changed, 1232 insertions(+), 9 deletions(-) create mode 100644 examples/inference/basic/basic_hy15_1080p.py create mode 100644 fastvideo/configs/models/upsamplers/__init__.py create mode 100644 fastvideo/configs/models/upsamplers/base.py create mode 100644 fastvideo/configs/models/upsamplers/hunyuan15.py create mode 100644 fastvideo/models/upsamplers/hunyuan15.py create mode 100644 fastvideo/pipelines/basic/hunyuan15/hunyuan15_2sr_pipeline.py create mode 100644 fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py create mode 100644 fastvideo/pipelines/basic/hunyuan15/hunyuan15_sr_pipeline.py create mode 100644 fastvideo/pipelines/stages/sr_denoising.py diff --git a/examples/inference/basic/basic_hy15.py b/examples/inference/basic/basic_hy15.py index 39d8772d3..e0b782237 100644 --- a/examples/inference/basic/basic_hy15.py +++ b/examples/inference/basic/basic_hy15.py @@ -39,4 +39,4 @@ def main(): if __name__ == "__main__": - main() + main() \ No newline at end of file diff --git a/examples/inference/basic/basic_hy15_1080p.py b/examples/inference/basic/basic_hy15_1080p.py new file mode 100644 index 000000000..6b346e8f9 --- /dev/null +++ b/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() diff --git a/fastvideo/configs/models/__init__.py b/fastvideo/configs/models/__init__.py index aa05ba62e..58b38d40c 100644 --- a/fastvideo/configs/models/__init__.py +++ b/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", ] diff --git a/fastvideo/configs/models/dits/hunyuanvideo15.py b/fastvideo/configs/models/dits/hunyuanvideo15.py index 6d7f77208..bc62df63f 100644 --- a/fastvideo/configs/models/dits/hunyuanvideo15.py +++ b/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\.(.*)$": diff --git a/fastvideo/configs/models/upsamplers/__init__.py b/fastvideo/configs/models/upsamplers/__init__.py new file mode 100644 index 000000000..b5b8332cb --- /dev/null +++ b/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" +] diff --git a/fastvideo/configs/models/upsamplers/base.py b/fastvideo/configs/models/upsamplers/base.py new file mode 100644 index 000000000..5430912bb --- /dev/null +++ b/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 diff --git a/fastvideo/configs/models/upsamplers/hunyuan15.py b/fastvideo/configs/models/upsamplers/hunyuan15.py new file mode 100644 index 000000000..dc16459c7 --- /dev/null +++ b/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 diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index 92482e238..a057628a3 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/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) diff --git a/fastvideo/configs/pipelines/hunyuan15.py b/fastvideo/configs/pipelines/hunyuan15.py index 29e1c5b45..2b9e76c50 100644 --- a/fastvideo/configs/pipelines/hunyuan15.py +++ b/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" diff --git a/fastvideo/configs/sample/base.py b/fastvideo/configs/sample/base.py index 8018323be..ab3131419 100644 --- a/fastvideo/configs/sample/base.py +++ b/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 diff --git a/fastvideo/configs/sample/hunyuan15.py b/fastvideo/configs/sample/hunyuan15.py index 941b9ea22..d50adb186 100644 --- a/fastvideo/configs/sample/hunyuan15.py +++ b/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 diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 33c06efcc..96fac3913 100644 --- a/fastvideo/fastvideo_args.py +++ b/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 diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index bf53d324c..8437e9821 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/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.""" diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index de60167bd..350618805 100644 --- a/fastvideo/models/registry.py +++ b/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__) diff --git a/fastvideo/models/upsamplers/hunyuan15.py b/fastvideo/models/upsamplers/hunyuan15.py new file mode 100644 index 000000000..0d31a517e --- /dev/null +++ b/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 \ No newline at end of file diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_2sr_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_2sr_pipeline.py new file mode 100644 index 000000000..fe7a5a106 --- /dev/null +++ b/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 diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py new file mode 100644 index 000000000..297a2637f --- /dev/null +++ b/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 diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_sr_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_sr_pipeline.py new file mode 100644 index 000000000..fd3bf87fc --- /dev/null +++ b/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 diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index f8d99ce10..d9a7f6419 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/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: diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 53e838acb..9a8f0471c 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/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 diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index b4d99c39e..f085d6749 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/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", diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index c09254365..c037cbad5 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/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 diff --git a/fastvideo/pipelines/stages/sr_denoising.py b/fastvideo/pipelines/stages/sr_denoising.py new file mode 100644 index 000000000..6311930bd --- /dev/null +++ b/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 diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 6d134b782..f6b7dc21c 100644 --- a/fastvideo/registry.py +++ b/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(