diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 6def015f76..2b353b7654 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -492,6 +492,12 @@ surfaces: trajectory_type: request.extensions.gen3c.trajectory_type movement_distance: request.extensions.gen3c.movement_distance camera_rotation: request.extensions.gen3c.camera_rotation + svi_ref_pad_num: request.extensions.svi.ref_pad_num + svi_ref_pad_cfg: request.extensions.svi.ref_pad_cfg + svi_num_clips: request.extensions.svi.num_clips + svi_num_motion_frames: request.extensions.svi.num_motion_frames + svi_seed_stride: request.extensions.svi.seed_stride + svi_clip_prompts: request.extensions.svi.clip_prompts prompt_attention_mask: request.extensions.hyworld.prompt_attention_mask negative_attention_mask: request.extensions.hyworld.negative_attention_mask camera_states: request.extensions.hunyuangamecraft.camera_states diff --git a/examples/inference/basic/basic_svi_i2v.py b/examples/inference/basic/basic_svi_i2v.py new file mode 100644 index 0000000000..61aa8beaf8 --- /dev/null +++ b/examples/inference/basic/basic_svi_i2v.py @@ -0,0 +1,94 @@ +from fastvideo import VideoGenerator + +MODEL_VARIANT = "shot" + +VARIANT_CONFIG = { + "shot": { + "lora_path": "vita-video-gen/svi-model/version-1.0/svi-shot.safetensors", + "image_url": + "https://raw.githubusercontent.com/vita-epfl/Stable-Video-Infinity/main/data/toy_test/shot/frame.jpg", + "prompts": [ + ("A sleek white motor yacht speeds across the turquoise blue sea, " + "leaving a dramatic wake of white foam behind it under a clear blue sky."), + ], + "num_clips": 2, + "num_motion_frames": 1, + "ref_pad_num": -1, + "height": 448, + }, + "film": { + "lora_path": "vita-video-gen/svi-model/version-1.0/svi-film-opt-10212025.safetensors", + "image_url": + "https://raw.githubusercontent.com/vita-epfl/Stable-Video-Infinity/main/data/toy_test/film/frame.jpg", + "prompts": [ + ("A Siamese kitten rests snugly inside a straw hat, its head slightly tilted " + "as it gazes curiously to the side."), + ("The Siamese kitten decides to explore the room and jumps out of the hat " + "onto the soft carpet below."), + ], + "num_clips": 2, + "num_motion_frames": 5, + "ref_pad_num": 0, + "height": 480, + }, + "tom": { + "lora_path": "vita-video-gen/svi-model/version-1.0/svi-tom.safetensors", + "image_url": + "https://raw.githubusercontent.com/vita-epfl/Stable-Video-Infinity/main/data/toy_test/tom/frame.png", + "prompts": [ + ("A static shot of the bright 1950s kitchen, turquoise cabinets and a chrome " + "sink glinting; Tom cat hovers over the counter, yellow eyes narrowed, while " + "Jerry mouse stands defiantly in a tiny milk puddle near a stack of purple plates."), + ("Close-up on Tom cat’s face: a wicked smirk creases his white muzzle; his black " + "brows angle into a sharp V as he crooks one claw toward Jerry mouse like a " + "menacing metronome."), + ], + "num_clips": 2, + "num_motion_frames": 1, + "ref_pad_num": 0, + "height": 560, + }, +} + +OUTPUT_PATH = "video_samples_svi" + + +def main(): + config = VARIANT_CONFIG[MODEL_VARIANT] + + generator = VideoGenerator.from_pretrained( + "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", + override_pipeline_cls_name="WanSVIImageToVideoPipeline", + lora_path=config["lora_path"], + lora_nickname=f"svi-{MODEL_VARIANT}", + num_gpus=1, + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=True, + flow_shift=5.0, + # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer + pin_cpu_memory=True, + ) + + generator.generate_video( + prompt=config["prompts"][0], + image_path=config["image_url"], + output_path=OUTPUT_PATH, + save_video=True, + height=config["height"], + width=832, + num_frames=81, + fps=24, + num_inference_steps=50, + guidance_scale=5.0, + seed=0, + svi_num_clips=config["num_clips"], + svi_clip_prompts=None if MODEL_VARIANT == "shot" else config["prompts"], + svi_num_motion_frames=config["num_motion_frames"], + svi_seed_stride=42, + svi_ref_pad_num=config["ref_pad_num"], + ) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/api/sampling_param.py b/fastvideo/api/sampling_param.py index 9db1de75e4..0fbb3896c7 100644 --- a/fastvideo/api/sampling_param.py +++ b/fastvideo/api/sampling_param.py @@ -112,6 +112,14 @@ class SamplingParam: movement_distance: float | None = None camera_rotation: str | None = None + # Stable-Video-Infinity image conditioning + svi_ref_pad_num: int | None = None # -1 = tile ref; 0 = zero pad; k>0 = ref for first k slots + svi_ref_pad_cfg: bool = False # widens y-mask to len(first_frames) instead of {first frame only} + svi_num_clips: int = 1 # >1 enables motion-frame chaining + svi_num_motion_frames: int = 1 # tail frames carried from clip K to clip K+1 (1=Shot/Tom, 5=Film) + svi_seed_stride: int = 42 + svi_clip_prompts: list[str] | None = None # optional one-to-one prompt list for multi-clip generation + # LTX-2 multi-modal CFG and STG. # Class-level defaults match the *distilled* LTX-2 schedule # (mirrors ``FastVideo-internal/.../LTX2DistilledSamplingParam``): diff --git a/fastvideo/pipelines/basic/wan/presets.py b/fastvideo/pipelines/basic/wan/presets.py index 7234bef129..8675eb9d12 100644 --- a/fastvideo/pipelines/basic/wan/presets.py +++ b/fastvideo/pipelines/basic/wan/presets.py @@ -120,6 +120,30 @@ }, ) +WAN_SVI_I2V_14B_480P = InferencePreset( + name="wan_svi_i2v_14b_480p", + version=1, + model_family="wan", + description="Stable-Video-Infinity multi-clip I2V on Wan 2.1 14B at 480p", + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "height": 448, + "width": 832, + "num_frames": 81, + "fps": 24, + "guidance_scale": 5.0, + "num_inference_steps": 50, + "negative_prompt": _NEGATIVE_PROMPT_EN, + # SVI-Shot defaults; override per variant in user kwargs. + "svi_num_clips": 1, + "svi_num_motion_frames": 1, + "svi_seed_stride": 42, + "svi_ref_pad_num": -1, + "svi_clip_prompts": None, + }, +) + # ------------------------------------------------------------------- # Wan 2.2 presets # ------------------------------------------------------------------- @@ -352,6 +376,7 @@ WAN_T2V_14B, WAN_I2V_14B_480P, WAN_I2V_14B_720P, + WAN_SVI_I2V_14B_480P, WAN_2_2_T2V_A14B, WAN_2_2_I2V_A14B, WAN_FUN_1_3B_INP, diff --git a/fastvideo/pipelines/basic/wan/wan_svi_i2v_pipeline.py b/fastvideo/pipelines/basic/wan/wan_svi_i2v_pipeline.py new file mode 100644 index 0000000000..ab11873aa8 --- /dev/null +++ b/fastvideo/pipelines/basic/wan/wan_svi_i2v_pipeline.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Wan I2V pipeline variant for Stable-Video-Infinity multi-clip inference. + +This module contains an implementation of the SVI-flavored Wan I2V +pipeline using the modular pipeline architecture. +""" + +import dataclasses + +import PIL.Image +import torch + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from fastvideo.pipelines.basic.wan.wan_i2v_pipeline import WanImageToVideoPipeline +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + +# isort: off +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, DenoisingStage, ImageEncodingStage, + InputValidationStage, LatentPreparationStage, SVIImageVAEEncodingStage, + TextEncodingStage, TimestepPreparationStage) +# isort: on + +logger = init_logger(__name__) + + +def _tensor_to_pil_list(frames: torch.Tensor) -> list[PIL.Image.Image]: + """Convert a video tensor to a list of PIL frames.""" + arr = (frames.detach().to(torch.float32).cpu().clamp(0, 1) * 255.0).to(torch.uint8) + arr = arr.permute(1, 2, 3, 0).numpy() + return [PIL.Image.fromarray(a) for a in arr] + + +def _validate_multiclip_frames(num_motion: int, num_frames: int) -> None: + """num_motion must be < num_frames, else follow-up clips stitch to empty.""" + if num_motion >= num_frames: + raise ValueError(f"svi_num_motion_frames ({num_motion}) must be smaller than num_frames ({num_frames}) " + "for multi-clip generation; otherwise stitched follow-up clips would be empty.") + + +def _stitch_clip_outputs(clip_outputs: list[torch.Tensor], num_motion: int) -> torch.Tensor: + """Stitch clips using the SVI overlap convention.""" + clips = [video[:, :, :-num_motion] for video in clip_outputs[:-1]] + return torch.cat([*clips, clip_outputs[-1]], dim=2) + + +def _resolve_clip_prompts( + prompt: str | list[str] | None, + clip_prompts: list[str] | None, + num_clips: int, +) -> list[str]: + """Return exactly one non-empty prompt per generated clip.""" + if clip_prompts is None: + if not isinstance(prompt, str) or not prompt.strip(): + raise ValueError("SVI requires a non-empty primary prompt") + return [prompt] * num_clips + + if len(clip_prompts) != num_clips: + raise ValueError(f"svi_clip_prompts must contain exactly svi_num_clips entries " + f"({num_clips}), but got {len(clip_prompts)}") + if any(not isinstance(item, str) or not item.strip() for item in clip_prompts): + raise ValueError("svi_clip_prompts entries must be non-empty strings") + return list(clip_prompts) + + +def _clip_seed(seed: int | None, clip_idx: int, seed_stride: int) -> int: + """Match the SVI reference by varying diffusion noise per clip.""" + if seed is None: + raise ValueError("SVI requires a seed") + return int(seed) + clip_idx * seed_stride + + +class WanSVIImageToVideoPipeline(WanImageToVideoPipeline): + """ + Pipeline for Stable-Video-Infinity multi-clip I2V generation on Wan 2.1. + """ + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + """Set up pipeline stages with proper dependency injection.""" + self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler( + num_train_timesteps=1000, + shift=fastvideo_args.pipeline_config.flow_shift, + ) + + self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage()) + self.add_stage( + stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + ) + if (self.get_module("image_encoder") is not None and self.get_module("image_processor") is not None): + self.add_stage( + stage_name="image_encoding_stage", + stage=ImageEncodingStage( + image_encoder=self.get_module("image_encoder"), + image_processor=self.get_module("image_processor"), + ), + ) + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")), + ) + self.add_stage( + stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + ), + ) + self.add_stage( + stage_name="image_latent_preparation_stage", + stage=SVIImageVAEEncodingStage(vae=self.get_module("vae")), + ) + self.add_stage( + stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + transformer_2=self.get_module("transformer_2"), + scheduler=self.get_module("scheduler"), + ), + ) + self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))) + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if not self.post_init_called: + self.post_init() + + n_steps = int(batch.num_inference_steps) + if batch.sigmas is None: + batch.sigmas = torch.linspace(1.0, 0.0, n_steps + 1)[:-1].tolist() + + num_clips = max(1, int(batch.svi_num_clips or 1)) + num_motion = max(1, int(batch.svi_num_motion_frames or 1)) + seed_stride = int(batch.svi_seed_stride) + + prompts = _resolve_clip_prompts(batch.prompt, batch.svi_clip_prompts, num_clips) + num_frames = int(batch.num_frames) if batch.num_frames is not None else 0 + if num_clips > 1: + _validate_multiclip_frames(num_motion, num_frames) + + # Resolve image_path before constructing the first motion window. + self.input_validation_stage(batch, fastvideo_args) + assert isinstance(batch.pil_image, PIL.Image.Image) + assert isinstance(batch.height, int) and isinstance(batch.width, int) + reference_frame = batch.pil_image.resize((batch.width, batch.height)) + random_ref = batch.svi_random_ref_frame or reference_frame + motion_frames = batch.svi_first_frames or [reference_frame] + + clip_outputs: list[torch.Tensor] = [] + for clip_idx in range(num_clips): + clip_batch = dataclasses.replace( + batch, + prompt=prompts[clip_idx], + seed=_clip_seed(batch.seed, clip_idx, seed_stride), + pil_image=motion_frames[0], + image_path=None, + svi_first_frames=motion_frames, + svi_random_ref_frame=random_ref, + prompt_embeds=[], + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_attention_mask=None, + clip_embedding_pos=None, + clip_embedding_neg=None, + image_embeds=[], + preprocessed_image=None, + latents=None, + image_latent=None, + noise_pred=None, + output=None, + timesteps=None, + timestep=None, + step_index=None, + is_prompt_processed=False, + ) + for stage in self.stages: + clip_batch = stage(clip_batch, fastvideo_args) + + assert clip_batch.output is not None + clip_outputs.append(clip_batch.output) + logger.info( + "SVI clip %d/%d generated with seed=%d, frames shape=%s", + clip_idx + 1, + num_clips, + clip_batch.seed, + clip_batch.output.shape, + ) + + if clip_idx + 1 < num_clips: + tail = clip_batch.output[0, :, -num_motion:, :, :] + motion_frames = _tensor_to_pil_list(tail) + + batch.output = _stitch_clip_outputs(clip_outputs, num_motion) + return batch + + +EntryClass = WanSVIImageToVideoPipeline diff --git a/fastvideo/pipelines/lora_pipeline.py b/fastvideo/pipelines/lora_pipeline.py index 5ee38cbc49..b993772c94 100644 --- a/fastvideo/pipelines/lora_pipeline.py +++ b/fastvideo/pipelines/lora_pipeline.py @@ -2,6 +2,7 @@ from collections import defaultdict from collections.abc import Hashable from contextlib import nullcontext +import re from typing import Any from collections.abc import Generator @@ -29,6 +30,14 @@ logger = init_logger(__name__) +def _normalize_lora_key(name: str) -> str: + """Normalize known DiffSynth/PEFT wrappers without rewriting module names.""" + name = name.removeprefix("diffusion_model.") + name = name.removeprefix("pipe.dit.") + name = re.sub(r"(\.lora_[AB])\.default(?=\.weight$)", r"\1", name) + return name.removesuffix(".weight") + + def _get_hook_ctx(module: nn.Module | None): if module is None: return nullcontext() @@ -324,8 +333,7 @@ def set_lora_adapter(self, to_merge_params: defaultdict[Hashable, dict[Any, Any]] = (defaultdict(dict)) for name, weight in lora_state_dict.items(): # Extract weights (lora_A, lora_B, and lora_alpha) - name = name.replace("diffusion_model.", "") - name = name.replace(".weight", "") + name = _normalize_lora_key(name) if "lora_alpha" in name: # Store alpha with minimal mapping - same processing as lora_A/lora_B diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 4859f7929b..52b5f8f005 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -156,6 +156,16 @@ class ForwardBatch: movement_distance: float | None = None camera_rotation: str | None = None + # Stable-Video-Infinity image conditioning + svi_first_frames: list[PIL.Image.Image] | None = None # Motion frames forwarded across clips + svi_random_ref_frame: PIL.Image.Image | None = None # Reference frame for ref-padded slots + svi_ref_pad_num: int | None = None # -1=tile ref; 0=zero pad; k>0=ref for first k, zero rest + svi_ref_pad_cfg: bool = False # widens y-mask to len(first_frames) instead of frame 0 only + svi_num_clips: int = 1 # >1 enables motion-frame chaining + svi_num_motion_frames: int = 1 # Tail frames carried clip K -> K+1 (1=Shot/Tom, 5=Film) + svi_seed_stride: int = 42 + svi_clip_prompts: list[str] | None = None # One prompt per clip; primary prompt is reused when omitted + # Latent dimensions height_latents: list[int] | int | None = None width_latents: list[int] | int | None = None diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index ae4b701f6d..f63a59e144 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -18,8 +18,9 @@ from fastvideo.pipelines.stages.image_encoding import (ImageEncodingStage, MatrixGame2ImageEncodingStage, MatrixGame2ImageVAEEncodingStage, MatrixGame3ImageVAEEncodingStage, RefImageEncodingStage, - ImageVAEEncodingStage, VideoVAEEncodingStage, - Hy15ImageEncodingStage, HYWorldImageEncodingStage) + ImageVAEEncodingStage, SVIImageVAEEncodingStage, + VideoVAEEncodingStage, Hy15ImageEncodingStage, + HYWorldImageEncodingStage) from fastvideo.pipelines.stages.gamecraft_image_encoding import (GameCraftImageVAEEncodingStage) from fastvideo.pipelines.stages.input_validation import InputValidationStage from fastvideo.pipelines.stages.latent_preparation import (Cosmos25LatentPreparationStage, CosmosLatentPreparationStage, @@ -95,6 +96,7 @@ "ImageVAEEncodingStage", "VideoVAEEncodingStage", "GameCraftImageVAEEncodingStage", + "SVIImageVAEEncodingStage", "TextEncodingStage", "Cosmos25TextEncodingStage", # LongCat stages diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py index 1e2be0443b..1431c3e4c3 100644 --- a/fastvideo/pipelines/stages/image_encoding.py +++ b/fastvideo/pipelines/stages/image_encoding.py @@ -7,6 +7,7 @@ - RefImageEncodingStage: Encodes reference image for Wan2.1 control pipeline - ImageVAEEncodingStage: Encodes images to latent space using VAE for I2V generation - VideoVAEEncodingStage: Encodes videos to latent space using VAE for V2V and control tasks +- SVIImageVAEEncodingStage: Encodes motion frames + reference padding for Stable-Video-Infinity """ import PIL @@ -866,6 +867,115 @@ def forward( return batch +class SVIImageVAEEncodingStage(ImageVAEEncodingStage): + """ + Stage for encoding motion frames and reference padding for Stable-Video-Infinity. + """ + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if fastvideo_args.mode != ExecutionMode.INFERENCE: + raise NotImplementedError("SVIImageVAEEncodingStage only supports inference mode") + + assert isinstance(batch.height, int) and isinstance(batch.width, int) + assert isinstance(batch.num_frames, int) + height, width, num_frames = batch.height, batch.width, batch.num_frames + + first_frames = batch.svi_first_frames + if not first_frames: + assert isinstance(batch.pil_image, PIL.Image.Image) + first_frames = [batch.pil_image] + random_ref_frame = batch.svi_random_ref_frame + if random_ref_frame is None: + random_ref_frame = first_frames[0] + ref_pad_num = batch.svi_ref_pad_num if batch.svi_ref_pad_num is not None else 0 + ref_pad_cfg = batch.svi_ref_pad_cfg + + num_condition_frames = len(first_frames) + remaining_frames = num_frames - num_condition_frames + if remaining_frames < 0: + raise ValueError(f"num_frames={num_frames} smaller than len(first_frames)={num_condition_frames}") + + device = get_local_torch_device() + self.vae = self.vae.to(device) + vae_scale = self.vae.spatial_compression_ratio + latent_height = height // vae_scale + latent_width = width // vae_scale + temporal_compression = self.vae.temporal_compression_ratio + + msk = torch.ones(1, num_frames, latent_height, latent_width, device=device, dtype=torch.float32) + if ref_pad_cfg: + msk[:, num_condition_frames:] = 0 + else: + msk[:, 1:] = 0 + msk = torch.concat( + [torch.repeat_interleave(msk[:, 0:1], repeats=temporal_compression, dim=1), msk[:, 1:]], + dim=1, + ) + msk = msk.view(1, msk.shape[1] // temporal_compression, temporal_compression, latent_height, latent_width) + msk = msk.transpose(1, 2)[0] # (4, T_lat, H/8, W/8) + + condition_frames = [ + self.preprocess(frame, vae_scale_factor=vae_scale, height=height, + width=width).to(device=device, dtype=torch.float32).unsqueeze(2) for frame in first_frames + ] + vae_input_condition = torch.cat(condition_frames, dim=2) + + if ref_pad_num < -1: + raise ValueError(f"Unsupported ref_pad_num={ref_pad_num} (expected -1, 0, or positive int)") + vae_input_pad = vae_input_condition.new_zeros( + vae_input_condition.shape[0], + vae_input_condition.shape[1], + remaining_frames, + height, + width, + ) + if ref_pad_num != 0 and remaining_frames: + pad_count = remaining_frames if ref_pad_num == -1 else min(ref_pad_num, remaining_frames) + ref = self.preprocess(random_ref_frame, vae_scale_factor=vae_scale, height=height, + width=width).to(device=device, dtype=torch.float32).unsqueeze(2) + vae_input_pad[:, :, :pad_count] = ref + + video_condition = torch.cat([vae_input_condition, vae_input_pad], dim=2) + assert video_condition.shape[2] == num_frames, video_condition.shape + + vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] + vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast + with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): + if fastvideo_args.pipeline_config.vae_tiling: + self.vae.enable_tiling() + if not vae_autocast_enabled: + video_condition = video_condition.to(vae_dtype) + encoder_output = self.vae.encode(video_condition) + + if batch.generator is None: + raise ValueError("Generator must be provided") + latent = self.retrieve_latents(encoder_output, batch.generator, sample_mode="argmax") + + if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None: + shift = self.vae.shift_factor + if isinstance(shift, torch.Tensor): + shift = shift.to(latent.device, latent.dtype) + latent = latent - shift + + scaling = self.vae.scaling_factor + if isinstance(scaling, torch.Tensor): + scaling = scaling.to(latent.device, latent.dtype) + latent = latent * scaling + + mask_batched = msk.unsqueeze(0).to(latent.device, latent.dtype) + batch.image_latent = torch.concat([mask_batched, latent], dim=1) + + if hasattr(self, "maybe_free_model_hooks"): + self.maybe_free_model_hooks() + self.vae.to("cpu") + return batch + + def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("image_latent", batch.image_latent, [V.is_tensor, V.with_dims(5)]) + return result + + class MatrixGame3ImageVAEEncodingStage(ImageVAEEncodingStage): def preprocess( diff --git a/fastvideo/tests/inference/lora/test_maybe_download_lora.py b/fastvideo/tests/inference/lora/test_maybe_download_lora.py new file mode 100644 index 0000000000..afc3f919e0 --- /dev/null +++ b/fastvideo/tests/inference/lora/test_maybe_download_lora.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Path-routing tests for fastvideo.utils.maybe_download_lora.""" +from __future__ import annotations + +from contextlib import contextmanager + +import huggingface_hub + +import fastvideo.utils as fv_utils +from fastvideo.pipelines.lora_pipeline import _normalize_lora_key +from fastvideo.utils import maybe_download_lora + + +def test_triple_slash_downloads_single_file(monkeypatch): + """org/repo/sub/file.safetensors -> hf_hub_download(repo_id=org/repo, filename=sub/file).""" + calls: list[dict] = [] + + def fake_hf_hub_download(*, repo_id, filename, local_dir=None, **kwargs): + calls.append({"repo_id": repo_id, "filename": filename, "local_dir": local_dir}) + return f"/cache/{repo_id}/{filename}" + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", fake_hf_hub_download) + + lock_events: list[str] = [] + + @contextmanager + def fake_lock(path): + lock_events.append(f"enter:{path}") + yield + lock_events.append(f"exit:{path}") + + monkeypatch.setattr(fv_utils, "get_lock", fake_lock) + + result = maybe_download_lora("vita-video-gen/svi-model/version-1.0/svi-shot.safetensors") + + assert len(calls) == 1 + assert calls[0]["repo_id"] == "vita-video-gen/svi-model" + assert calls[0]["filename"] == "version-1.0/svi-shot.safetensors" + assert result == "/cache/vita-video-gen/svi-model/version-1.0/svi-shot.safetensors" + assert lock_events == [ + "enter:vita-video-gen/svi-model/version-1.0/svi-shot.safetensors", + "exit:vita-video-gen/svi-model/version-1.0/svi-shot.safetensors", + ] + + +def test_triple_slash_keeps_only_first_two_segments_as_repo(monkeypatch): + """A deeper nesting still maps repo_id to the first two segments.""" + captured: dict = {} + + def fake_hf_hub_download(*, repo_id, filename, local_dir=None, **kwargs): + captured["repo_id"] = repo_id + captured["filename"] = filename + return "ok" + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", fake_hf_hub_download) + + maybe_download_lora("org/repo/a/b/c/weights.safetensors") + + assert captured["repo_id"] == "org/repo" + assert captured["filename"] == "a/b/c/weights.safetensors" + + +def test_local_file_short_circuits(tmp_path, monkeypatch): + """An existing local .safetensors file is returned verbatim, no download.""" + + def boom(*args, **kwargs): + raise AssertionError("hf_hub_download must not be called for a local file") + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", boom) + + f = tmp_path / "version-1.0" / "svi-shot.safetensors" + f.parent.mkdir(parents=True) + f.write_bytes(b"\x00") + + assert maybe_download_lora(str(f)) == str(f) + + +def test_plain_repo_id_falls_through_to_repo_download(monkeypatch): + """A two-segment HF id (no .safetensors suffix) does NOT hit the single-file branch.""" + + def boom(*args, **kwargs): + raise AssertionError("two-segment repo id must not take the single-file branch") + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", boom) + monkeypatch.setattr(fv_utils, "maybe_download_model", lambda *a, **k: "/cache/org/repo") + monkeypatch.setattr(fv_utils, "_best_guess_weight_name", lambda *a, **k: "adapter.safetensors") + + result = maybe_download_lora("org/repo") + + assert result == "/cache/org/repo/adapter.safetensors" + + +def test_normalize_lora_key_only_strips_peft_default_adapter_segment(): + assert _normalize_lora_key("pipe.dit.blocks.0.attn.lora_A.default.weight") == ( + "blocks.0.attn.lora_A" + ) + assert _normalize_lora_key("blocks.default.attn.lora_A.weight") == ( + "blocks.default.attn.lora_A" + ) + + +def test_normalize_lora_key_only_strips_pipe_dit_prefix(): + assert _normalize_lora_key("module.pipe.dit.attn.lora_B.weight") == ( + "module.pipe.dit.attn.lora_B" + ) diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index b972458cfc..d15190a6ff 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -367,7 +367,8 @@ def run_unit_test(): run_test( "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ " "./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ " - "./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/worker/ " + "./fastvideo/tests/stages/ ./fastvideo/tests/pipelines/ ./fastvideo/tests/ops/ " + "./fastvideo/tests/worker/ " "./fastvideo/tests/training/test_trackers.py " "./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py ./fastvideo/tests/modal/test_pr_test.py " "--ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py " diff --git a/fastvideo/tests/pipelines/test_svi_multiclip_stitch.py b/fastvideo/tests/pipelines/test_svi_multiclip_stitch.py new file mode 100644 index 0000000000..9b1dd52f8f --- /dev/null +++ b/fastvideo/tests/pipelines/test_svi_multiclip_stitch.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for SVI multi-clip inference.""" +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from PIL import Image + +import fastvideo.pipelines.basic.wan.wan_svi_i2v_pipeline as svi_pipeline +from fastvideo.pipelines.basic.wan.wan_svi_i2v_pipeline import ( + WanSVIImageToVideoPipeline, + _clip_seed, + _resolve_clip_prompts, + _stitch_clip_outputs, + _validate_multiclip_frames, +) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +def _clip(num_frames: int, fill: float) -> torch.Tensor: + # (B=1, C=3, T=num_frames, H=2, W=2) + return torch.full((1, 3, num_frames, 2, 2), fill, dtype=torch.float32) + + +def test_stitch_drops_motion_tail_from_non_final_clips(): + num_motion = 2 + clips = [_clip(9, 0.0), _clip(9, 1.0), _clip(9, 2.0)] + + out = _stitch_clip_outputs(clips, num_motion) + + expected_t = 9 + (9 - num_motion) + (9 - num_motion) + assert out.shape == (1, 3, expected_t, 2, 2) + + assert torch.all(out[:, :, :7] == 0.0) + assert torch.all(out[:, :, 7:7 + 7] == 1.0) + assert torch.all(out[:, :, 7 + 7:] == 2.0) + + +def test_stitch_single_clip_is_identity(): + clip = _clip(5, 3.0) + out = _stitch_clip_outputs([clip], num_motion=2) + assert torch.equal(out, clip) + + +def test_stitch_num_motion_one_drops_single_frame(): + clips = [_clip(4, 0.0), _clip(4, 1.0)] + out = _stitch_clip_outputs(clips, num_motion=1) + assert out.shape[2] == 4 + 3 + + +def test_validate_rejects_motion_ge_frames(): + with pytest.raises(ValueError, match="must be smaller than num_frames"): + _validate_multiclip_frames(num_motion=5, num_frames=5) + with pytest.raises(ValueError, match="must be smaller than num_frames"): + _validate_multiclip_frames(num_motion=8, num_frames=5) + + +def test_validate_accepts_motion_lt_frames(): + # Should not raise. + _validate_multiclip_frames(num_motion=1, num_frames=81) + _validate_multiclip_frames(num_motion=5, num_frames=6) + + +def test_clip_seed_offsets_each_chunk(): + assert [_clip_seed(0, idx, 42) for idx in range(3)] == [0, 42, 84] + assert [_clip_seed(7, idx, 0) for idx in range(3)] == [7, 7, 7] + + +def test_svi_scheduler_uses_official_shift(monkeypatch): + monkeypatch.setattr(svi_pipeline, "DenoisingStage", lambda **_kwargs: SimpleNamespace()) + pipeline = object.__new__(WanSVIImageToVideoPipeline) + pipeline.modules = {} + pipeline._stages = [] + pipeline._stage_name_mapping = {} + + args = SimpleNamespace(pipeline_config=SimpleNamespace(flow_shift=5.0)) + pipeline.create_pipeline_stages(args) + + assert pipeline.modules["scheduler"].shift == 5.0 + + +def test_resolve_clip_prompts_reuses_primary_prompt_by_default(): + assert _resolve_clip_prompts("one prompt", None, 3) == ["one prompt"] * 3 + + +def test_resolve_clip_prompts_accepts_exact_per_clip_prompts(): + prompts = ["first", "second"] + assert _resolve_clip_prompts("fallback", prompts, 2) == prompts + + +def test_resolve_clip_prompts_rejects_wrong_length_or_empty_entries(): + with pytest.raises(ValueError, match="exactly svi_num_clips"): + _resolve_clip_prompts("fallback", ["only one"], 2) + with pytest.raises(ValueError, match="non-empty strings"): + _resolve_clip_prompts("fallback", ["first", ""], 2) + + +def test_single_clip_uses_resized_svi_reference(): + pipeline = object.__new__(WanSVIImageToVideoPipeline) + pipeline.post_init_called = True + pipeline.input_validation_stage = lambda batch, _args: batch + observed = [] + + def fake_stage(batch, _args): + observed.append((batch.pil_image.size, batch.svi_first_frames[0].size, + batch.svi_random_ref_frame.size)) + batch.output = _clip(4, 0.5) + return batch + + pipeline._stages = [fake_stage] + batch = ForwardBatch( + data_type="i2v", + prompt="prompt", + seed=0, + height=4, + width=6, + num_frames=4, + num_inference_steps=2, + svi_num_clips=1, + pil_image=Image.new("RGB", (2, 2)), + ) + + output = pipeline.forward(batch, SimpleNamespace()) + + assert output.output.shape == (1, 3, 4, 2, 2) + assert observed == [((6, 4), (6, 4), (6, 4))] + + +def test_multiclip_forward_uses_motion_head_and_keeps_padding_reference(): + pipeline = object.__new__(WanSVIImageToVideoPipeline) + pipeline.post_init_called = True + pipeline.input_validation_stage = lambda batch, _args: batch + + observed = [] + + def fake_stage(batch, _args): + assert isinstance(batch.pil_image, Image.Image) + assert batch.svi_first_frames + assert isinstance(batch.svi_random_ref_frame, Image.Image) + observed.append(( + batch.prompt, + batch.seed, + batch.pil_image.getpixel((0, 0)), + batch.svi_first_frames[0].getpixel((0, 0)), + batch.svi_random_ref_frame.getpixel((0, 0)), + )) + values = [0.0, 0.25, 0.5, 0.75] if len(observed) == 1 else [0.1, 0.2, 0.3, 0.4] + batch.output = torch.tensor(values).view(1, 1, 4, 1, 1).expand(1, 3, 4, 2, 2) + return batch + + pipeline._stages = [fake_stage] + batch = ForwardBatch( + data_type="i2v", + prompt="fallback", + seed=10, + height=2, + width=2, + num_frames=4, + num_inference_steps=2, + svi_num_clips=2, + svi_num_motion_frames=2, + svi_seed_stride=42, + svi_clip_prompts=["first", "second"], + pil_image=Image.new("RGB", (2, 2), color=(7, 7, 7)), + ) + + output = pipeline.forward(batch, SimpleNamespace()) + + assert output.output is not None + assert output.output.shape == (1, 3, 6, 2, 2) + assert observed == [ + ("first", 10, (7, 7, 7), (7, 7, 7), (7, 7, 7)), + ("second", 52, (127, 127, 127), (127, 127, 127), (7, 7, 7)), + ] diff --git a/fastvideo/tests/stages/test_svi_image_vae_encoding.py b/fastvideo/tests/stages/test_svi_image_vae_encoding.py new file mode 100644 index 0000000000..2fa905bb5f --- /dev/null +++ b/fastvideo/tests/stages/test_svi_image_vae_encoding.py @@ -0,0 +1,169 @@ +from unittest.mock import MagicMock + +import numpy as np +import PIL.Image +import pytest +import torch + +from fastvideo.fastvideo_args import ExecutionMode, FastVideoArgs +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.image_encoding import SVIImageVAEEncodingStage + + +def _diffsynth_preprocess(img: PIL.Image.Image) -> torch.Tensor: + """Reference (diffsynth) preprocess: uint8 * (2/255) - 1, permute to CHW.""" + arr = np.array(img, dtype=np.float32) * (2 / 255) - 1 + return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0) + + +def _reference_pre_vae( + first_frames: list[PIL.Image.Image], + random_ref_frame: PIL.Image.Image, + num_frames: int, + height: int, + width: int, + ref_pad_num: int, + ref_pad_cfg: bool, + vae_scale: int, + temporal_compression: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Mirror of encode_images_adaptive mask + video_condition construction.""" + num_condition_frames = len(first_frames) + remaining_frames = num_frames - num_condition_frames + random_ref = _diffsynth_preprocess(random_ref_frame) + + msk = torch.ones(1, num_frames, height // vae_scale, width // vae_scale, dtype=torch.float32) + if ref_pad_cfg: + msk[:, num_condition_frames:] = 0 + else: + msk[:, 1:] = 0 + msk = torch.concat( + [torch.repeat_interleave(msk[:, 0:1], repeats=temporal_compression, dim=1), msk[:, 1:]], + dim=1, + ) + msk = msk.view(1, msk.shape[1] // temporal_compression, temporal_compression, height // vae_scale, + width // vae_scale) + msk = msk.transpose(1, 2)[0] + + if len(first_frames) > 1: + frame_tensors = [_diffsynth_preprocess(f) for f in first_frames] + vae_input_condition = torch.cat(frame_tensors, dim=0).permute(1, 0, 2, 3) + else: + vae_input_condition = _diffsynth_preprocess(first_frames[0]).transpose(0, 1) + + if ref_pad_num == 0: + vae_input_pad = torch.zeros(3, remaining_frames, height, width, dtype=torch.float32) + elif ref_pad_num > 0: + pad_imgs = [random_ref.transpose(0, 1)] * ref_pad_num + if remaining_frames > ref_pad_num: + pad_imgs += [torch.zeros(3, 1, height, width, dtype=torch.float32)] * (remaining_frames - ref_pad_num) + vae_input_pad = torch.cat(pad_imgs, dim=1) + elif ref_pad_num == -1: + vae_input_pad = random_ref.transpose(0, 1).repeat(1, remaining_frames, 1, 1) + else: + raise ValueError(ref_pad_num) + + return msk, torch.concat([vae_input_condition, vae_input_pad], dim=1) + + +class _CaptureVAE: + """Mock VAE that records its encode() input and returns a zero latent.""" + + spatial_compression_ratio = 8 + temporal_compression_ratio = 4 + scaling_factor = 1.0 + shift_factor = None + + def __init__(self): + self.captured_input: torch.Tensor | None = None + + def to(self, *_a, **_kw): + return self + + def enable_tiling(self): + pass + + def encode(self, x: torch.Tensor): + self.captured_input = x.detach().clone() + t_lat = (x.shape[2] - 1) // self.temporal_compression_ratio + 1 + h_lat = x.shape[3] // self.spatial_compression_ratio + w_lat = x.shape[4] // self.spatial_compression_ratio + latent = torch.zeros(x.shape[0], 16, t_lat, h_lat, w_lat, device=x.device, dtype=torch.float32) + out = MagicMock() + out.sample = lambda _g: latent + 1 + out.mode = lambda: latent + return out + + +def _make_pil(h: int, w: int, seed: int) -> PIL.Image.Image: + arr = np.random.default_rng(seed).integers(0, 256, size=(h, w, 3), dtype=np.uint8) + return PIL.Image.fromarray(arr) + + +def _make_fastvideo_args() -> FastVideoArgs: + args = MagicMock(spec=FastVideoArgs) + args.mode = ExecutionMode.INFERENCE + args.disable_autocast = True + args.pipeline_config = MagicMock() + args.pipeline_config.vae_precision = "fp32" + args.pipeline_config.vae_tiling = False + return args + + +@pytest.mark.parametrize( + "num_motion, ref_pad_num, ref_pad_cfg", + [ + (1, -1, False), + (5, 0, False), + (1, 3, False), + (5, -1, True), + ], +) +def test_pre_vae_parity(num_motion: int, ref_pad_num: int, ref_pad_cfg: bool): + height, width, num_frames = 48, 96, 17 + + first_frames = [_make_pil(height, width, seed=i) for i in range(num_motion)] + random_ref = _make_pil(height, width, seed=999) + + expected_msk, expected_video = _reference_pre_vae( + first_frames=first_frames, + random_ref_frame=random_ref, + num_frames=num_frames, + height=height, + width=width, + ref_pad_num=ref_pad_num, + ref_pad_cfg=ref_pad_cfg, + vae_scale=8, + temporal_compression=4, + ) + + vae = _CaptureVAE() + stage = SVIImageVAEEncodingStage(vae=vae) # type: ignore[arg-type] + + batch = ForwardBatch( + data_type="video", + generator=torch.Generator().manual_seed(0), + height=height, + width=width, + num_frames=num_frames, + pil_image=first_frames[0], + svi_first_frames=first_frames, + svi_random_ref_frame=random_ref, + svi_ref_pad_num=ref_pad_num, + svi_ref_pad_cfg=ref_pad_cfg, + ) + stage.forward(batch, _make_fastvideo_args()) + + assert vae.captured_input is not None + # Stage feeds (1, 3, T, H, W); reference produces (3, T, H, W). + # .cpu() keeps the comparison device-agnostic when get_local_torch_device() picks CUDA. + actual_video = vae.captured_input.squeeze(0).cpu() + assert actual_video.shape == expected_video.shape + # FastVideo preprocess is (uint8/255) -> 2x-1; diffsynth is uint8*(2/255) - 1. Same math, ~1 ULP apart. + torch.testing.assert_close(actual_video, expected_video, atol=2e-7, rtol=1e-5) + + assert batch.image_latent is not None + actual_msk = batch.image_latent[0, :4].cpu() + assert actual_msk.shape == expected_msk.shape + torch.testing.assert_close(actual_msk, expected_msk, atol=0.0, rtol=0.0) + assert torch.count_nonzero(batch.image_latent[:, 4:]) == 0 diff --git a/fastvideo/utils.py b/fastvideo/utils.py index c0f66cd17e..74b431b5b4 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -591,7 +591,7 @@ def maybe_download_lora(model_name_or_path: str, local_dir: str | None = None, d model_name_or_path: Local path or Hugging Face Hub model ID local_dir: Local directory to save the model download: Whether to download the model from Hugging Face Hub - + Returns: Local path to the model """ @@ -600,6 +600,17 @@ def maybe_download_lora(model_name_or_path: str, local_dir: str | None = None, d if os.path.isfile(model_name_or_path): return model_name_or_path + # Handle "org/repo//.safetensors" by downloading the single + # specified file, for HF repos that ship several LoRA variants side by side. + parts = model_name_or_path.split("/") + if (len(parts) >= 3 and not model_name_or_path.startswith(("/", ".")) + and model_name_or_path.endswith(".safetensors")): + from huggingface_hub import hf_hub_download + repo_id = "/".join(parts[:2]) + filename = "/".join(parts[2:]) + with get_lock(model_name_or_path): + return hf_hub_download(repo_id=repo_id, filename=filename, local_dir=local_dir) + local_path = maybe_download_model(model_name_or_path, local_dir, download) weight_name = _best_guess_weight_name(model_name_or_path, file_extension=".safetensors")