From dbddfab6d2bb294d4de9d365743c4115d17d89ff Mon Sep 17 00:00:00 2001 From: KyleShao <71746088+KyleShao1016@users.noreply.github.com> Date: Thu, 15 Jan 2026 15:09:05 -0800 Subject: [PATCH] [feat] Introduce Cosmos 2.5 Text2World pipeline (#974) --- .../inference/basic/basic_cosmos2_5_t2w.py | 42 + fastvideo/configs/models/encoders/__init__.py | 3 +- fastvideo/configs/models/encoders/reason1.py | 72 + fastvideo/configs/models/vaes/__init__.py | 2 + fastvideo/configs/models/vaes/cosmos2_5vae.py | 223 ++ fastvideo/configs/pipelines/__init__.py | 3 +- fastvideo/configs/pipelines/cosmos2_5.py | 89 + fastvideo/configs/pipelines/registry.py | 8 +- fastvideo/configs/sample/cosmos2_5.py | 19 + fastvideo/configs/sample/registry.py | 14 +- .../models/encoders/qwen2_5_vl_custom.py | 2105 +++++++++++++++++ fastvideo/models/encoders/reason1.py | 353 +++ fastvideo/models/loader/component_loader.py | 88 +- fastvideo/models/loader/fsdp_load.py | 23 +- fastvideo/models/registry.py | 10 +- .../scheduling_flow_unipc_multistep.py | 69 +- fastvideo/models/vaes/cosmos25wanvae.py | 735 ++++++ .../basic/cosmos/cosmos2_5_pipeline.py | 61 + fastvideo/pipelines/pipeline_registry.py | 1 + fastvideo/pipelines/stages/__init__.py | 15 +- fastvideo/pipelines/stages/decoding.py | 41 +- fastvideo/pipelines/stages/denoising.py | 139 ++ .../pipelines/stages/latent_preparation.py | 205 ++ fastvideo/pipelines/stages/text_encoding.py | 66 + .../pipelines/stages/timestep_preparation.py | 29 + .../pt_to_safetensors.py | 106 + 26 files changed, 4459 insertions(+), 62 deletions(-) create mode 100644 examples/inference/basic/basic_cosmos2_5_t2w.py create mode 100644 fastvideo/configs/models/encoders/reason1.py create mode 100644 fastvideo/configs/models/vaes/cosmos2_5vae.py create mode 100644 fastvideo/configs/pipelines/cosmos2_5.py create mode 100644 fastvideo/configs/sample/cosmos2_5.py create mode 100644 fastvideo/models/encoders/qwen2_5_vl_custom.py create mode 100644 fastvideo/models/encoders/reason1.py create mode 100644 fastvideo/models/vaes/cosmos25wanvae.py create mode 100644 fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py create mode 100644 scripts/checkpoint_conversion/pt_to_safetensors.py diff --git a/examples/inference/basic/basic_cosmos2_5_t2w.py b/examples/inference/basic/basic_cosmos2_5_t2w.py new file mode 100644 index 000000000..d2d11a0ef --- /dev/null +++ b/examples/inference/basic/basic_cosmos2_5_t2w.py @@ -0,0 +1,42 @@ +from fastvideo import VideoGenerator + + +def main(): + # Point this to your local diffusers model dir (or replace with a HF model ID). + model_path = "KyleShao/Cosmos-Predict2.5-2B-Diffusers" + + generator = VideoGenerator.from_pretrained( + model_path, + num_gpus=1, + use_fsdp_inference=False, # set True if GPU is out of memory + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=True, + pin_cpu_memory=True, + ) + + prompt = ( + "A high-definition video captures the precision of robotic welding in an industrial setting. The first frame showcases a robotic arm, equipped with a welding torch, positioned over a large metal structure. The welding process is in full swing, with bright sparks and intense light illuminating the scene, creating a vivid display of blue and white hues. A significant amount of smoke billows around the welding area, partially obscuring the view but emphasizing the heat and activity. The background reveals parts of the workshop environment, including a ventilation system and various pieces of machinery, indicating a busy and functional industrial workspace. As the video progresses, the robotic arm maintains its steady position, continuing the welding process and moving to its left. The welding torch consistently emits sparks and light, and the smoke continues to rise, diffusing slightly as it moves upward. The metal surface beneath the torch shows ongoing signs of heating and melting. The scene retains its industrial ambiance, with the welding sparks and smoke dominating the visual field, underscoring the ongoing nature of the welding operation." + ) + + video = generator.generate_video( + prompt, + negative_prompt="", + height=704, + width=1280, + num_frames=77, + num_inference_steps=35, + guidance_scale=7.0, + fps=24, + output_path="outputs_video/cosmos2_5_t2w.mp4", + save_video=True, + ) + + generator.shutdown() + + +if __name__ == "__main__": + main() + + + diff --git a/fastvideo/configs/models/encoders/__init__.py b/fastvideo/configs/models/encoders/__init__.py index 86a3e3334..f943114e8 100644 --- a/fastvideo/configs/models/encoders/__init__.py +++ b/fastvideo/configs/models/encoders/__init__.py @@ -7,10 +7,11 @@ from fastvideo.configs.models.encoders.clip import ( from fastvideo.configs.models.encoders.llama import LlamaConfig from fastvideo.configs.models.encoders.t5 import T5Config, T5LargeConfig from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig +from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1Config __all__ = [ "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", - "Qwen2_5_VLConfig" + "Qwen2_5_VLConfig", "Reason1ArchConfig", "Reason1Config" ] diff --git a/fastvideo/configs/models/encoders/reason1.py b/fastvideo/configs/models/encoders/reason1.py new file mode 100644 index 000000000..825f69f07 --- /dev/null +++ b/fastvideo/configs/models/encoders/reason1.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Config for Reason1 (Qwen2.5-VL) text encoder.""" + +from dataclasses import dataclass, field +from typing import Any + +from fastvideo.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig + + +@dataclass +class Reason1ArchConfig(TextEncoderArchConfig): + """Architecture settings (defaults match Qwen2.5-VL-7B-Instruct).""" + + architectures: list[str] = field( + default_factory=lambda: ["Qwen2_5_VLForConditionalGeneration"]) + model_type: str = "qwen2_5_vl" + + vocab_size: int = 152064 + hidden_size: int = 3584 + num_hidden_layers: int = 28 + num_attention_heads: int = 28 + num_key_value_heads: int = 4 + intermediate_size: int = 18944 + + text_len: int = 512 + hidden_state_skip_layer: int = 0 + bos_token_id: int = 151643 + pad_token_id: int = 151643 + eos_token_id: int = 151645 + + image_token_id: int = 151655 + video_token_id: int = 151656 + vision_token_id: int = 151654 + vision_start_token_id: int = 151652 + vision_end_token_id: int = 151653 + + vision_config: dict[str, Any] | None = None + + rope_theta: float = 1000000.0 + rope_scaling: dict[str, Any] | None = field(default_factory=lambda: { + "type": "mrope", + "mrope_section": [16, 24, 24] + }) + max_position_embeddings: int = 128000 + max_window_layers: int = 28 + + embedding_concat_strategy: str = "mean_pooling" + n_layers_per_group: int = 5 + num_embedding_padding_tokens: int = 512 + + attention_dropout: float = 0.0 + hidden_act: str = "silu" + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-6 + + use_sliding_window: bool = False + sliding_window: int = 32768 + + tie_word_embeddings: bool = False + use_cache: bool = False + output_hidden_states: bool = True + + torch_dtype: str = "bfloat16" + _attn_implementation: str = "flash_attention_2" + + +@dataclass +class Reason1Config(TextEncoderConfig): + """Reason1 text encoder config.""" + + arch_config: Reason1ArchConfig = field(default_factory=Reason1ArchConfig) + tokenizer_type: str = "Qwen/Qwen2.5-VL-7B-Instruct" diff --git a/fastvideo/configs/models/vaes/__init__.py b/fastvideo/configs/models/vaes/__init__.py index d39254d21..ad1b2f206 100644 --- a/fastvideo/configs/models/vaes/__init__.py +++ b/fastvideo/configs/models/vaes/__init__.py @@ -1,4 +1,5 @@ from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig +from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig from fastvideo.configs.models.vaes.stepvideovae import StepVideoVAEConfig @@ -9,5 +10,6 @@ __all__ = [ "WanVAEConfig", "StepVideoVAEConfig", "CosmosVAEConfig", + "Cosmos25VAEConfig", "Hunyuan15VAEConfig", ] diff --git a/fastvideo/configs/models/vaes/cosmos2_5vae.py b/fastvideo/configs/models/vaes/cosmos2_5vae.py new file mode 100644 index 000000000..1b48cbc1a --- /dev/null +++ b/fastvideo/configs/models/vaes/cosmos2_5vae.py @@ -0,0 +1,223 @@ +"""Cosmos 2.5 (Wan2.1-style) VAE config and checkpoint-key mapping.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field + +import torch + +from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class Cosmos25VAEArchConfig(VAEArchConfig): + _name_or_path: str = "" + base_dim: int = 96 + decoder_base_dim: int | None = None + z_dim: int = 16 + dim_mult: tuple[int, ...] = (1, 2, 4, 4) + num_res_blocks: int = 2 + attn_scales: tuple[float, ...] = () + temperal_downsample: tuple[bool, ...] = (False, True, True) + dropout: float = 0.0 + is_residual: bool = False + in_channels: int = 3 + out_channels: int = 3 + patch_size: int | None = None + scale_factor_temporal: int = 4 + scale_factor_spatial: int = 8 + clip_output: bool = True + + latents_mean: tuple[float, ...] = ( + -0.7571, + -0.7089, + -0.9113, + 0.1075, + -0.1745, + 0.9653, + -0.1517, + 1.5508, + 0.4134, + -0.0715, + 0.5517, + -0.3632, + -0.1922, + -0.9497, + 0.2503, + -0.2921, + ) + latents_std: tuple[float, ...] = ( + 2.8184, + 1.4541, + 2.3275, + 2.6558, + 1.2196, + 1.7708, + 2.6052, + 2.0743, + 3.2687, + 2.1526, + 2.8652, + 1.5579, + 1.6382, + 1.1253, + 2.8251, + 1.9160, + ) + + # Simple 1:1 renames. More complex decoder remapping is handled by + # `map_official_key()`. + param_names_mapping: dict[str, str] = field( + default_factory=lambda: { + r"^conv1\.(.*)$": r"quant_conv.\1", + r"^conv2\.(.*)$": r"post_quant_conv.\1", + r"^encoder\.conv1\.(.*)$": r"encoder.conv_in.\1", + r"^decoder\.conv1\.(.*)$": r"decoder.conv_in.\1", + r"^encoder\.head\.0\.gamma$": r"encoder.norm_out.gamma", + r"^encoder\.head\.2\.(.*)$": r"encoder.conv_out.\1", + r"^decoder\.head\.0\.gamma$": r"decoder.norm_out.gamma", + r"^decoder\.head\.2\.(.*)$": r"decoder.conv_out.\1", + }) + + @staticmethod + def map_official_key(key: str) -> str | None: + """Map a single official checkpoint key into FastVideo key space.""" + + def map_residual_subkey(prefix: str, sub: str) -> str | None: + if re.match(r"^residual\.0\.gamma$", sub): + return f"{prefix}.norm1.gamma" + m = re.match(r"^residual\.2\.(weight|bias)$", sub) + if m: + return f"{prefix}.conv1.{m.group(1)}" + if re.match(r"^residual\.3\.gamma$", sub): + return f"{prefix}.norm2.gamma" + m = re.match(r"^residual\.6\.(weight|bias)$", sub) + if m: + return f"{prefix}.conv2.{m.group(1)}" + m = re.match(r"^shortcut\.(weight|bias)$", sub) + if m: + return f"{prefix}.conv_shortcut.{m.group(1)}" + return None + + def map_attn_subkey(prefix: str, sub: str) -> str | None: + if re.match(r"^norm\.gamma$", sub): + return f"{prefix}.norm.gamma" + m = re.match(r"^to_qkv\.(weight|bias)$", sub) + if m: + return f"{prefix}.to_qkv.{m.group(1)}" + m = re.match(r"^proj\.(weight|bias)$", sub) + if m: + return f"{prefix}.proj.{m.group(1)}" + return None + + def map_resample_subkey(prefix: str, sub: str) -> str | None: + m = re.match(r"^resample\.1\.(weight|bias)$", sub) + if m: + return f"{prefix}.resample.1.{m.group(1)}" + m = re.match(r"^time_conv\.(weight|bias)$", sub) + if m: + return f"{prefix}.time_conv.{m.group(1)}" + return None + + m = re.match(r"^conv1\.(weight|bias)$", key) + if m: + return f"quant_conv.{m.group(1)}" + m = re.match(r"^conv2\.(weight|bias)$", key) + if m: + return f"post_quant_conv.{m.group(1)}" + m = re.match(r"^(encoder|decoder)\.conv1\.(weight|bias)$", key) + if m: + return f"{m.group(1)}.conv_in.{m.group(2)}" + m = re.match(r"^(encoder|decoder)\.head\.0\.gamma$", key) + if m: + return f"{m.group(1)}.norm_out.gamma" + m = re.match(r"^(encoder|decoder)\.head\.2\.(weight|bias)$", key) + if m: + return f"{m.group(1)}.conv_out.{m.group(2)}" + + m = re.match(r"^(encoder|decoder)\.middle\.0\.(.*)$", key) + if m: + return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.0", + m.group(2)) + m = re.match(r"^(encoder|decoder)\.middle\.1\.(.*)$", key) + if m: + return map_attn_subkey(f"{m.group(1)}.mid_block.attentions.0", + m.group(2)) + m = re.match(r"^(encoder|decoder)\.middle\.2\.(.*)$", key) + if m: + return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.1", + m.group(2)) + + m = re.match(r"^encoder\.downsamples\.(\d+)\.(.*)$", key) + if m: + idx = int(m.group(1)) + sub = m.group(2) + if sub.startswith("residual.") or sub.startswith("shortcut."): + return map_residual_subkey(f"encoder.down_blocks.{idx}", sub) + if sub.startswith("resample.") or sub.startswith("time_conv."): + return map_resample_subkey(f"encoder.down_blocks.{idx}", sub) + return None + + m = re.match(r"^decoder\.upsamples\.(\d+)\.(.*)$", key) + if m: + uidx = int(m.group(1)) + sub = m.group(2) + + if uidx in (0, 1, 2): + block_i, res_i = 0, uidx + elif uidx == 3: + block_i, res_i = 0, None + elif uidx in (4, 5, 6): + block_i, res_i = 1, uidx - 4 + elif uidx == 7: + block_i, res_i = 1, None + elif uidx in (8, 9, 10): + block_i, res_i = 2, uidx - 8 + elif uidx == 11: + block_i, res_i = 2, None + elif uidx in (12, 13, 14): + block_i, res_i = 3, uidx - 12 + else: + return None + + if res_i is None: + return map_resample_subkey( + f"decoder.up_blocks.{block_i}.upsamplers.0", + sub, + ) + + return map_residual_subkey( + f"decoder.up_blocks.{block_i}.resnets.{res_i}", + sub, + ) + + return None + + temporal_compression_ratio: int = 4 + spatial_compression_ratio: int = 8 + + def __post_init__(self): + self.scaling_factor: torch.Tensor = 1.0 / torch.tensor( + self.latents_std).view(1, self.z_dim, 1, 1, 1) + self.shift_factor: torch.Tensor = torch.tensor(self.latents_mean).view( + 1, self.z_dim, 1, 1, 1) + self.temporal_compression_ratio = self.scale_factor_temporal + self.spatial_compression_ratio = self.scale_factor_spatial + + +@dataclass +class Cosmos25VAEConfig(VAEConfig): + """Cosmos2.5 VAE config.""" + + arch_config: Cosmos25VAEArchConfig = field( + default_factory=Cosmos25VAEArchConfig) + + use_feature_cache: bool = True + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False + + def __post_init__(self): + self.blend_num_frames = (self.tile_sample_min_num_frames - + self.tile_sample_stride_num_frames) * 2 diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index e28901889..8120ed53e 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -1,6 +1,7 @@ from fastvideo.configs.pipelines.base import (PipelineConfig, SlidingTileAttnConfig) 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.registry import ( @@ -15,5 +16,5 @@ __all__ = [ "Hunyuan15T2V480PConfig", "Hunyuan15T2V720PConfig", "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig", "StepVideoT2VConfig", "SelfForcingWanT2V480PConfig", - "CosmosConfig", "get_pipeline_config_cls_from_name" + "CosmosConfig", "Cosmos25Config", "get_pipeline_config_cls_from_name" ] diff --git a/fastvideo/configs/pipelines/cosmos2_5.py b/fastvideo/configs/pipelines/cosmos2_5.py new file mode 100644 index 000000000..3ab01e161 --- /dev/null +++ b/fastvideo/configs/pipelines/cosmos2_5.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +from collections.abc import Callable +from dataclasses import dataclass, field + +import torch + +from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig +from fastvideo.configs.models.dits import Cosmos25VideoConfig +from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25ArchConfig +from fastvideo.configs.models.encoders import BaseEncoderOutput +from fastvideo.configs.models.encoders.reason1 import Reason1Config, Reason1ArchConfig +from fastvideo.configs.models.vaes import Cosmos25VAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig, STA_Mode + + +def _identity_preprocess_text(prompt: str) -> str: + return prompt + + +def reason1_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor: + hidden_states = getattr(outputs, "hidden_states", None) + if hidden_states is None: + raise ValueError("Reason1 postprocess requires outputs.hidden_states") + + hs = list(hidden_states)[1:] + normed = [] + for h in hs: + h = h.float() + h = (h - h.mean(dim=-1, keepdim=True)) / (h.std(dim=-1, keepdim=True) + + 1e-8) + normed.append(h) + return torch.cat(normed, dim=-1).to(hidden_states[0].dtype) + + +@dataclass +class Cosmos25Config(PipelineConfig): + """Configuration for Cosmos 2.5 (Predict2.5) video generation pipeline.""" + + dit_config: DiTConfig = field(default_factory=lambda: Cosmos25VideoConfig( + arch_config=Cosmos25ArchConfig( + num_attention_heads=16, + attention_head_dim=128, + in_channels=16, + out_channels=16, + num_layers=28, + patch_size=[1, 2, 2], + max_size=[128, 240, 240], + rope_scale=[1.0, 3.0, 3.0], + text_embed_dim=1024, + mlp_ratio=4.0, + adaln_lora_dim=256, + use_adaln_lora=True, + concat_padding_mask=True, + extra_pos_embed_type=None, + use_crossattn_projection=True, + rope_enable_fps_modulation=False, + qk_norm="rms_norm", + ))) + + vae_config: VAEConfig = field(default_factory=Cosmos25VAEConfig) + + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (Reason1Config(arch_config=Reason1ArchConfig( + embedding_concat_strategy="full_concat")), )) + + preprocess_text_funcs: tuple[Callable[[str], str], ...] = field( + default_factory=lambda: (_identity_preprocess_text, )) + postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], + ...] = field(default_factory=lambda: + (reason1_postprocess_text, )) + + dit_precision: str = "bf16" + vae_precision: str = "bf16" + text_encoder_precisions: tuple[str, ...] = field( + default_factory=lambda: ("bf16", )) + + embedded_cfg_scale: float = 0.0 + flow_shift: float = 5.0 + + vae_tiling: bool = False + vae_sp: bool = False + + STA_mode: STA_Mode = STA_Mode.NONE + skip_time_steps: int = 0 + + def __post_init__(self): + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + self._vae_latent_dim = 16 diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index 063b72c4d..0b4c148df 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -6,6 +6,7 @@ from collections.abc import Callable 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.stepvideo import StepVideoT2VConfig @@ -55,6 +56,7 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = { "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config, "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config, "nvidia/Cosmos-Predict2-2B-Video2World": CosmosConfig, + "KyleShao/Cosmos-Predict2.5-2B-Diffusers": Cosmos25Config, "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGameI2V480PConfig, "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGameI2V480PConfig, "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers": MatrixGameI2V480PConfig, @@ -94,7 +96,10 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = { "stepvideo": lambda id: "stepvideo" in id.lower(), "cosmos": - lambda id: "cosmos" in id.lower(), + lambda id: "cosmos" in id.lower() and ("2.5" not in id.lower( + ) and "2_5" not in id.lower() and "25" not in id.lower()), + "cosmos25": + lambda id: "cosmos25" in id.lower(), "turbodiffusion": lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(), # Add other pipeline architecture detectors @@ -105,6 +110,7 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = { "longcatimagetovideo": LongCatT2V480PConfig, "longcatvideocontinuation": LongCatT2V480PConfig, "longcat": LongCatT2V480PConfig, + "cosmos25": Cosmos25Config, "hunyuan": HunyuanConfig, # Base Hunyuan config as fallback for any Hunyuan variant "matrixgame": MatrixGameI2V480PConfig, diff --git a/fastvideo/configs/sample/cosmos2_5.py b/fastvideo/configs/sample/cosmos2_5.py new file mode 100644 index 000000000..48287b303 --- /dev/null +++ b/fastvideo/configs/sample/cosmos2_5.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass + +from fastvideo.configs.sample.base import SamplingParam + + +@dataclass +class Cosmos_Predict2_5_2B_Diffusers_SamplingParam(SamplingParam): + """Defaults for Cosmos 2.5 (Predict2.5) text-to-video diffusers-format model.""" + + height: int = 480 + width: int = 832 + num_frames: int = 121 + fps: int = 24 + + guidance_scale: float = 7.0 + # Official Cosmos2.5 sampling uses empty string as unconditional. + negative_prompt: str = "" + num_inference_steps: int = 35 diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py index 07eedcbb4..89673bed8 100644 --- a/fastvideo/configs/sample/registry.py +++ b/fastvideo/configs/sample/registry.py @@ -9,6 +9,7 @@ from fastvideo.configs.sample.hunyuan15 import Hunyuan15_480P_SamplingParam, Hun from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam from fastvideo.configs.sample.cosmos import Cosmos_Predict2_2B_Video2World_SamplingParam +from fastvideo.configs.sample.cosmos2_5 import Cosmos_Predict2_5_2B_Diffusers_SamplingParam # isort: off from fastvideo.configs.sample.wan import ( @@ -96,6 +97,10 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = { "nvidia/Cosmos-Predict2-2B-Video2World": Cosmos_Predict2_2B_Video2World_SamplingParam, + # Cosmos2.5 + "KyleShao/Cosmos-Predict2.5-2B-Diffusers": + Cosmos_Predict2_5_2B_Diffusers_SamplingParam, + # MatrixGame2.0 models "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGame2_SamplingParam, @@ -135,6 +140,10 @@ SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = { lambda id: "matrixgame" in id.lower() or "matrix-game" in id.lower(), "turbodiffusion": lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(), + "cosmos25": + lambda id: "cosmos2_5" in id.lower(), + "cosmos": + lambda id: "cosmos" in id.lower() and "2_5" not in id.lower(), # Add other pipeline architecture detectors } @@ -153,6 +162,8 @@ SAMPLING_FALLBACK_PARAM: dict[str, Any] = { "matrixgame": MatrixGame2_SamplingParam, "turbodiffusion": TurboDiffusionT2V_1_3B_SamplingParam, # Default to T2V for fallback + "cosmos25": Cosmos_Predict2_5_2B_Diffusers_SamplingParam, + "cosmos": Cosmos_Predict2_2B_Video2World_SamplingParam, # Other fallbacks by architecture } @@ -176,9 +187,6 @@ def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None: if os.path.exists(pipeline_name_or_path): config = verify_model_config_and_directory(pipeline_name_or_path) - logger.warning( - "FastVideo may not correctly identify the optimal sampling param for this model, as the local directory may have been renamed." - ) else: config = maybe_download_model_index(pipeline_name_or_path) diff --git a/fastvideo/models/encoders/qwen2_5_vl_custom.py b/fastvideo/models/encoders/qwen2_5_vl_custom.py new file mode 100644 index 000000000..78a0e7a19 --- /dev/null +++ b/fastvideo/models/encoders/qwen2_5_vl_custom.py @@ -0,0 +1,2105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +import warnings +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import omegaconf +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed._tensor import DTensor +from torch.distributed.device_mesh import DeviceMesh +from torch.nn import CrossEntropyLoss +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache +from transformers.modeling_attn_mask_utils import AttentionMaskConverter +from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + +try: + from torch.distributed.tensor import Shard +except ImportError: + Shard = None # type: ignore[assignment] + warnings.warn("torch.distributed.tensor.Shard is not available; context-parallel features may be disabled.") + +try: + from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig +except ImportError: + warnings.warn("transformers is too old; Qwen2.5-VL config classes are not available.") + Qwen2_5_VLConfig = dict + Qwen2_5_VLVisionConfig = dict + +from transformers.utils import ( + is_flash_attn_2_available, + is_flash_attn_greater_or_equal_2_10, + is_torchdynamo_compiling, + logging, +) + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_varlen_func + from flash_attn.layers.rotary import apply_rotary_emb + +else: + flash_attn_varlen_func = None + apply_rotary_emb = None + + +if is_flash_attn_2_available(): + from transformers.modeling_flash_attention_utils import _flash_attention_forward +else: + _flash_attention_forward = None + flash_attn_varlen_func = None + +if not is_flash_attn_2_available(): + warnings.warn( + "flash_attn_2 not available. Falling back to SDPA attention for Qwen2.5-VL. " + "Install flash-attn for better performance.", + RuntimeWarning, + ) + +logger = logging.get_logger(__name__) + +class Qwen2_5_VLMLP(nn.Module): + def __init__(self, config, bias: bool = False): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_state): + return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) + + +class Qwen2_5_VisionPatchEmbed(nn.Module): + def __init__( + self, + patch_size: int = 14, + temporal_patch_size: int = 2, + in_channels: int = 3, + embed_dim: int = 1152, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.temporal_patch_size = temporal_patch_size + self.in_channels = in_channels + self.embed_dim = embed_dim + + kernel_size = [temporal_patch_size, patch_size, patch_size] + self.proj = nn.Conv3d( + in_channels, + embed_dim, + kernel_size=kernel_size, + stride=kernel_size, + bias=False, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + target_dtype = self.proj.weight.dtype + hidden_states = hidden_states.view( + -1, + self.in_channels, + self.temporal_patch_size, + self.patch_size, + self.patch_size, + ) + hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim) + return hidden_states + + +class Qwen2_5_VisionRotaryEmbedding(nn.Module): + def __init__(self, dim: int, theta: float = 10000.0) -> None: + super().__init__() + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.dim = dim + + def init_weights(self, buffer_device: Optional[torch.device] = None): + if buffer_device is None: + device = self.inv_freq.device + else: + device = buffer_device + self.inv_freq = 1.0 / (10000.0 ** (torch.arange(0, self.dim, 2, dtype=torch.float) / self.dim)).to(device) + + def forward(self, seqlen: int) -> torch.Tensor: + seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype) + freqs = torch.outer(seq, self.inv_freq) + return freqs + + +class Qwen2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + Qwen2RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class Qwen2_5_VLPatchMerger(nn.Module): + def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None: + super().__init__() + self.hidden_size = context_dim * (spatial_merge_size**2) + self.ln_q = Qwen2RMSNorm(context_dim, eps=1e-6) + self.mlp = nn.Sequential( + nn.Linear(self.hidden_size, self.hidden_size), + nn.GELU(), + nn.Linear(self.hidden_size, dim), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.mlp(self.ln_q(x).view(-1, self.hidden_size)) + return x + + def init_weights(self, buffer_device: Optional[torch.device] = None): + pass + + +def apply_rotary_pos_emb_flashatt(tensor: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + tensor_ = tensor.float() + cos = freqs.cos().float() + sin = freqs.sin().float() + output = apply_rotary_emb(tensor_, cos, sin).type_as(tensor) + return output + + +class Qwen2_5_VLVisionFlashAttention2(nn.Module): + def __init__(self, dim: int, num_heads: int = 16) -> None: + super().__init__() + self.num_heads = num_heads + self.qkv = nn.Linear(dim, dim * 3, bias=True) + self.proj = nn.Linear(dim, dim) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: torch.Tensor = None, + ) -> torch.Tensor: + seq_length = hidden_states.shape[0] + q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + q = apply_rotary_pos_emb_flashatt(q.unsqueeze(0), rotary_pos_emb).squeeze(0) + k = apply_rotary_pos_emb_flashatt(k.unsqueeze(0), rotary_pos_emb).squeeze(0) + + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() + attn_output = flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen).reshape( + seq_length, -1 + ) + attn_output = self.proj(attn_output) + return attn_output + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb_vision(tensor: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + orig_dtype = tensor.dtype + tensor = tensor.float() + cos = freqs.cos() + sin = freqs.sin() + cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float() + sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float() + output = (tensor * cos) + (rotate_half(tensor) * sin) + output = output.to(orig_dtype) + return output + + +class Qwen2_5_VLVisionSdpaAttention(nn.Module): + def __init__(self, dim: int, num_heads: int = 16) -> None: + super().__init__() + self.num_heads = num_heads + self.qkv = nn.Linear(dim, dim * 3, bias=True) + self.proj = nn.Linear(dim, dim) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: torch.Tensor = None, + ) -> torch.Tensor: + seq_length = hidden_states.shape[0] + q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0) + k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0) + + attention_mask = torch.zeros([1, seq_length, seq_length], device=q.device, dtype=torch.bool) + for i in range(1, len(cu_seqlens)): + attention_mask[ + ..., + cu_seqlens[i - 1] : cu_seqlens[i], + cu_seqlens[i - 1] : cu_seqlens[i], + ] = True + q = q.transpose(0, 1) + k = k.transpose(0, 1) + v = v.transpose(0, 1) + attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0) + attn_output = attn_output.transpose(0, 1) + attn_output = attn_output.reshape(seq_length, -1) + attn_output = self.proj(attn_output) + return attn_output + + +QWEN2_5_VL_VISION_ATTENTION_CLASSES = { + "flash_attention_2": Qwen2_5_VLVisionFlashAttention2, + "sdpa": Qwen2_5_VLVisionSdpaAttention, +} + +# If FlashAttention2 is not available, transparently fall back to SDPA. +if not is_flash_attn_2_available(): + QWEN2_5_VL_VISION_ATTENTION_CLASSES["flash_attention_2"] = ( + Qwen2_5_VLVisionSdpaAttention + ) + +class Qwen2_5_VLVisionBlock(nn.Module): + def __init__(self, config, attn_implementation: str = "sdpa") -> None: + super().__init__() + self.norm1 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) + self.norm2 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) + self.attn = QWEN2_5_VL_VISION_ATTENTION_CLASSES[attn_implementation]( + config.hidden_size, num_heads=config.num_heads + ) + self.mlp = Qwen2_5_VLMLP(config, bias=True) + + def forward(self, hidden_states, cu_seqlens, rotary_pos_emb) -> torch.Tensor: + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states), + cu_seqlens=cu_seqlens, + rotary_pos_emb=rotary_pos_emb, + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +class Qwen2_5_VisionTransformerPretrainedModel(nn.Module): + config_class = Qwen2_5_VLVisionConfig + _no_split_modules = ["Qwen2_5_VLVisionBlock"] + + def __init__(self, config) -> None: + super().__init__() + + self.dtype = torch.bfloat16 if config.torch_dtype == "bfloat16" else torch.float32 + + self.spatial_merge_size = config.spatial_merge_size + self.patch_size = config.patch_size + self.fullatt_block_indexes = config.fullatt_block_indexes + self.window_size = config.window_size + self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size + + self.patch_embed = Qwen2_5_VisionPatchEmbed( + patch_size=config.patch_size, + temporal_patch_size=config.temporal_patch_size, + in_channels=config.in_channels, + embed_dim=config.hidden_size, + ) + + head_dim = config.hidden_size // config.num_heads + self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2) + + self.blocks = nn.ModuleList( + [Qwen2_5_VLVisionBlock(config, config._attn_implementation) for _ in range(config.depth)] + ) + self.merger = Qwen2_5_VLPatchMerger( + dim=config.out_hidden_size, + context_dim=config.hidden_size, + spatial_merge_size=config.spatial_merge_size, + ) + self.gradient_checkpointing = False + + def init_weights(self, buffer_device: Optional[torch.device] = None): + self.rotary_pos_emb.init_weights(buffer_device) + + def rot_pos_emb(self, grid_thw): + pos_ids = [] + for t, h, w in grid_thw: + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) + hpos_ids = hpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + hpos_ids = hpos_ids.permute(0, 2, 1, 3) + hpos_ids = hpos_ids.flatten() + + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) + wpos_ids = wpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + wpos_ids = wpos_ids.permute(0, 2, 1, 3) + wpos_ids = wpos_ids.flatten() + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + pos_ids = torch.cat(pos_ids, dim=0) + max_grid_size = grid_thw[:, 1:].max() + rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) + return rotary_pos_emb + + def get_window_index(self, grid_thw): + window_index: list = [] + cu_window_seqlens: list = [0] + window_index_id = 0 + vit_merger_window_size = self.window_size // self.spatial_merge_size // self.patch_size + + for grid_t, grid_h, grid_w in grid_thw: + llm_grid_h, llm_grid_w = ( + grid_h // self.spatial_merge_size, + grid_w // self.spatial_merge_size, + ) + index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w) + pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size + pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size + num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size + num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size + index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", -100) + index_padded = index_padded.reshape( + grid_t, + num_windows_h, + vit_merger_window_size, + num_windows_w, + vit_merger_window_size, + ) + index_padded = index_padded.permute(0, 1, 3, 2, 4).reshape( + grid_t, + num_windows_h * num_windows_w, + vit_merger_window_size, + vit_merger_window_size, + ) + seqlens = (index_padded != -100).sum([2, 3]).reshape(-1) + index_padded = index_padded.reshape(-1) + index_new = index_padded[index_padded != -100] + window_index.append(index_new + window_index_id) + cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1] + cu_window_seqlens.extend(cu_seqlens_tmp.tolist()) + window_index_id += (grid_t * llm_grid_h * llm_grid_w).item() + window_index = torch.cat(window_index, dim=0) + + return window_index, cu_window_seqlens + + def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: + """ + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_size)`): + The final hidden states of the model. + grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): + The temporal, height and width of feature shape of each image in LLM. + + Returns: + `torch.Tensor`: hidden_states. + """ + hidden_states = self.patch_embed(hidden_states) + rotary_pos_emb = self.rot_pos_emb(grid_thw) + window_index, cu_window_seqlens = self.get_window_index(grid_thw) + cu_window_seqlens = torch.tensor( + cu_window_seqlens, + device=hidden_states.device, + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens) + + seq_len, _ = hidden_states.size() + hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + hidden_states = hidden_states[window_index, :, :] + hidden_states = hidden_states.reshape(seq_len, -1) + rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + rotary_pos_emb = rotary_pos_emb[window_index, :, :] + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) + + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + dim=0, + # Select dtype based on the following factors: + # - FA2 requires that cu_seqlens_q must have dtype int32 + # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw + # See https://github.com/huggingface/transformers/pull/34852 for more information + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + + for layer_num, blk in enumerate(self.blocks): + if layer_num in self.fullatt_block_indexes: + cu_seqlens_now = cu_seqlens + else: + cu_seqlens_now = cu_window_seqlens + if self.gradient_checkpointing and self.training: + hidden_states = self._gradient_checkpointing_func( + blk.__call__, hidden_states, cu_seqlens_now, rotary_pos_emb + ) + else: + hidden_states = blk( + hidden_states, + cu_seqlens=cu_seqlens_now, + rotary_pos_emb=rotary_pos_emb, + ) + + hidden_states = self.merger(hidden_states) + reverse_indices = torch.argsort(window_index) + hidden_states = hidden_states[reverse_indices, :] + + return hidden_states + + +class Qwen2_5_VLRotaryEmbedding(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig, device=None): + super().__init__() + # BC: "rope_type" was originally "type" + if hasattr(config, "rope_scaling") and config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + def init_weights(self, buffer_device: Optional[torch.device] = None): + if buffer_device is None: + device = self.inv_freq.device + else: + device = buffer_device + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.inv_freq = inv_freq + + def _dynamic_frequency_update(self, position_ids, device): + raise NotImplementedError("Dynamic frequency update is not implemented for Qwen2_5_VL") + # """ + # dynamic RoPE layers should recompute `inv_freq` in the following situations: + # 1 - growing beyond the cached sequence length (allow scaling) + # 2 - the current sequence length is in the original scale (avoid losing precision with small sequences) + # """ + # seq_len = torch.max(position_ids) + 1 + # if seq_len > self.max_seq_len_cached: # growth + # inv_freq, self.attention_scaling = self.rope_init_fn( + # self.config, device, seq_len=seq_len, **self.rope_kwargs + # ) + # self.register_buffer("inv_freq", inv_freq, persistent=False) + # self.max_seq_len_cached = seq_len + + # if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset + # self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) + # self.max_seq_len_cached = self.original_max_seq_len + + @torch.no_grad() + def forward(self, x, position_ids): + if "dynamic" in self.rope_type: + self._dynamic_frequency_update(position_ids, device=x.device) + + # Core RoPE block. In contrast to other models, Qwen2_5_VL has different position ids for thw grids + # So we expand the inv_freq to shape (3, ...) + inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) + position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions) + # Force float32 (see https://github.com/huggingface/transformers/pull/29285) + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + + # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention + cos = cos * self.attention_scaling + sin = sin * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class Qwen2MLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + # in llama self.w2(F.silu(self.w1(x)) * self.w3(x)) + # i.e. w2 is self.down_proj, w1 is self.gate_proj, w3 is self.up_proj + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1): + """Applies Rotary Position Embedding with Multimodal Sections to the query and key tensors (https://qwenlm.github.io/blog/qwen2-vl/). + + Explanation: + Multimodal 3D rotary position embedding is an extension to 1D rotary position embedding. The input embedding + sequence contains vision (images / videos) embedding and text embedding or just contains text embedding. For + vision embedding part, we apply rotary position embedding on temporal, height and width dimension seperately. + Here we split the channel dimension to 3 chunks for the temporal, height and width rotary position embedding. + For text embedding part, we just apply 1D rotary position embedding. The three rotary position index (temporal, + height and width) of text embedding is always the same, so the text embedding rotary position embedding has no + difference with modern LLMs. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + The position indices of the tokens corresponding to the query and key tensors. For example, this can be + used to pass offsetted position ids when working with a KV-cache. + mrope_section(`List(int)`): + Multimodal rope section is for channel dimension of temporal, height and width in rope calculation. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + if isinstance(mrope_section, omegaconf.listconfig.ListConfig): + mrope_section = list(mrope_section) + mrope_section = mrope_section * 2 + cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim) + sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +class Qwen2_5_VLAttention(nn.Module): + """ + Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer + and "Generating Long Sequences with Sparse Transformers". + """ + + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.is_causal = True + self.attention_dropout = config.attention_dropout + self.rope_scaling = config.rope_scaling + + if (self.head_dim * self.num_heads) != self.hidden_size: + raise ValueError( + f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" + f" and `num_heads`: {self.num_heads})." + ) + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) + + def init_weights(self): + self.rotary_emb.init_weights() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_multimodal_rotary_pos_emb( + query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] + ) + + if past_key_value is not None: + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attention_mask is not None: # no matter the length, we just slice it + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + # Fix precision issues in Qwen2-VL float16 inference + # Replace inf values with zeros in attention weights to prevent NaN propagation + if query_states.dtype == torch.float16: + attn_weights = torch.where(torch.isinf(attn_weights), torch.zeros_like(attn_weights), attn_weights) + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, -1) + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): + """ + Qwen2_5_VL flash attention module, following Qwen2_5_VL attention module. This module inherits from `Qwen2_5_VLAttention` + as the weights of the module stays untouched. The only required change would be on the forward pass + where it needs to correctly call the public API of flash attention and deal with padding tokens + in case the input contains any of them. Additionally, for sliding window attention, we apply SWA only to the bottom + config.max_window_layers layers. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + cp_mesh: DeviceMesh | None = None, + ): + """ + Args: + cp_mesh (DeviceMesh, optional): Device mesh over which context parallelism is done. + """ + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + + # Because the input can be padded, the absolute sequence length depends on the max position id. + cos, sin = position_embeddings + query_states, key_states = apply_multimodal_rotary_pos_emb( + query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] + ) + + if past_key_value is not None: + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + dropout_rate = 0.0 if not self.training else self.attention_dropout + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in float16 just to be sure everything works as expected. + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + # Reashape to the expected shape for Flash Attention + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + if ( + self.config.use_sliding_window + and getattr(self.config, "sliding_window", None) is not None + and self.layer_idx >= self.config.max_window_layers + ): + sliding_window = self.config.sliding_window + else: + sliding_window = None + if cp_mesh is not None: + raise NotImplementedError( + "CP is not supported for flash attention2, _flash_attention_forward will produce wrong output if query_states is sharded" + ) + + attn_output = _flash_attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + dropout=dropout_rate, + sliding_window=sliding_window, + is_causal=self.is_causal, + use_top_left_mask=self._flash_attn_uses_top_left_mask, + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention): + """ + Qwen2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from + `Qwen2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to + SDPA API. + """ + + # Adapted from Qwen2Attention.forward + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + cp_mesh: DeviceMesh | None = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if output_attentions: + logger.warning_once( + "Qwen2_5_VLModel is using Qwen2_5_VLSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " + 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + return super().forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_multimodal_rotary_pos_emb( + query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] + ) + + if past_key_value is not None: + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + causal_mask = attention_mask + if attention_mask is not None and attention_mask.ndim == 4: # no matter the length, we just slice it + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + + # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, + # Reference: https://github.com/pytorch/pytorch/issues/112577. + if query_states.device.type == "cuda" and attention_mask is not None: + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. + is_causal = True if causal_mask is None and q_len > 1 else False + if cp_mesh is not None: + key_states = DTensor.from_local(key_states, cp_mesh, [Shard(2)]).full_tensor() + value_states = DTensor.from_local(value_states, cp_mesh, [Shard(2)]).full_tensor() + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=causal_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + is_causal=is_causal, + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, self.hidden_size) + + attn_output = self.o_proj(attn_output) + + return attn_output, None, past_key_value + + +QWEN2_5_VL_ATTENTION_CLASSES = { + "flash_attention_2": Qwen2_5_VLFlashAttention2, + "sdpa": Qwen2_5_VLSdpaAttention, +} + + +# If FlashAttention2 is not available, transparently fall back to SDPA. +if not is_flash_attn_2_available(): + QWEN2_5_VL_ATTENTION_CLASSES["flash_attention_2"] = Qwen2_5_VLSdpaAttention + +class Qwen2_5_VLDecoderLayer(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + if config.use_sliding_window and config._attn_implementation != "flash_attention_2": + logger.warning_once( + f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " + "unexpected results may be encountered." + ) + self.self_attn = QWEN2_5_VL_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx) + self.mlp = Qwen2MLP(config) + self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + cp_mesh: DeviceMesh | None = None, + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): attention mask of size + `(batch, sequence_length)` where padding elements are indicated by 0. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. + position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): + Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, + with `head_dim` being the embedding dimension of each attention head. + cp_mesh (DeviceMesh, optional): Device mesh over which context parallelism is done. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states, + attention_mask, + position_ids, + past_key_value, + output_attentions, + use_cache, + cache_position, + position_embeddings, + cp_mesh=cp_mesh, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + return outputs + + +class Qwen2_5_VLModel(nn.Module): + def __init__(self, config: Qwen2_5_VLConfig): + super().__init__() + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen2_5_VLDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self._attn_implementation = config._attn_implementation + self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) + + self.gradient_checkpointing = False + self.dtype = config.torch_dtype + self.cp_mesh = None + + def init_weights(self, buffer_device: Optional[torch.device] = None): + self.rotary_emb.init_weights(buffer_device) + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def set_cp_mesh(self, cp_mesh): + self.cp_mesh = cp_mesh + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + assert not self.gradient_checkpointing + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + # torch.jit.trace() doesn't support cache objects in the output + if use_cache and past_key_values is None and not torch.jit.is_tracing(): + past_key_values = DynamicCache() + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, + ) + + # the hard coded `3` is for temporal, height and width. + if position_ids is None: + position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1) + elif position_ids.dim() == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + + if self.cp_mesh is None: + causal_mask = self._update_causal_mask( + attention_mask, + inputs_embeds, + cache_position, + past_key_values, + output_attentions, + ) + + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + layer_id = 0 + seqlen = hidden_states.shape[1] + if self.config._attn_implementation == "sdpa": + causal_mask = torch.full((seqlen, seqlen), float("-inf"), device=hidden_states.device).triu_(1) + causal_mask = causal_mask.to(hidden_states.dtype) + if self.cp_mesh is not None: + seq_range = self._seq_range(seqlen) + position_embeddings = ( + position_embeddings[0][:, :, seq_range[0] : seq_range[1], :], + position_embeddings[1][:, :, seq_range[0] : seq_range[1], :], + ) + hidden_states = hidden_states[:, seq_range[0] : seq_range[1], :] + position_ids = position_ids[:, :, seq_range[0] : seq_range[1]] + cache_position = cache_position[seq_range[0] : seq_range[1]] + causal_mask = causal_mask[seq_range[0] : seq_range[1]] + assert past_key_values is None, "not support cp with past_key_values" + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + layer_id += 1 + if self.gradient_checkpointing and self.training: + raise NotImplementedError("Gradient checkpointing is not implemented here") + # layer_outputs = self._gradient_checkpointing_func( + # decoder_layer.__call__, + # hidden_states, + # causal_mask, + # position_ids, + # past_key_values, + # output_attentions, + # use_cache, + # cache_position, + # position_embeddings, + # ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, # None + position_ids=position_ids, + past_key_value=past_key_values, # None + output_attentions=output_attentions, # bool + use_cache=use_cache, # bool + cache_position=cache_position, # None + position_embeddings=position_embeddings, + cp_mesh=self.cp_mesh, # device mesh + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + def _seq_range(self, seqlen) -> tuple[int, int]: + if self.cp_mesh is not None: + assert seqlen % self.cp_mesh.size() == 0, f"seqlen: {seqlen}, mesh size: {self.cp_mesh.size()}" + local_seqlen = seqlen // self.cp_mesh.size() + cp_rank = self.cp_mesh.get_local_rank() + return (cp_rank * local_seqlen, (cp_rank + 1) * local_seqlen) + else: + return (0, seqlen) + + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool, + ): + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and past_key_values is not None: + is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0] + if is_padding_right: + raise ValueError( + "You are attempting to perform batched generation with padding_side='right'" + " this may lead to unexpected behaviour for Flash Attention version of Qwen2_5_VL. Make sure to " + " call `tokenizer.padding_side = 'left'` before tokenizing the input. " + ) + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + using_static_cache = isinstance(past_key_values, StaticCache) + using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if ( + self.config._attn_implementation == "sdpa" + and not (using_static_cache or using_sliding_window_cache) + and not output_attentions + ): + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + sliding_window=self.config.sliding_window, + is_training=self.training, + ): + return None + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + # SlidingWindowCache or StaticCache + if using_sliding_window_cache or using_static_cache: + target_length = past_key_values.get_max_cache_shape() + # DynamicCache or no cache + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). + causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=target_length, + dtype=dtype, + device=device, + cache_position=cache_position, + batch_size=input_tensor.shape[0], + config=self.config, + past_key_values=past_key_values, + ) + + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type in ["cuda", "xpu"] + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + + @staticmethod + def _prepare_4d_causal_attention_mask_with_cache_position( + attention_mask: torch.Tensor, + sequence_length: int, + target_length: int, + dtype: torch.dtype, + device: torch.device, + cache_position: torch.Tensor, + batch_size: int, + config: Qwen2_5_VLConfig, + past_key_values: Cache, + ): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. + + Args: + attention_mask (`torch.Tensor`): + A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. + sequence_length (`int`): + The sequence length being processed. + target_length (`int`): + The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. + dtype (`torch.dtype`): + The dtype to use for the 4D attention mask. + device (`torch.device`): + The device to plcae the 4D attention mask on. + cache_position (`torch.Tensor`): + Indices depicting the position of the input sequence tokens in the sequence. + batch_size (`torch.Tensor`): + Batch size. + config (`Qwen2_5_VLConfig`): + The model's configuration class + past_key_values (`Cache`): + The cache class that is being used currently to generate + """ + if attention_mask is not None and attention_mask.dim() == 4: + # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. + causal_mask = attention_mask + else: + min_dtype = torch.finfo(dtype).min + causal_mask = torch.full( + (sequence_length, target_length), + fill_value=min_dtype, + dtype=dtype, + device=device, + ) + diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + if config.sliding_window is not None: + # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also + # the check is needed to verify is current checkpoint was trained with sliding window or not + if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length: + sliding_attend_mask = torch.arange(target_length, device=device) <= ( + cache_position.reshape(-1, 1) - config.sliding_window + ) + diagonal_attend_mask.bitwise_or_(sliding_attend_mask) + causal_mask *= diagonal_attend_mask + causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + if attention_mask.shape[-1] > target_length: + attention_mask = attention_mask[:, :target_length] + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( + causal_mask.device + ) + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + return causal_mask + + +@dataclass +class Qwen2_5_VLCausalLMOutputWithPast(ModelOutput): + """ + Base class for Qwen2_5_VL causal language model (or autoregressive) outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. + """ + + loss: Optional[torch.FloatTensor] = None + logits: torch.FloatTensor = None + past_key_values: Optional[List[torch.FloatTensor]] = None + hidden_states: Optional[Tuple[torch.FloatTensor]] = None + attentions: Optional[Tuple[torch.FloatTensor]] = None + rope_deltas: Optional[torch.LongTensor] = None + + +class Qwen2_5_VLForConditionalGenerationSimple(nn.Module): + _tied_weights_keys = ["lm_head.weight"] + config_class = Qwen2_5_VLConfig + _no_split_modules = ["Qwen2VLDecoderLayer", "Qwen2_5_VLVisionBlock"] + + def __init__(self, config): + super().__init__() + self.config = config + self.visual = Qwen2_5_VisionTransformerPretrainedModel(config.vision_config) + + self.model = Qwen2_5_VLModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.rope_deltas = None + + def init_weights(self, buffer_device: Optional[torch.device] = None): + self.model.init_weights(buffer_device) + self.visual.init_weights(buffer_device) + + def can_generate(self): + return True + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def get_rope_index( + self, + input_ids: Optional[torch.LongTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Calculate the 3D rope index based on image and video's temporal, height and width in LLM. + + Explanation: + Each embedding sequence contains vision embedding and text embedding or just contains text embedding. + + For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. + Examples: + input_ids: [T T T T T], here T is for text. + temporal position_ids: [0, 1, 2, 3, 4] + height position_ids: [0, 1, 2, 3, 4] + width position_ids: [0, 1, 2, 3, 4] + + For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part + and 1D rotary position embeddin for text part. + Examples: + Temporal (Time): 3 patches, representing different segments of the video in time. + Height: 2 patches, dividing each frame vertically. + Width: 2 patches, dividing each frame horizontally. + We also have some important parameters: + fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. + tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity. + temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. + interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs. + input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. + vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] + vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] + vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] + text temporal position_ids: [101, 102, 103, 104, 105] + text height position_ids: [101, 102, 103, 104, 105] + text width position_ids: [101, 102, 103, 104, 105] + Here we calculate the text start position_ids as the max vision position_ids plus 1. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): + The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + Returns: + position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) + mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) + """ + spatial_merge_size = self.config.vision_config.spatial_merge_size + image_token_id = self.config.image_token_id + video_token_id = self.config.video_token_id + vision_start_token_id = self.config.vision_start_token_id + mrope_position_deltas = [] + if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None): + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_index, video_index = 0, 0 + attention_mask = attention_mask.to(total_input_ids.device) + for i, input_ids in enumerate(total_input_ids): + input_ids = input_ids[attention_mask[i] == 1] + image_nums, video_nums = 0, 0 + vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) + vision_tokens = input_ids[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + input_tokens = input_ids.tolist() + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + second_per_grid_t = 0 + image_index += 1 + remain_images -= 1 + ed = ed_image + + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + if second_per_grid_ts is not None: + second_per_grid_t = second_per_grid_ts[video_index] + else: + second_per_grid_t = 1.0 + video_index += 1 + remain_videos -= 1 + ed = ed_video + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + text_len = ed - st + + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + st_idx + ) + + range_tensor = torch.arange(llm_grid_t).view(-1, 1) + expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) + + time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second + + time_tensor_long = time_tensor.long() + t_index = time_tensor_long.flatten() + + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() + llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + text_len = len(input_tokens) - st + llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + return position_ids, mrope_position_deltas + else: + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device) + max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0] + mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] + else: + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, 1, -1) + .expand(3, input_ids.shape[0], -1) + ) + mrope_position_deltas = torch.zeros( + [input_ids.shape[0], 1], + device=input_ids.device, + dtype=input_ids.dtype, + ) + + return position_ids, mrope_position_deltas + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + ) -> Union[Tuple, Qwen2_5_VLCausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + >>> model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") + >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") + + >>> messages = [ + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": "What is shown in this image?"}, + ], + }, + ] + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos]) + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if inputs_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) + # This is a trick to handle TP for LLM but no TP for vision encoder, we need to convert DTensor to regular tensor later + is_inputs_embeds_dtensor = isinstance(inputs_embeds, DTensor) # This is True for TP>1, False for TP=1 + if is_inputs_embeds_dtensor: + target_device_mesh = inputs_embeds.device_mesh + target_placements = inputs_embeds.placements + inputs_embeds = inputs_embeds.full_tensor() + + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_features = image_embeds.shape[0] + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + + mask = input_ids == self.config.image_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + image_mask = mask_expanded.to(inputs_embeds.device) + + image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + + mask = input_ids == self.config.video_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + video_mask = mask_expanded.to(inputs_embeds.device) + + video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + if is_inputs_embeds_dtensor: + inputs_embeds = ( + DTensor.from_local(inputs_embeds, device_mesh=target_device_mesh) + .redistribute(placements=target_placements) + .to_local() + ) + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + # if we get 4D attention mask we cannot calculate rope deltas anymore. + if position_ids is None and (attention_mask is None or attention_mask.ndim == 2): + # calculate RoPE index once per generation in the pre-fill stage only + if ( + (cache_position is not None and cache_position[0] == 0) + or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ): + position_ids, rope_deltas = self.get_rope_index( + input_ids, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + attention_mask, + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + batch_size, seq_length, _ = inputs_embeds.shape + delta = ( + (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) if cache_position is not None else 0 + ) + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + outputs = self.model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + + loss = None + if labels is not None: + # Upcast to float if we need to compute the loss to avoid potential precision issues + logits = logits.float() + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) # + outputs[1:] + return (loss,) + output if loss is not None else output + + return Qwen2_5_VLCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=self.rope_deltas, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + cache_position=None, + position_ids=None, + use_cache=True, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + second_per_grid_ts=None, + **kwargs, + ): + # Overwritten -- in specific circumstances we don't want to forward image inputs to the model + + # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens + # Exception 1: when passing input_embeds, input_ids may be missing entries + # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here + # Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case. + # (we can't check exception 3 while compiling) + # Exception 4: If input_embeds are passed then slice it through `cache_position`, to keep only the unprocessed tokens and + # generate the first token for each sequence. Later use the generated Input ids for continuation. + if past_key_values is not None: + if inputs_embeds is not None and input_ids.shape[1] == 0: # Exception 4 + inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :] + elif inputs_embeds is not None or ( # Exception 1 + is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1] + ): # Exception 3 + input_ids = input_ids[:, -cache_position.shape[0] :] + elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2) + input_ids = input_ids[:, cache_position] + + if cache_position[0] != 0: + pixel_values = None + pixel_values_videos = None + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]: + model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None} + else: + model_inputs = {"input_ids": input_ids, "inputs_embeds": None} + + if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2: + if model_inputs["inputs_embeds"] is not None: + batch_size, sequence_length, _ = inputs_embeds.shape + device = inputs_embeds.device + else: + batch_size, sequence_length = input_ids.shape + device = input_ids.device + + attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=past_key_values.get_max_cache_shape(), + dtype=self.lm_head.weight.dtype, + device=device, + cache_position=cache_position, + batch_size=batch_size, + config=self.config, + past_key_values=past_key_values, + ) + + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": use_cache, + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "pixel_values_videos": pixel_values_videos, + "image_grid_thw": image_grid_thw, + "video_grid_thw": video_grid_thw, + "cache_position": cache_position, + "second_per_grid_ts": second_per_grid_ts, + } + ) + return model_inputs + + +__all__ = ["Qwen2_5_VLForConditionalGenerationSimple"] + + +def get_rope_index( + model_config, + input_ids: Optional[torch.LongTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Calculate the 3D rope index based on image and video's temporal, height and width in LLM. + + Explanation: + Each embedding sequence contains vision embedding and text embedding or just contains text embedding. + + For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. + Examples: + input_ids: [T T T T T], here T is for text. + temporal position_ids: [0, 1, 2, 3, 4] + height position_ids: [0, 1, 2, 3, 4] + width position_ids: [0, 1, 2, 3, 4] + + For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part + and 1D rotary position embeddin for text part. + Examples: + Temporal (Time): 3 patches, representing different segments of the video in time. + Height: 2 patches, dividing each frame vertically. + Width: 2 patches, dividing each frame horizontally. + We also have some important parameters: + fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. + tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity. + temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. + interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs. + input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. + vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] + vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] + vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] + text temporal position_ids: [101, 102, 103, 104, 105] + text height position_ids: [101, 102, 103, 104, 105] + text width position_ids: [101, 102, 103, 104, 105] + Here we calculate the text start position_ids as the max vision position_ids plus 1. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): + The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + Returns: + position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) + mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) + """ + spatial_merge_size = model_config.vision_config.spatial_merge_size + image_token_id = model_config.image_token_id + video_token_id = model_config.video_token_id + vision_start_token_id = model_config.vision_start_token_id + mrope_position_deltas = [] + if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None): + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_index, video_index = 0, 0 + attention_mask = attention_mask.to(total_input_ids.device) + for i, input_ids in enumerate(total_input_ids): + input_ids = input_ids[attention_mask[i] == 1] + image_nums, video_nums = 0, 0 + vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) + vision_tokens = input_ids[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + input_tokens = input_ids.tolist() + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + second_per_grid_t = 0 + image_index += 1 + remain_images -= 1 + ed = ed_image + + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + if second_per_grid_ts is not None: + second_per_grid_t = second_per_grid_ts[video_index] + else: + second_per_grid_t = 1.0 + video_index += 1 + remain_videos -= 1 + ed = ed_video + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + text_len = ed - st + + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + st_idx + ) + + range_tensor = torch.arange(llm_grid_t, device=input_ids.device).view(-1, 1) + expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) + + time_tensor = expanded_range * second_per_grid_t * model_config.vision_config.tokens_per_second + + time_tensor_long = time_tensor.long() + t_index = time_tensor_long.flatten() + + h_index = ( + torch.arange(llm_grid_h, device=input_ids.device) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w, device=input_ids.device) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + text_len = len(input_tokens) - st + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + st_idx + ) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + return position_ids, mrope_position_deltas + else: + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device) + max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0] + mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] + else: + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, 1, -1) + .expand(3, input_ids.shape[0], -1) + ) + mrope_position_deltas = torch.zeros( + [input_ids.shape[0], 1], + device=input_ids.device, + dtype=input_ids.dtype, + ) + + return position_ids, mrope_position_deltas diff --git a/fastvideo/models/encoders/reason1.py b/fastvideo/models/encoders/reason1.py new file mode 100644 index 000000000..48eff9286 --- /dev/null +++ b/fastvideo/models/encoders/reason1.py @@ -0,0 +1,353 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reason1 (Qwen2.5-VL) text encoder.""" + +import os +from dataclasses import dataclass +from collections.abc import Iterable + +import torch +from transformers import AutoProcessor + +from fastvideo.configs.models.encoders import BaseEncoderOutput, Reason1Config +from fastvideo.logger import init_logger +from fastvideo.models.encoders.base import TextEncoder +from fastvideo.models.loader.weight_utils import default_weight_loader +from fastvideo.platforms import AttentionBackendEnum + +from fastvideo.models.encoders.qwen2_5_vl_custom import ( + Qwen2_5_VLForConditionalGenerationSimple, + Qwen2_5_VLConfig, + get_rope_index, +) + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class _WeightsSource: + """Mimic `TextEncoderLoader.Source` (avoid import cycles).""" + + model_or_path: str + prefix: str = "" + fall_back_to_pt: bool = True + allow_patterns_overrides: list[str] | None = None + + + + +class Reason1TextEncoder(TextEncoder): + """Reason1 (Qwen2.5-VL) text encoder.""" + + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ) + + def __init__(self, config: Reason1Config, prefix: str = "", checkpoint_path: str | None = None): + super().__init__(config) + + self.prefix = prefix + self.quant_config = None # For future quantization support + + self.embedding_concat_strategy = config.arch_config.embedding_concat_strategy + self.n_layers_per_group = config.arch_config.n_layers_per_group + self.num_embedding_padding_tokens = config.arch_config.num_embedding_padding_tokens + + config_path = checkpoint_path if checkpoint_path else config.tokenizer_type + + logger.info("Initializing Reason1TextEncoder (Qwen2.5-VL) from %s", config_path) + try: + from transformers import AutoConfig as HFAutoConfig + hf_config = HFAutoConfig.from_pretrained( + config_path, + trust_remote_code=True, + ) + except Exception as e: + logger.warning("Failed to load HF config from %s (%s). Using default Qwen2.5-VL-7B config.", + config_path, e) + hf_config = Qwen2_5_VLConfig( + hidden_size=3584, + intermediate_size=18944, + max_window_layers=28, + num_attention_heads=28, + num_hidden_layers=28, + num_key_value_heads=4, + tie_word_embeddings=False, + vocab_size=152064, + ) + + hf_config.output_hidden_states = True + + if hasattr(config.arch_config, '_attn_implementation') and config.arch_config._attn_implementation: + hf_config._attn_implementation = config.arch_config._attn_implementation + else: + hf_config._attn_implementation = "flash_attention_2" + logger.info("Reason1 attention implementation: %s", getattr(hf_config, "_attn_implementation", None)) + + with torch.device("meta"): + self.model = Qwen2_5_VLForConditionalGenerationSimple(hf_config) + + self.processor = AutoProcessor.from_pretrained( + config_path, + trust_remote_code=True, + ) + + weights_override = os.getenv("FASTVIDEO_REASON1_WEIGHTS_PATH") + if weights_override: + self.secondary_weights = ( + _WeightsSource( + model_or_path=weights_override, + prefix="", + fall_back_to_pt=True, + allow_patterns_overrides=None, + ), + ) + logger.info("Reason1TextEncoder: overlaying weights from %s", weights_override) + + self._weights_loaded = False + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> BaseEncoderOutput: + # Cosmos2.5 alignment: keep attention_mask=None. + outputs = self.model( + input_ids=input_ids, + attention_mask=None, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_hidden_states=True, + return_dict=True, + pixel_values=kwargs.get('pixel_values', None), + pixel_values_videos=kwargs.get('pixel_values_videos', None), + image_grid_thw=kwargs.get('image_grid_thw', None), + video_grid_thw=kwargs.get('video_grid_thw', None), + ) + + hidden_states = outputs.hidden_states + last_hidden_state = hidden_states[-1] + + return BaseEncoderOutput( + last_hidden_state=last_hidden_state, + hidden_states=hidden_states if output_hidden_states else None, + attention_mask=None, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + first_weight = None + weights_list = [] + for name, weight in weights: + if first_weight is None: + first_weight = weight + self.model = self.model.to_empty(device=weight.device) + self.model.init_weights(buffer_device=weight.device) + weights_list.append((name, weight)) + + params_dict = dict(self.model.named_parameters()) + loaded_params: set[str] = set() + skipped_weights = {"lm_head": 0, "visual": 0, "decoder": 0} + + for name, loaded_weight in weights_list: + if "lm_head" in name: + skipped_weights["lm_head"] += 1 + continue + if "visual" in name: + skipped_weights["visual"] += 1 + continue + if "decoder" in name: + skipped_weights["decoder"] += 1 + continue + + # Handle stacked params mapping (for quantized models) + for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + param_name_with_prefix = f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}" + loaded_params.add(param_name_with_prefix) + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + param_name_with_prefix = f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}" + loaded_params.add(param_name_with_prefix) + if first_weight is not None: + self.model = self.model.to(first_weight.device) + + all_params = set(f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}" + for name in params_dict.keys()) + loaded_params.update(all_params) + + # Mark weights as loaded + self._weights_loaded = True + return loaded_params + + def compute_text_embeddings_online( + self, + data_batch: dict[str, list[str]], + input_caption_key: str, + ) -> torch.Tensor: + prompts = data_batch[input_caption_key] + return self.compute_text_embeddings(prompts) + + def compute_text_embeddings( + self, + prompts: list[str], + device: str | torch.device = "cuda", + ) -> torch.Tensor: + """Compute embeddings for a list of prompts.""" + input_ids_batch = [] + + tok = getattr(self.processor, "tokenizer", None) + if tok is None: + raise RuntimeError("Reason1TextEncoder requires processor.tokenizer") + pad_id = getattr(tok, "pad_id", None) + if pad_id is None: + pad_id = getattr(tok, "pad_token_id", None) + if pad_id is None: + pad_id = getattr(self.model.config, "pad_token_id", None) + if pad_id is None: + pad_id = 0 + + for prompt in prompts: + conversations = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a helpful assistant who will provide prompts to an image generator.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": prompt, + } + ], + }, + ] + + try: + tokenizer_output = tok.apply_chat_template( + conversations, + tokenize=True, + add_generation_prompt=False, + add_vision_id=False, + ) + except TypeError: + tokenizer_output = tok.apply_chat_template( + conversations, + tokenize=True, + add_generation_prompt=False, + ) + + if isinstance(tokenizer_output, dict) and "input_ids" in tokenizer_output: + input_ids = tokenizer_output["input_ids"] + if hasattr(input_ids, "tolist"): + input_ids = input_ids.tolist() + else: + input_ids = tokenizer_output + if hasattr(input_ids, "tolist"): + input_ids = input_ids.tolist() + if isinstance(input_ids, list) and len(input_ids) == 1 and isinstance( + input_ids[0], list): + input_ids = input_ids[0] + if not isinstance(input_ids, list): + raise RuntimeError( + f"Unexpected chat_template output type: {type(tokenizer_output)}" + ) + + if self.num_embedding_padding_tokens > len(input_ids): + pad_len = self.num_embedding_padding_tokens - len(input_ids) + input_ids = input_ids + [pad_id] * pad_len + else: + input_ids = input_ids[:self.num_embedding_padding_tokens] + + input_ids = torch.LongTensor(input_ids).to(device=device) + input_ids_batch.append(input_ids) + + input_ids_batch = torch.stack(input_ids_batch, dim=0) + + # Cosmos2.5 alignment: keep attention_mask=None. + target_device = input_ids_batch.device + try: + embed_device = self.model.model.embed_tokens.weight.device # type: ignore[attr-defined] + except Exception: + embed_device = None + if embed_device is not None and embed_device != target_device: + self.model = self.model.to(target_device) + + with torch.no_grad(): + position_ids, _ = get_rope_index( + self.model.config, + input_ids_batch, + image_grid_thw=None, + video_grid_thw=None, + second_per_grid_ts=None, + attention_mask=None, + ) + position_ids = position_ids.to(target_device) + + outputs = self.model.model( + input_ids=input_ids_batch, + position_ids=position_ids, + attention_mask=None, + output_hidden_states=True, + return_dict=True, + use_cache=False, + ) + hidden_states = outputs.hidden_states + + normalized_hidden_states = [] + for layer_idx in range(1, len(hidden_states)): + normalized_state = self._mean_normalize(hidden_states[layer_idx]) + normalized_hidden_states.append(normalized_state) + + if self.embedding_concat_strategy == "full_concat": + text_embeddings = torch.cat(normalized_hidden_states, dim=-1) + elif self.embedding_concat_strategy == "mean_pooling": + text_embeddings = torch.stack(normalized_hidden_states).mean(dim=0) + elif self.embedding_concat_strategy == "pool_every_n_layers_and_concat": + pooled_embeddings = [] + for i in range(0, len(normalized_hidden_states), self.n_layers_per_group): + group = normalized_hidden_states[i : i + self.n_layers_per_group] + pooled = torch.stack(group).mean(dim=0) + pooled_embeddings.append(pooled) + text_embeddings = torch.cat(pooled_embeddings, dim=-1) + else: + raise ValueError( + f"Unknown embedding_concat_strategy: {self.embedding_concat_strategy}" + ) + + return text_embeddings + @staticmethod + def _mean_normalize(tensor: torch.Tensor) -> torch.Tensor: + return (tensor - tensor.mean(dim=-1, keepdim=True)) / ( + tensor.std(dim=-1, keepdim=True) + 1e-8 + ) + diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index c76ddc95b..2b08c5b69 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -15,8 +15,7 @@ import torch.distributed as dist import torch.nn as nn from safetensors.torch import load_file as safetensors_load_file from torch.distributed import init_device_mesh -from transformers import AutoImageProcessor, AutoTokenizer -from transformers import UMT5EncoderModel +from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer from transformers.utils import SAFE_WEIGHTS_INDEX_NAME from fastvideo.configs.models import EncoderConfig @@ -91,10 +90,12 @@ class ComponentLoader(ABC): if module_type in module_loaders: loader_cls, expected_library = module_loaders[module_type] - # Assert that the library matches what's expected for this module type - assert transformers_or_diffusers == expected_library, ( - f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}" - ) + # Allow fastvideo.* libraries for custom implementations (e.g. Cosmos2_5Pipeline) + # that aren't available in diffusers/transformers yet + is_fastvideo_module = transformers_or_diffusers.startswith("fastvideo.") + if not is_fastvideo_module: + # Assert that the library matches what's expected for this module type + assert transformers_or_diffusers == expected_library, f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}" return loader_cls() # For unknown module types, use a generic loader @@ -449,6 +450,33 @@ class TokenizerLoader(ComponentLoader): """Load the tokenizer based on the model path, and inference args.""" logger.info("Loading tokenizer from %s", model_path) + # Cosmos2.5 stores an AutoProcessor config in `tokenizer/config.json` (not a tokenizer + # config). Use its `_name_or_path` (e.g. Qwen/Qwen2.5-VL-7B-Instruct) as the source. + tokenizer_cfg_path = os.path.join(model_path, "config.json") + if os.path.exists(tokenizer_cfg_path): + try: + with open(tokenizer_cfg_path, "r") as f: + tokenizer_cfg = json.load(f) + if isinstance(tokenizer_cfg, dict) and ( + tokenizer_cfg.get("_class_name") == "AutoProcessor" + or "processor_type" in tokenizer_cfg + ): + src = tokenizer_cfg.get("_name_or_path", "") + if isinstance(src, str) and src.strip(): + processor = AutoProcessor.from_pretrained( + src.strip(), + trust_remote_code=True, + ) + logger.info( + "Loaded tokenizer/processor from %s: %s", + src, + processor.__class__.__name__, + ) + return processor + except Exception: + # If parsing fails, fall through to AutoTokenizer below. + pass + tokenizer = AutoTokenizer.from_pretrained( model_path, # "/tokenizer" # in v0, this was same string as encoder_name "ClipTextModel" @@ -491,19 +519,40 @@ class VAELoader(ComponentLoader): if fastvideo_args.pipeline_config.vae_precision else torch.bfloat16 ): + # Cosmos2.5 uses a Wan2.1 VAE stored as `tokenizer.safetensors` under the VAE folder. + is_cosmos25 = fastvideo_args.pipeline_config.__class__.__name__ == "Cosmos25Config" + if class_name == "AutoencoderKLWan" and is_cosmos25: + from fastvideo.models.vaes.cosmos25wanvae import Cosmos25WanVAE + + dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] + vae = Cosmos25WanVAE(device=target_device, dtype=dtype) + + weight_path = os.path.join(model_path, "tokenizer.safetensors") + if not os.path.exists(weight_path): + raise FileNotFoundError( + f"Missing Cosmos2.5 VAE weights: {weight_path}" + ) + sd = safetensors_load_file(weight_path) + vae.load_state_dict(sd, strict=False) + return vae.eval() + vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) vae = vae_cls(vae_config).to(target_device) # Find all safetensors files safetensors_list = glob.glob( - os.path.join(str(model_path), "*.safetensors") - ) - loaded = {} - for sf_file in safetensors_list: - loaded.update(safetensors_load_file(sf_file)) - vae.load_state_dict( - loaded, strict=False - ) # We might only load encoder or decoder + os.path.join(str(model_path), "*.safetensors")) + if not safetensors_list: + raise ValueError(f"No safetensors files found in {model_path}") + # Common case: a single `.safetensors` checkpoint file. + # Some models may be sharded into multiple files; in that case we merge. + 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)) + vae.load_state_dict(loaded, strict=False) return vae.eval() @@ -581,7 +630,17 @@ class TransformerLoader(ComponentLoader): ] # Load the model using FSDP loader + logger.info("Loading model from %s, default_dtype: %s", cls_name, + default_dtype) assert fastvideo_args.hsdp_shard_dim is not None + # Cosmos2.5 checkpoints can include extra entries not present in the + # instantiated model (e.g. pos_embedder ranges / *_extra_state). Load + # non-strictly for Cosmos2.5 only; keep upstream strict behavior for others. + strict_load = not ( + cls_name.startswith("Cosmos25") + or cls_name == "Cosmos25Transformer3DModel" + or getattr(fastvideo_args.pipeline_config, "prefix", "") == "Cosmos25" + ) model = maybe_load_fsdp_model( model_cls=model_cls, init_params={"config": dit_config, "hf_config": hf_config}, @@ -589,6 +648,7 @@ class TransformerLoader(ComponentLoader): device=get_local_torch_device(), hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim, hsdp_shard_dim=fastvideo_args.hsdp_shard_dim, + strict=strict_load, cpu_offload=fastvideo_args.dit_cpu_offload, pin_cpu_memory=fastvideo_args.pin_cpu_memory, fsdp_inference=fastvideo_args.use_fsdp_inference, diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 072e70713..4459f5166 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -67,6 +67,7 @@ def maybe_load_fsdp_model( default_dtype: torch.dtype, param_dtype: torch.dtype, reduce_dtype: torch.dtype, + strict: bool = True, cpu_offload: bool = False, fsdp_inference: bool = False, output_dtype: torch.dtype | None = None, @@ -141,7 +142,7 @@ def maybe_load_fsdp_model( weight_iterator, device, default_dtype, - strict=True, + strict=strict, cpu_offload=cpu_offload, param_names_mapping=param_names_mapping_fn, ) @@ -293,6 +294,26 @@ def load_model_from_full_model_state_dict( for target_param_name, full_tensor in custom_param_sd.items(): meta_sharded_param = meta_sd.get(target_param_name) if meta_sharded_param is None: + # Some checkpoints include extra entries that are not part of the + # instantiated model's state_dict (e.g. `_extra_state` keys from + # some FSDP checkpoint formats). These can be safely skipped. + if (target_param_name.endswith("._extra_state") + or target_param_name.endswith("_extra_state")): + logger.warning( + "Skipping non-parameter checkpoint key: %s", + target_param_name, + ) + continue + + # For non-strict loads, treat this as an "unexpected key" and skip it + # (mirrors torch.nn.Module.load_state_dict(strict=False)). + if not strict: + logger.warning( + "Skipping unexpected checkpoint key (not present in model): %s", + target_param_name, + ) + continue + raise ValueError( f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect." ) diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index cbe2685a6..4fe9b1f39 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -30,8 +30,9 @@ _TEXT_TO_VIDEO_DIT_MODELS = { "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), "StepVideoModel": ("dits", "stepvideo", "StepVideoModel"), "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"), - "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"), - "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"), + "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"), + "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"), # Wrapper (Phase 1) + "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"), # Native (Phase 2) } _IMAGE_TO_VIDEO_DIT_MODELS = { @@ -50,6 +51,9 @@ _TEXT_ENCODER_MODELS = { "STEP1TextEncoder": ("encoders", "stepllm", "STEP1TextEncoder"), "BertModel": ("encoders", "clip", "CLIPTextModel"), "Qwen2_5_VLTextModel": ("encoders", "qwen2_5", "Qwen2_5_VLTextModel"), + "Reason1TextEncoder": ("encoders", "reason1", "Reason1TextEncoder"), + "Qwen2_5_VLForConditionalGeneration": + ("encoders", "reason1", "Reason1TextEncoder"), } _IMAGE_ENCODER_MODELS: dict[str, tuple] = { @@ -72,6 +76,8 @@ _SCHEDULERS = { "FlowMatchEulerDiscreteScheduler"), "UniPCMultistepScheduler": ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"), + "FlowUniPCMultistepScheduler": + ("schedulers", "scheduling_flow_unipc_multistep", "FlowUniPCMultistepScheduler"), "SelfForcingFlowMatchScheduler": ("schedulers", "scheduling_self_forcing_flow_match", "SelfForcingFlowMatchScheduler"), diff --git a/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py b/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py index ce111ddc5..b82921574 100644 --- a/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py +++ b/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py @@ -109,6 +109,9 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): sigmas = 1.0 - alphas sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32) + # Needed when final_sigmas_type == "sigma_min" (kept for compatibility). + self.alphas_cumprod = torch.from_numpy(alphas).to(dtype=torch.float32) + if not use_dynamic_shifting: # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution assert shift is not None @@ -171,6 +174,8 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): sigmas: list[float] | None = None, mu: float | None | None = None, shift: float | None | None = None, + use_karras_sigmas: bool | None = None, + use_kerras_sigma: bool | None = None, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -186,21 +191,44 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`" ) - if sigmas is None: - assert num_inference_steps is not None - sigmas = np.linspace(self.sigma_max, self.sigma_min, - num_inference_steps + - 1).copy()[:-1] # pyright: ignore + + # Cosmos official uses `use_kerras_sigma=True` and a specific EDM sigma schedule. + # Some external code uses the misspelling `use_kerras_sigma`; support both. + if use_karras_sigmas is None and use_kerras_sigma is not None: + use_karras_sigmas = use_kerras_sigma + + if use_karras_sigmas: + # Force to use the exact sigma used in official EDM sampler: + # sigma_max=200, sigma_min=0.01, rho=7 + sigma_max = 200.0 + sigma_min = 0.01 + rho = 7.0 + # Match the official Cosmos implementation: Karras/EDM schedule with + # `num_inference_steps + 1` points, then `final_sigmas_type="zero"` + # appends the terminal sigma (0.0). + ramp = np.arange(num_inference_steps + 1, + dtype=np.float32) / float(num_inference_steps) + min_inv_rho = sigma_min**(1 / rho) + max_inv_rho = sigma_max**(1 / rho) + sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho + # Convert EDM sigma to flow-matching sigma in [0, 1). + sigmas = sigmas / (1.0 + sigmas) + else: + if sigmas is None: + assert num_inference_steps is not None + sigmas = np.linspace(self.sigma_max, self.sigma_min, + num_inference_steps + + 1).copy()[:-1] # pyright: ignore if self.config.use_dynamic_shifting: assert mu is not None sigmas = self.time_shift(mu, 1.0, sigmas) # pyright: ignore else: - if shift is None: - shift = self.config.shift - assert isinstance(sigmas, np.ndarray) - sigmas = shift * sigmas / (1 + - (shift - 1) * sigmas) # pyright: ignore + if not use_karras_sigmas: + if shift is None: + shift = self.config.shift + assert isinstance(sigmas, np.ndarray) + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) # pyright: ignore if self.config.final_sigmas_type == "sigma_min": sigma_last = ((1 - self.alphas_cumprod[0]) / @@ -418,8 +446,12 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + # Numerical safety + eps = 1e-12 + lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log( + torch.clamp(sigma_t, min=eps)) + lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log( + torch.clamp(sigma_s0, min=eps)) h = lambda_t - lambda_s0 device = sample.device @@ -430,7 +462,8 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): si = self.step_index - i # pyright: ignore mi = model_output_list[-(i + 1)] alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) - lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log( + torch.clamp(sigma_si, min=eps)) rk = (lambda_si - lambda_s0) / h rks.append(rk) assert mi is not None @@ -563,8 +596,11 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + eps = 1e-12 + lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log( + torch.clamp(sigma_t, min=eps)) + lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log( + torch.clamp(sigma_s0, min=eps)) h = lambda_t - lambda_s0 device = this_sample.device @@ -575,7 +611,8 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): si = self.step_index - (i + 1) # pyright: ignore mi = model_output_list[-(i + 1)] alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) - lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log( + torch.clamp(sigma_si, min=eps)) rk = (lambda_si - lambda_s0) / h rks.append(rk) assert mi is not None diff --git a/fastvideo/models/vaes/cosmos25wanvae.py b/fastvideo/models/vaes/cosmos25wanvae.py new file mode 100644 index 000000000..98428af66 --- /dev/null +++ b/fastvideo/models/vaes/cosmos25wanvae.py @@ -0,0 +1,735 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +""" +Cosmos 2.5 / Wan2.1 VAE adapter. + +Why this exists: +- Cosmos2.5 uses a Wan2.1-style VAE, but the *diffusion model* operates in a + **normalized latent space**: + z_norm = (z - mean) / std + + Meanwhile, FastVideo's `AutoencoderKLWan` operates in the VAE's native latent + space (denormalized): + z = z_norm * std + mean + +This adapter provides a single, stable interface for FastVideo pipelines: +- `encode(x)` returns an object with `.mean` / `.sample()` / `.mode()` +- `decode(z)` returns a tensor in pixel space + +It also exposes flags used by pipeline stages to avoid double (de)normalization: +- `handles_latent_norm = True` -> stages should NOT normalize encoder latents +- `handles_latent_denorm = True` -> stages should NOT denormalize before decode +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + + +@dataclass +class _TensorLatentDist: + """Minimal distribution-like wrapper used by pipeline stages.""" + + mean: torch.Tensor + + def mode(self) -> torch.Tensor: + return self.mean + + def sample(self, generator: Any | None = None) -> torch.Tensor: # generator for API compatibility + # The official interface encodes deterministically; for compatibility we + # return the mean. (Stochastic posterior sampling isn't required for + # Cosmos2.5 inference.) + _ = generator + return self.mean + + +class Cosmos25WanVAEAdapter(nn.Module): + """ + Adapter that makes a Wan2.1-style VAE follow Cosmos2.5's latent contract: + - `encode()` returns **normalized** latents + - `decode()` expects **normalized** latents + """ + + # Pipeline stage hints (see latent_preparation.py / decoding.py / image_encoding.py) + handles_latent_norm: bool = True + handles_latent_denorm: bool = True + latent_norm_mode: str = "internal" # informational + + def __init__( + self, + inner: Any, + *, + latents_mean: Optional[torch.Tensor] = None, + latents_std: Optional[torch.Tensor] = None, + ) -> None: + super().__init__() + self.inner = inner + + # Preserve `config` when available; some pipeline utilities expect it. + self.config = getattr(inner, "config", None) + + # If not provided, try to derive from `config.latents_mean/std`. + cfg = self.config + if latents_mean is None and cfg is not None and hasattr(cfg, "latents_mean"): + latents_mean = torch.tensor(cfg.latents_mean, dtype=torch.float32).view(1, -1, 1, 1, 1) + if latents_std is None and cfg is not None and hasattr(cfg, "latents_std"): + latents_std = torch.tensor(cfg.latents_std, dtype=torch.float32).view(1, -1, 1, 1, 1) + + if latents_mean is None or latents_std is None: + raise RuntimeError( + "Cosmos25WanVAEAdapter requires latents_mean/latents_std (either passed explicitly or available on inner.config)." + ) + + # Register as buffers so `.to(...)` moves them with the module. + self.register_buffer("_latents_mean", latents_mean, persistent=False) + self.register_buffer("_latents_std", latents_std, persistent=False) + + def _to_latent_stats(self, like: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + mean = self._latents_mean.to(device=like.device, dtype=like.dtype) + std = self._latents_std.to(device=like.device, dtype=like.dtype) + return mean, std + + def get_latent_num_frames(self, num_pixel_frames: int) -> int: + # Keep parity with official interface. + if hasattr(self.inner, "get_latent_num_frames"): + return int(self.inner.get_latent_num_frames(num_pixel_frames)) + return 1 + (num_pixel_frames - 1) // 4 + + def encode(self, x: torch.Tensor) -> _TensorLatentDist: + """ + Returns *normalized* latents (Cosmos contract). + """ + enc_out = self.inner.encode(x) + + # Support common encoder output shapes: + # - DiagonalGaussianDistribution (FastVideo VAE): has `.mean` / `.sample()` / `.mode()` + # - diffusers EncoderOutput: has `.latent_dist` + # - raw tensor + if hasattr(enc_out, "latent_dist"): + dist = enc_out.latent_dist + z_mean = dist.mode() if hasattr(dist, "mode") else dist.mean + elif hasattr(enc_out, "mode") and hasattr(enc_out, "mean"): + z_mean = enc_out.mode() + elif isinstance(enc_out, torch.Tensor): + z_mean = enc_out + else: + attrs = [a for a in dir(enc_out) if not a.startswith("_")] + raise RuntimeError( + f"Unsupported VAE encoder output type: {type(enc_out)}. attrs={attrs}" + ) + + mean, std = self._to_latent_stats(z_mean) + z_norm = (z_mean - mean) / std + return _TensorLatentDist(z_norm) + + def decode(self, z: torch.Tensor) -> torch.Tensor: + """ + Expects *normalized* latents (Cosmos contract). + """ + mean, std = self._to_latent_stats(z) + z_denorm = z * std + mean + out = self.inner.decode(z_denorm) + return out.sample if hasattr(out, "sample") else out + + +# +# Official-like Wan2.1 VAE implementation (ported from cosmos_predict2 wan2pt1.py) +# ------------------------------------------------------------------------------- +# Motivation: +# - We already solved checkpoint *key mapping* and can load official weights. +# - Remaining output drift vs the official tokenizer is largely decoder-side. +# - FastVideo's `AutoencoderKLWan` uses a different temporal upsample path +# (`DupUp3D` + `first_chunk` slicing), while the official tokenizer uses +# `Resample(mode="upsample3d")` with a time-conv + interleave reshape. +# +# This section ports the core modules (CausalConv3d/Resample/etc.) so we can run +# a VAE that is behaviorally closer to the official implementation WITHOUT +# importing any official repo classes at runtime. +# + +CACHE_T = 2 + + +class Cosmos25CausalConv3d(nn.Conv3d): + """ + Official-like causal 3D convolution. + + Matches `CausalConv3d` in the official tokenizer: uses explicit F.pad and + supports a `cache_x` prefix for causal chunking. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + # padding order for F.pad: (W_left, W_right, H_left, H_right, T_left, T_right) + self._padding: tuple[int, ...] = ( + self.padding[2], + self.padding[2], + self.padding[1], + self.padding[1], + 2 * self.padding[0], + 0, + ) + self.padding = (0, 0, 0) + + def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor: + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + return super().forward(x) + + +class Cosmos25RMSNorm(nn.Module): + """Official-like RMS_norm (uses learnable gamma and optional bias).""" + + def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None: + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dim = 1 if self.channel_first else -1 + return F.normalize(x, dim=dim) * self.scale * self.gamma + self.bias + + +class Cosmos25Upsample(nn.Upsample): + """Official-like Upsample that is safe for bf16 (casts to fp32 internally).""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: # type: ignore[override] + return super().forward(x.float()).type_as(x) + + +class Cosmos25Resample(nn.Module): + """ + Official-like Resample used for both spatial and temporal up/downsampling. + """ + + def __init__(self, dim: int, mode: str) -> None: + assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d") + super().__init__() + self.dim = dim + self.mode = mode + + if mode == "upsample2d": + self.resample = nn.Sequential( + Cosmos25Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim // 2, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Cosmos25Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim // 2, 3, padding=1), + ) + self.time_conv = Cosmos25CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = Cosmos25CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + else: + self.resample = nn.Identity() + + def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor: + b, c, t, h, w = x.size() + + # Temporal upsample uses a time-conv and then interleaves frames. + if self.mode == "upsample3d" and feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = "Rep" + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep": + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep": + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + + if feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + + t = x.shape[2] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.resample(x) + x = rearrange(x, "(b t) c h w -> b c t h w", t=t) + + # Temporal downsample: time_conv consumes last-frame cache. + if self.mode == "downsample3d" and feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + return x + + +class Cosmos25ResidualBlock(nn.Module): + def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.0) -> None: + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.residual = nn.Sequential( + Cosmos25RMSNorm(in_dim, images=False), + nn.SiLU(), + Cosmos25CausalConv3d(in_dim, out_dim, 3, padding=1), + Cosmos25RMSNorm(out_dim, images=False), + nn.SiLU(), + nn.Dropout(dropout), + Cosmos25CausalConv3d(out_dim, out_dim, 3, padding=1), + ) + self.shortcut = Cosmos25CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + + def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor: + h = self.shortcut(x) + for layer in self.residual: + if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class Cosmos25AttentionBlock(nn.Module): + """Official-like causal self-attention with a single head.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.dim = dim + self.norm = Cosmos25RMSNorm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + nn.init.zeros_(self.proj.weight) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + identity = x + b, c, t, h, w = x.size() + x2 = rearrange(x, "b c t h w -> (b t) c h w") + x2 = self.norm(x2) + q, k, v = ( + self.to_qkv(x2) + .reshape(b * t, 1, c * 3, -1) + .permute(0, 1, 3, 2) + .contiguous() + .chunk(3, dim=-1) + ) + x2 = F.scaled_dot_product_attention(q, k, v) + x2 = x2.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + x2 = self.proj(x2) + x2 = rearrange(x2, "(b t) c h w-> b c t h w", t=t) + return x2 + identity + + +class Cosmos25Encoder3d(nn.Module): + def __init__( + self, + dim: int = 96, + z_dim: int = 32, + dim_mult: list[int] = [1, 2, 4, 4], + num_res_blocks: int = 2, + attn_scales: list[float] = [], + temperal_downsample: list[bool] = [False, True, True], + dropout: float = 0.0, + ) -> None: + super().__init__() + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + self.conv1 = Cosmos25CausalConv3d(3, dims[0], 3, padding=1) + + downsamples: list[nn.Module] = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + for _ in range(num_res_blocks): + downsamples.append(Cosmos25ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + downsamples.append(Cosmos25AttentionBlock(out_dim)) + in_dim = out_dim + + if i != len(dim_mult) - 1: + mode = "downsample3d" if temperal_downsample[i] else "downsample2d" + downsamples.append(Cosmos25Resample(out_dim, mode=mode)) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + self.middle = nn.Sequential( + Cosmos25ResidualBlock(out_dim, out_dim, dropout), + Cosmos25AttentionBlock(out_dim), + Cosmos25ResidualBlock(out_dim, out_dim, dropout), + ) + + self.head = nn.Sequential( + Cosmos25RMSNorm(out_dim, images=False), + nn.SiLU(), + Cosmos25CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor: + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) # type: ignore[misc] + else: + x = layer(x) # type: ignore[misc] + + for layer in self.middle: + if isinstance(layer, Cosmos25ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) # type: ignore[misc] + + for layer in self.head: + if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) # type: ignore[misc] + return x + + +class Cosmos25Decoder3d(nn.Module): + def __init__( + self, + dim: int = 96, + z_dim: int = 16, + dim_mult: list[int] = [1, 2, 4, 4], + num_res_blocks: int = 2, + attn_scales: list[float] = [], + temperal_upsample: list[bool] = [False, True, True], + dropout: float = 0.0, + ) -> None: + super().__init__() + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2 ** (len(dim_mult) - 2) + + self.conv1 = Cosmos25CausalConv3d(z_dim, dims[0], 3, padding=1) + self.middle = nn.Sequential( + Cosmos25ResidualBlock(dims[0], dims[0], dropout), + Cosmos25AttentionBlock(dims[0]), + Cosmos25ResidualBlock(dims[0], dims[0], dropout), + ) + + upsamples: list[nn.Module] = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + if i in (1, 2, 3): + in_dim = in_dim // 2 + for _ in range(num_res_blocks + 1): + upsamples.append(Cosmos25ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + upsamples.append(Cosmos25AttentionBlock(out_dim)) + in_dim = out_dim + if i != len(dim_mult) - 1: + mode = "upsample3d" if temperal_upsample[i] else "upsample2d" + upsamples.append(Cosmos25Resample(out_dim, mode=mode)) + scale *= 2.0 + self.upsamples = nn.Sequential(*upsamples) + + self.head = nn.Sequential( + Cosmos25RMSNorm(out_dim, images=False), + nn.SiLU(), + Cosmos25CausalConv3d(out_dim, 3, 3, padding=1), + ) + + def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor: + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if isinstance(layer, Cosmos25ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) # type: ignore[misc] + + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) # type: ignore[misc] + else: + x = layer(x) # type: ignore[misc] + + for layer in self.head: + if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) # type: ignore[misc] + return x + + +def _count_cosmos25_conv3d(model: nn.Module) -> int: + return sum(1 for m in model.modules() if isinstance(m, Cosmos25CausalConv3d)) + + +class Cosmos25WanVAE(nn.Module): + """ + A FastVideo-native copy of the *official-like* Wan2.1 VAE core. + + Key properties: + - Module naming matches official tokenizer (`encoder`, `decoder`, `conv1`, `conv2`) + so it can consume `tokenizer.pth` keys directly. + - `encode()` returns **normalized** latents and `decode()` expects **normalized** + latents (Cosmos2.5 contract), matching `Wan2pt1VAEInterface`. + """ + + handles_latent_norm: bool = True + handles_latent_denorm: bool = True + + def __init__( + self, + *, + device: torch.device | str = "cpu", + dtype: torch.dtype = torch.float32, + temporal_window: int = 4, + latents_mean: Optional[torch.Tensor] = None, + latents_std: Optional[torch.Tensor] = None, + ) -> None: + super().__init__() + + # Official hyperparams for Cosmos2.5 tokenizer (Wan2.1 VAE). + cfg = dict( + dim=96, + z_dim=16, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0, + temporal_window=temporal_window, + ) + self.z_dim = 16 + self.temporal_window = temporal_window + + self.encoder = Cosmos25Encoder3d( + dim=cfg["dim"], + z_dim=cfg["z_dim"] * 2, + dim_mult=cfg["dim_mult"], + num_res_blocks=cfg["num_res_blocks"], + attn_scales=cfg["attn_scales"], + temperal_downsample=cfg["temperal_downsample"], + dropout=cfg["dropout"], + ) + self.conv1 = Cosmos25CausalConv3d(self.z_dim * 2, self.z_dim * 2, 1) + self.conv2 = Cosmos25CausalConv3d(self.z_dim, self.z_dim, 1) + self.decoder = Cosmos25Decoder3d( + dim=cfg["dim"], + z_dim=cfg["z_dim"], + dim_mult=cfg["dim_mult"], + num_res_blocks=cfg["num_res_blocks"], + attn_scales=cfg["attn_scales"], + temperal_upsample=list(cfg["temperal_downsample"])[::-1], + dropout=cfg["dropout"], + ) + + # Default Cosmos2.5 latent stats (shared with configs). + if latents_mean is None: + latents_mean = torch.tensor( + [ + -0.7571, + -0.7089, + -0.9113, + 0.1075, + -0.1745, + 0.9653, + -0.1517, + 1.5508, + 0.4134, + -0.0715, + 0.5517, + -0.3632, + -0.1922, + -0.9497, + 0.2503, + -0.2921, + ], + dtype=torch.float32, + ).view(1, 16, 1, 1, 1) + if latents_std is None: + latents_std = torch.tensor( + [ + 2.8184, + 1.4541, + 2.3275, + 2.6558, + 1.2196, + 1.7708, + 2.6052, + 2.0743, + 3.2687, + 2.1526, + 2.8652, + 1.5579, + 1.6382, + 1.1253, + 2.8251, + 1.9160, + ], + dtype=torch.float32, + ).view(1, 16, 1, 1, 1) + + self.register_buffer("_latents_mean", latents_mean, persistent=False) + self.register_buffer("_latents_std", latents_std, persistent=False) + + self.to(device=device, dtype=dtype) + self.clear_cache() + + def clear_cache(self) -> None: + # Decoder cache + self._conv_num = _count_cosmos25_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map: list[Any] = [None] * self._conv_num + # Encoder cache + self._enc_conv_num = _count_cosmos25_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map: list[Any] = [None] * self._enc_conv_num + + def _scale(self, like: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + mean = self._latents_mean.to(device=like.device, dtype=like.dtype) + std = self._latents_std.to(device=like.device, dtype=like.dtype) + return mean, 1.0 / std + + def _i0_encode(self, x: torch.Tensor) -> torch.Tensor: + return self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) + + def _i0_decode(self, x: torch.Tensor) -> torch.Tensor: + return self.decoder(x[:, :, 0:1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) + + def encode(self, x: torch.Tensor) -> _TensorLatentDist: + """ + Encode to *normalized* latents (Cosmos contract). + """ + self.clear_cache() + t = x.shape[2] + iters = 1 + (t - 1) // self.temporal_window + + for i in range(iters): + self._enc_conv_idx = [0] + if i == 0: + out = self._i0_encode(x) + else: + out_ = self.encoder( + x[:, :, 1 + self.temporal_window * (i - 1) : 1 + self.temporal_window * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], 2) + + if (t - 1) % self.temporal_window: + self._enc_conv_idx = [0] + out_ = self.encoder( + x[:, :, 1 + self.temporal_window * (iters - 1) :, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], 2) + + mu, _log_var = self.conv1(out).chunk(2, dim=1) + mean, inv_std = self._scale(mu) + z_norm = (mu - mean) * inv_std + self.clear_cache() + return _TensorLatentDist(z_norm) + + def decode(self, latent: torch.Tensor) -> torch.Tensor: + """ + Decode from *normalized* latents (Cosmos contract). + """ + self.clear_cache() + mean, inv_std = self._scale(latent) + z = latent / inv_std + mean # z = z_norm * std + mean + + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self._i0_decode(x) + else: + out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) + + self.clear_cache() + return out + + # --- Interface helpers (match official Wan2pt1VAEInterface) --- + def get_latent_num_frames(self, num_pixel_frames: int) -> int: + return 1 + (int(num_pixel_frames) - 1) // 4 + + def get_pixel_num_frames(self, num_latent_frames: int) -> int: + return (int(num_latent_frames) - 1) * 4 + 1 + + @property + def spatial_compression_factor(self) -> int: + return 8 + + @property + def temporal_compression_factor(self) -> int: + return 4 + + @property + def latent_ch(self) -> int: + return 16 + + + + + + + + + diff --git a/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py b/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py new file mode 100644 index 000000000..c21720100 --- /dev/null +++ b/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos 2.5 pipeline entry (staged pipeline).""" + +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, + Cosmos25DenoisingStage, + Cosmos25LatentPreparationStage, + DecodingStage, InputValidationStage, + Cosmos25TextEncodingStage, + Cosmos25TimestepPreparationStage) + +logger = init_logger(__name__) + + +class Cosmos2_5Pipeline(ComposedPipelineBase): + """Cosmos 2.5 video generation pipeline.""" + + _required_config_modules = [ + "text_encoder", "tokenizer", "vae", "transformer", "scheduler", + "safety_checker" + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + logger.info("Creating Cosmos 2.5 pipeline stages...") + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + self.add_stage( + stage_name="prompt_encoding_stage", + stage=Cosmos25TextEncodingStage( + text_encoder=self.get_module("text_encoder"), ), + ) + + self.add_stage(stage_name="conditioning_stage", + stage=ConditioningStage()) + + self.add_stage(stage_name="timestep_preparation_stage", + stage=Cosmos25TimestepPreparationStage( + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="latent_preparation_stage", + stage=Cosmos25LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + vae=self.get_module("vae"))) + + self.add_stage(stage_name="denoising_stage", + stage=Cosmos25DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + logger.info("Cosmos 2.5 pipeline stages created") + + +# Entry point for pipeline registry +EntryClass = Cosmos2_5Pipeline diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index 4728fa8dd..a2f325dd2 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -29,6 +29,7 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = { "HunyuanVideoPipeline": "hunyuan", "HunyuanVideo15Pipeline": "hunyuan15", "Cosmos2VideoToWorldPipeline": "cosmos", + "Cosmos2_5Pipeline": "cosmos", "MatrixGamePipeline": "matrixgame", "MatrixGameCausalDMDPipeline": "matrixgame", "LongCatPipeline": "longcat", diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index e5116c34d..a8136bafe 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -10,7 +10,8 @@ from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.pipelines.stages.causal_denoising import CausalDMDDenosingStage from fastvideo.pipelines.stages.conditioning import ConditioningStage from fastvideo.pipelines.stages.decoding import DecodingStage -from fastvideo.pipelines.stages.denoising import (CosmosDenoisingStage, +from fastvideo.pipelines.stages.denoising import (Cosmos25DenoisingStage, + CosmosDenoisingStage, DenoisingStage, DmdDenoisingStage) from fastvideo.pipelines.stages.encoding import EncodingStage @@ -19,14 +20,16 @@ from fastvideo.pipelines.stages.image_encoding import ( ImageVAEEncodingStage, VideoVAEEncodingStage, Hy15ImageEncodingStage) from fastvideo.pipelines.stages.input_validation import InputValidationStage from fastvideo.pipelines.stages.latent_preparation import ( - CosmosLatentPreparationStage, LatentPreparationStage) + Cosmos25LatentPreparationStage, CosmosLatentPreparationStage, + LatentPreparationStage) from fastvideo.pipelines.stages.matrixgame_denoising import ( MatrixGameCausalDenoisingStage) from fastvideo.pipelines.stages.stepvideo_encoding import ( StepvideoPromptEncodingStage) -from fastvideo.pipelines.stages.text_encoding import TextEncodingStage +from fastvideo.pipelines.stages.text_encoding import (Cosmos25TextEncodingStage, + TextEncodingStage) from fastvideo.pipelines.stages.timestep_preparation import ( - TimestepPreparationStage) + Cosmos25TimestepPreparationStage, TimestepPreparationStage) # LongCat stages from fastvideo.pipelines.stages.longcat_video_vae_encoding import LongCatVideoVAEEncodingStage @@ -37,14 +40,17 @@ __all__ = [ "PipelineStage", "InputValidationStage", "TimestepPreparationStage", + "Cosmos25TimestepPreparationStage", "LatentPreparationStage", "CosmosLatentPreparationStage", + "Cosmos25LatentPreparationStage", "ConditioningStage", "DenoisingStage", "DmdDenoisingStage", "CausalDMDDenosingStage", "MatrixGameCausalDenoisingStage", "CosmosDenoisingStage", + "Cosmos25DenoisingStage", "EncodingStage", "DecodingStage", "ImageEncodingStage", @@ -54,6 +60,7 @@ __all__ = [ "ImageVAEEncodingStage", "VideoVAEEncodingStage", "TextEncodingStage", + "Cosmos25TextEncodingStage", "StepvideoPromptEncodingStage", # LongCat stages "LongCatVideoVAEEncodingStage", diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index 7de5d8960..7450ff226 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/pipelines/stages/decoding.py @@ -51,39 +51,42 @@ class DecodingStage(PipelineStage): return result def _denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor: - # denormalization for MatrixGame VAE - # z = z * std + mean during decode - if (hasattr(self.vae.config, 'latents_mean') - and hasattr(self.vae.config, 'latents_std')): - # Convert config values to tensors - latents_mean = torch.tensor(self.vae.config.latents_mean, + """Convert normalized latents into the VAE's expected latent space.""" + # Some VAEs handle latent (de)normalization internally. + if bool(getattr(self.vae, "handles_latent_denorm", False)): + return latents + + cfg = getattr(self.vae, "config", None) + + # MatrixGame-style: z = z * std + mean + if (cfg is not None and hasattr(cfg, "latents_mean") + and hasattr(cfg, "latents_std")): + latents_mean = torch.tensor(cfg.latents_mean, device=latents.device, dtype=latents.dtype).view( 1, -1, 1, 1, 1) - - latents_std = torch.tensor(self.vae.config.latents_std, + latents_std = torch.tensor(cfg.latents_std, device=latents.device, dtype=latents.dtype).view( 1, -1, 1, 1, 1) + return latents * latents_std + latents_mean - # Apply denormalization: z = z * std + mean - latents = latents * latents_std + latents_mean - elif hasattr(self.vae, 'scaling_factor'): - # Standard VAE scaling + # Diffusers-style: scaling_factor (+ optional shift_factor) + if hasattr(self.vae, "scaling_factor"): if isinstance(self.vae.scaling_factor, torch.Tensor): latents = latents / self.vae.scaling_factor.to( latents.device, latents.dtype) else: latents = latents / self.vae.scaling_factor - # Apply shifting if needed - if (hasattr(self.vae, "shift_factor") - and self.vae.shift_factor is not None): + if hasattr(self.vae, + "shift_factor") and self.vae.shift_factor is not None: if isinstance(self.vae.shift_factor, torch.Tensor): - latents += self.vae.shift_factor.to(latents.device, - latents.dtype) + latents = latents + self.vae.shift_factor.to( + latents.device, latents.dtype) else: - latents += self.vae.shift_factor + latents = latents + self.vae.shift_factor + return latents @torch.no_grad() @@ -273,4 +276,4 @@ class DecodingStage(PipelineStage): del pipeline.modules["vae"] fastvideo_args.model_loaded["vae"] = False - return batch + return batch \ No newline at end of file diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 7669b26df..a6c235b5d 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -1018,6 +1018,145 @@ class CosmosDenoisingStage(DenoisingStage): return result +class Cosmos25DenoisingStage(CosmosDenoisingStage): + """Denoising stage for Cosmos 2.5 DiT (expects 1D/2D timestep, not 5D).""" + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + 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 + + extra_step_kwargs = self.prepare_extra_func_kwargs( + self.scheduler.step, + { + "generator": batch.generator, + "eta": batch.eta + }, + ) + + if hasattr(self.transformer, 'module'): + transformer_dtype = next(self.transformer.module.parameters()).dtype + else: + transformer_dtype = next(self.transformer.parameters()).dtype + target_dtype = transformer_dtype + autocast_enabled = (target_dtype != torch.float32 + ) and not fastvideo_args.disable_autocast + + latents = batch.latents + if latents is None: + raise ValueError( + "latents must be provided for Cosmos25DenoisingStage") + guidance_scale = batch.guidance_scale + + # Use timesteps prepared by Cosmos25TimestepPreparationStage when available. + if batch.timesteps is None: + self.scheduler.set_timesteps(batch.num_inference_steps, + device=latents.device) + timesteps = self.scheduler.timesteps + else: + timesteps = batch.timesteps.to(latents.device) + + # Match official behavior: pass fps as a tensor. + fps_val = batch.fps if isinstance(batch.fps, int | float) else 24 + fps_tensor = torch.tensor([fps_val], + device=latents.device, + dtype=target_dtype) + + # Cosmos2.5 denoises a 4D latent (C,T,H,W) and the scheduler.step expects (B,C,T,H,W). + latents_4d = latents[0] + + # Masks from latent prep stage + condition_mask = batch.cond_mask.to(target_dtype) if hasattr( + batch, 'cond_mask') else None + padding_mask = batch.padding_mask.to(target_dtype) if hasattr( + batch, 'padding_mask') else None + if condition_mask is None: + _, t, h, w = latents_4d.shape + condition_mask = torch.zeros(1, + 1, + t, + h, + w, + device=latents.device, + dtype=target_dtype) + if padding_mask is None: + _, _, h, w = latents_4d.shape + padding_mask = torch.ones(1, + 1, + h, + w, + device=latents.device, + dtype=target_dtype) + + # Cosmos2.5 timestep scaling (see compare_pipelines.py): t * 0.001 + timestep_scale = 0.001 + + with self.progress_bar(total=len(timesteps)) as progress_bar: + for i, t in enumerate(timesteps): + t_val = float(t) + timestep_val = t_val * timestep_scale + timestep = torch.tensor([[timestep_val]], + device=latents.device, + dtype=target_dtype) + + model_hidden_states = latents_4d.unsqueeze(0) + + with ( + set_forward_context(current_timestep=int(t_val), + attn_metadata=None, + forward_batch=batch), + torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled), + ): + cond_v = self.transformer( + hidden_states=model_hidden_states.to(target_dtype), + encoder_hidden_states=batch.prompt_embeds[0].to( + target_dtype), + timestep=timestep, + fps=fps_tensor, + condition_mask=condition_mask, + padding_mask=padding_mask, + return_dict=False, + )[0] + + if batch.do_classifier_free_guidance and batch.negative_prompt_embeds: + uncond_v = self.transformer( + hidden_states=model_hidden_states.to(target_dtype), + encoder_hidden_states=batch. + negative_prompt_embeds[0].to(target_dtype), + timestep=timestep, + fps=fps_tensor, + condition_mask=condition_mask, + padding_mask=padding_mask, + return_dict=False, + )[0] + v = uncond_v + guidance_scale * (cond_v - uncond_v) + else: + v = cond_v + + prev = self.scheduler.step(v.unsqueeze(0), + t, + latents_4d.unsqueeze(0), + **extra_step_kwargs, + return_dict=False)[0] + latents_4d = prev.squeeze(0) + + progress_bar.update() + + batch.latents = latents_4d.unsqueeze(0) + return batch + + class DmdDenoisingStage(DenoisingStage): """ Denoising stage for DMD. diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/pipelines/stages/latent_preparation.py index a38515b46..4af23ed2d 100644 --- a/fastvideo/pipelines/stages/latent_preparation.py +++ b/fastvideo/pipelines/stages/latent_preparation.py @@ -5,6 +5,7 @@ Latent preparation stage for diffusion pipelines. from typing import Any +import numpy as np import torch from diffusers.utils.torch_utils import randn_tensor @@ -427,6 +428,210 @@ class CosmosLatentPreparationStage(PipelineStage): return batch + +class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage): + """Latent preparation for Cosmos 2.5 DiT input conventions.""" + + @staticmethod + def _arch_invariant_randn( + shape: tuple[int, ...], + *, + seed: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + """Architecture-invariant RNG (matches cosmos_predict2.misc.arch_invariant_rand).""" + rng = np.random.RandomState(seed) + arr = rng.standard_normal(shape).astype(np.float32) + return torch.from_numpy(arr).to(device=device, dtype=dtype) + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + # Differences vs `CosmosLatentPreparationStage`: channel convention, seed usage, + # and `padding_mask` for concat_padding_mask=True. + + # Determine batch size + if isinstance(batch.prompt, list): + batch_size = len(batch.prompt) + elif batch.prompt is not None: + batch_size = 1 + else: + batch_size = batch.prompt_embeds[0].shape[0] + + batch_size *= batch.num_videos_per_prompt + + # Match `compare_pipelines.py`: initialize noise in fp32, then run the + # denoising computation in bf16. + dtype = torch.float32 + device = get_local_torch_device() + generator = batch.generator + latents = batch.latents + num_frames = batch.num_frames + height = batch.height + width = batch.width + + if height is None or width is None: + raise ValueError("Height and width must be provided") + + vae_scale_factor_spatial = 8 + vae_scale_factor_temporal = 4 + + latent_height = height // 8 + latent_width = width // vae_scale_factor_spatial + num_latent_frames = (num_frames - 1) // vae_scale_factor_temporal + 1 + + # Cosmos 2.5 convention: transformer in_channels == latent channels + num_channels_latents = self.transformer.config.in_channels + + shape = (batch_size, num_channels_latents, num_latent_frames, + latent_height, latent_width) + + init_latents = None + conditioning_latents = None + video = None + + if hasattr(batch, 'video') and batch.video is not None: + video = batch.video + elif hasattr(batch, 'pil_image') and batch.pil_image is not None: + vae_scale_factor_spatial = 8 + image_processor = ImageProcessor( + vae_scale_factor=vae_scale_factor_spatial) + processed_image = image_processor.preprocess( + batch.pil_image, height, width) + video = processed_image.unsqueeze(2) + video = video.to(device=device, dtype=torch.bfloat16) + elif hasattr( + batch, + 'preprocessed_image') and batch.preprocessed_image is not None: + if isinstance(batch.preprocessed_image, torch.Tensor): + if batch.preprocessed_image.dim() == 4: + video = batch.preprocessed_image.unsqueeze(2) + elif batch.preprocessed_image.dim() == 5: + video = batch.preprocessed_image + else: + logger.info( + "CosmosLatentPreparationStage - No video input sources found") + + if video is not None: + num_cond_frames = video.size(2) + if num_cond_frames >= num_frames: + num_cond_latent_frames = (num_frames - + 1) // vae_scale_factor_temporal + 1 + video = video[:, :, -num_frames:] + else: + num_cond_latent_frames = (num_cond_frames - + 1) // vae_scale_factor_temporal + 1 + num_padding_frames = num_frames - num_cond_frames + last_frame = video[:, :, -1:] + padding = last_frame.repeat(1, 1, num_padding_frames, 1, 1) + video = torch.cat([video, padding], dim=2) + + if self.vae is not None: + self.vae = self.vae.to(device) + self.vae = self.vae.to(dtype=video.dtype) + + def retrieve_latents( + encoder_output: Any, + generator: Any | None = None) -> torch.Tensor: + if hasattr(encoder_output, "latent_dist"): + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + elif hasattr(encoder_output, "sample"): + return encoder_output.sample(generator) + elif isinstance(encoder_output, torch.Tensor): + return encoder_output + else: + attrs = [ + attr for attr in dir(encoder_output) + if not attr.startswith('_') + ] + raise AttributeError( + f"Could not access latents of provided encoder_output. Available attributes: {attrs}" + ) + + if isinstance(generator, list): + init_latents = [ + retrieve_latents(self.vae.encode(video[i].unsqueeze(0)), + generator=torch.Generator( + device="cpu").manual_seed(100)) + for i in range(batch_size) + ] + else: + init_latents = [ + retrieve_latents( + self.vae.encode(vid.unsqueeze(0)), + torch.Generator(device="cpu").manual_seed(100)) + for vid in video + ] + + init_latents = torch.cat(init_latents, dim=0).to(dtype) + + cfg = getattr(self.vae, "config", None) + if (not bool(getattr(self.vae, "handles_latent_norm", False)) + and cfg is not None and hasattr(cfg, 'latents_mean') + and hasattr(cfg, 'latents_std')): + latents_mean = torch.tensor(cfg.latents_mean).view( + 1, cfg.z_dim, 1, 1, 1).to(device, dtype) + latents_std = torch.tensor(cfg.latents_std).view( + 1, cfg.z_dim, 1, 1, 1).to(device, dtype) + init_latents = (init_latents - latents_mean + ) / latents_std * self.scheduler.sigma_data + + conditioning_latents = init_latents + self.vae.to("cpu") + else: + num_cond_latent_frames = 0 + + if latents is None: + seed = int(batch.seed if batch.seed is not None else 0) + # Use arch-invariant RNG to match Cosmos2.5 reference sampling. + latents_fp32 = self._arch_invariant_randn(shape, + seed=seed, + device=device, + dtype=torch.float32) + latents = latents_fp32.to(torch.bfloat16) + else: + # If latents are supplied, keep compute dtype consistent with Cosmos sampling. + latents = latents.to(device=device, dtype=torch.bfloat16) + + # Cosmos2.5 starts from unit Gaussian noise (no extra sigma_max scaling). + + padding_shape = (batch_size, 1, num_latent_frames, latent_height, + latent_width) + ones_padding = latents.new_ones(padding_shape) + zeros_padding = latents.new_zeros(padding_shape) + + cond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1) + cond_indicator[:, :, :num_cond_latent_frames] = 1.0 + cond_mask = cond_indicator * ones_padding + ( + 1 - cond_indicator) * zeros_padding + + uncond_indicator = None + uncond_mask = None + if batch.do_classifier_free_guidance: + uncond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1) + uncond_indicator[:, :, :num_cond_latent_frames] = 1.0 + uncond_mask = uncond_indicator * ones_padding + ( + 1 - uncond_indicator) * zeros_padding + + # Cosmos 2.5 requires a spatial padding mask when concat_padding_mask=True + padding_mask = latents.new_ones(batch_size, 1, latent_height, + latent_width) + + batch.latents = latents + batch.raw_latent_shape = latents.shape + batch.conditioning_latents = conditioning_latents + batch.cond_indicator = cond_indicator + batch.uncond_indicator = uncond_indicator + batch.cond_mask = cond_mask + batch.uncond_mask = uncond_mask + batch.padding_mask = padding_mask + return batch + def adjust_video_length(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> int: """ diff --git a/fastvideo/pipelines/stages/text_encoding.py b/fastvideo/pipelines/stages/text_encoding.py index 612749bdc..74b115ff0 100644 --- a/fastvideo/pipelines/stages/text_encoding.py +++ b/fastvideo/pipelines/stages/text_encoding.py @@ -328,3 +328,69 @@ class TextEncodingStage(PipelineStage): lambda x: not batch.do_classifier_free_guidance or V. list_of_tensors_with_min_dims(x, 2)) return result + + +class Cosmos25TextEncodingStage(PipelineStage): + """Cosmos 2.5 text encoding stage. + + Cosmos 2.5 uses Reason1 (Qwen2.5-VL) and relies on the encoder's + `compute_text_embeddings_online()`. + """ + + def __init__(self, text_encoder) -> None: + super().__init__() + self.text_encoder = text_encoder + + @torch.no_grad() + def forward(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> ForwardBatch: + assert batch.prompt is not None + prompts = [batch.prompt] if isinstance(batch.prompt, + str) else batch.prompt + + encoder = self.text_encoder + if not hasattr(encoder, "compute_text_embeddings_online"): + raise RuntimeError( + "Cosmos25TextEncodingStage requires text_encoder.compute_text_embeddings_online()" + ) + + with set_forward_context(current_timestep=0, attn_metadata=None): + prompt_embeds = encoder.compute_text_embeddings_online( + {"text": prompts}, "text") + + batch.prompt_embeds = [prompt_embeds] + + if batch.do_classifier_free_guidance: + neg = batch.negative_prompt + neg_prompts = ([neg] * + len(prompts)) if isinstance(neg, str) else neg + with set_forward_context(current_timestep=0, attn_metadata=None): + neg_embeds = encoder.compute_text_embeddings_online( + {"text": neg_prompts}, "text") + batch.negative_prompt_embeds = [neg_embeds] + else: + batch.negative_prompt_embeds = [] + + return batch + + def verify_input(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("prompt", batch.prompt, V.string_or_list_strings) + result.add_check( + "negative_prompt", + batch.negative_prompt, + lambda x: + (not batch.do_classifier_free_guidance) or isinstance(x, str), + ) + return result + + def verify_output(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("prompt_embeds", batch.prompt_embeds, + V.list_of_tensors_min_dims(2)) + 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 diff --git a/fastvideo/pipelines/stages/timestep_preparation.py b/fastvideo/pipelines/stages/timestep_preparation.py index ca79c858b..c1074cd41 100644 --- a/fastvideo/pipelines/stages/timestep_preparation.py +++ b/fastvideo/pipelines/stages/timestep_preparation.py @@ -123,3 +123,32 @@ class TimestepPreparationStage(PipelineStage): result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.with_dims(1)]) return result + + +class Cosmos25TimestepPreparationStage(TimestepPreparationStage): + """Cosmos 2.5 timestep preparation with scheduler-specific kwargs.""" + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + scheduler = self.scheduler + device = get_local_torch_device() + num_inference_steps = batch.num_inference_steps + + extra_kwargs: dict = {} + sig = inspect.signature(scheduler.set_timesteps) + if "shift" in sig.parameters: + extra_kwargs["shift"] = fastvideo_args.pipeline_config.flow_shift + # Prefer the canonical diffusers kwarg name if available. + if "use_karras_sigmas" in sig.parameters: + extra_kwargs["use_karras_sigmas"] = True + elif "use_kerras_sigma" in sig.parameters: + extra_kwargs["use_kerras_sigma"] = True + + scheduler.set_timesteps(num_inference_steps, + device=device, + **extra_kwargs) + batch.timesteps = scheduler.timesteps + return batch diff --git a/scripts/checkpoint_conversion/pt_to_safetensors.py b/scripts/checkpoint_conversion/pt_to_safetensors.py new file mode 100644 index 000000000..25f0909c9 --- /dev/null +++ b/scripts/checkpoint_conversion/pt_to_safetensors.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""Convert a PyTorch checkpoint (.pt) to a safetensors file.""" + +import argparse +from pathlib import Path + +import torch +from safetensors.torch import save_file + + +def convert_pt_to_safetensors( + input_path: str, + output_path: str, + key: str | None = None, + force: bool = False, + skip_patterns: list[str] | None = None, +): + input_path = Path(input_path) + output_path = Path(output_path) + + if not input_path.exists(): + raise FileNotFoundError(f"Input file not found: {input_path}") + + if output_path.exists() and not force: + raise FileExistsError( + f"Output file already exists: {output_path}. Use --force to overwrite." + ) + + checkpoint = torch.load(input_path, map_location="cpu") + + state_dict: dict[str, torch.Tensor] + if isinstance(checkpoint, dict): + if key is not None: + if key not in checkpoint: + raise KeyError(f"Key {key!r} not found in checkpoint.") + state_dict = checkpoint[key] + else: + for k in ("state_dict", "model_state_dict", "model", "ema"): + if k in checkpoint: + state_dict = checkpoint[k] + break + else: + state_dict = checkpoint + else: + state_dict = checkpoint + + if not isinstance(state_dict, dict): + raise TypeError(f"Expected a dict state_dict, got {type(state_dict)}") + + if skip_patterns: + state_dict = { + k: v + for k, v in state_dict.items() + if not any(pat in k for pat in skip_patterns) + } + + output_path.parent.mkdir(parents=True, exist_ok=True) + save_file(state_dict, str(output_path)) + + +def main(): + parser = argparse.ArgumentParser( + description="Convert a PyTorch checkpoint (.pt) to safetensors." + ) + parser.add_argument( + "input", + type=str, + help="Path to input .pt checkpoint file" + ) + parser.add_argument( + "output", + type=str, + help="Path to output .safetensors file" + ) + parser.add_argument( + "--key", + type=str, + default=None, + help="Optional key to extract from checkpoint dict (e.g. 'model_state_dict')" + ) + parser.add_argument( + "--force", + action="store_true", + help="Overwrite output file if it exists" + ) + parser.add_argument( + "--skip-pattern", + action="append", + dest="skip_patterns", + help="Parameter name patterns to skip (can be used multiple times)" + ) + + args = parser.parse_args() + + convert_pt_to_safetensors( + args.input, + args.output, + args.key, + args.force, + args.skip_patterns + ) + + +if __name__ == "__main__": + main() +