diff --git a/examples/inference/basic/basic_turbodiffusion_i2v.py b/examples/inference/basic/basic_turbodiffusion_i2v.py new file mode 100644 index 000000000..c2b21059c --- /dev/null +++ b/examples/inference/basic/basic_turbodiffusion_i2v.py @@ -0,0 +1,41 @@ +import os + +# Set SLA attention backend BEFORE fastvideo imports +os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLA_ATTN" + +from fastvideo import VideoGenerator + +# Use local model path +MODEL_PATH = "loayrashid/TurboWan2.2-I2V-A14B-Diffusers" +OUTPUT_PATH = "video_samples_turbodiffusion_i2v" + + +def main() -> None: + # TurboDiffusion I2V: 1-4 step image-to-video generation + # Uses RCM scheduler with sigma_max=200 for I2V + generator = VideoGenerator.from_pretrained( + MODEL_PATH, + num_gpus=2, + override_pipeline_cls_name="TurboDiffusionI2VPipeline", + ) + + # Example prompt and image for I2V + prompt = ("Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside.") + + # Use an example image path + image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG" + + + video = generator.generate_video( + prompt, + image_path=image_path, + output_path=OUTPUT_PATH, + save_video=True, + num_inference_steps=4, + seed=42, + guidance_scale=1.0, + ) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index f5c8f4626..ff6bdfe49 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -88,6 +88,8 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = { lambda id: "stepvideo" in id.lower(), "cosmos": lambda id: "cosmos" in id.lower(), + "turbodiffusion": + lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(), # Add other pipeline architecture detectors } @@ -106,7 +108,8 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = { "wanimagetovideo": WanI2V480PConfig, "wandmdpipeline": FastWan2_1_T2V_480P_Config, "wancausaldmdpipeline": SelfForcingWanT2V480PConfig, - "stepvideo": StepVideoT2VConfig + "stepvideo": StepVideoT2VConfig, + "turbodiffusion": Wan2_2_I2V_A14B_Config, # Other fallbacks by architecture } diff --git a/fastvideo/pipelines/basic/turbodiffusion/__init__.py b/fastvideo/pipelines/basic/turbodiffusion/__init__.py index e69de29bb..0df303a96 100644 --- a/fastvideo/pipelines/basic/turbodiffusion/__init__.py +++ b/fastvideo/pipelines/basic/turbodiffusion/__init__.py @@ -0,0 +1,6 @@ +from fastvideo.pipelines.basic.turbodiffusion.turbodiffusion_pipeline import ( + TurboDiffusionPipeline, ) +from fastvideo.pipelines.basic.turbodiffusion.turbodiffusion_i2v_pipeline import ( + TurboDiffusionI2VPipeline, ) + +__all__ = ["TurboDiffusionPipeline", "TurboDiffusionI2VPipeline"] diff --git a/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_i2v_pipeline.py b/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_i2v_pipeline.py new file mode 100644 index 000000000..7dd75224c --- /dev/null +++ b/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_i2v_pipeline.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +TurboDiffusion I2V (Image-to-Video) Pipeline Implementation. + +This module contains an implementation of the TurboDiffusion I2V pipeline +for 1-4 step image-to-video generation using rCM (recurrent Consistency Model) +sampling with SLA (Sparse-Linear Attention). + +Key differences from T2V: +- Uses dual models (high/low noise) with boundary switching +- sigma_max=200 (vs 80 for T2V) +- Mask conditioning with encoded first frame +""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_rcm import RCMScheduler +from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, + DenoisingStage, ImageVAEEncodingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage) + +logger = init_logger(__name__) + + +class TurboDiffusionI2VPipeline(LoRAPipeline, ComposedPipelineBase): + """ + TurboDiffusion I2V pipeline for 1-4 step image-to-video generation. + + Uses RCM scheduler, SLA attention, and dual model switching for + high-quality I2V generation. + """ + + _required_config_modules = [ + "text_encoder", "tokenizer", "vae", "transformer", "transformer_2", + "scheduler" + ] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + # Use RCM scheduler with higher sigma_max for I2V + logger.info( + "Initializing RCM scheduler for TurboDiffusion I2V (sigma_max=200)") + self.modules["scheduler"] = RCMScheduler(sigma_max=200.0) + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + self.add_stage(stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + )) + + 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", None))) + + # I2V: Encode initial image to latent space + self.add_stage(stage_name="image_latent_preparation_stage", + stage=ImageVAEEncodingStage(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", None), + scheduler=self.get_module("scheduler"), + vae=self.get_module("vae"), + pipeline=self)) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"), + pipeline=self)) + + +EntryClass = TurboDiffusionI2VPipeline diff --git a/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py b/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py index c0918080c..e2dbfe7a2 100644 --- a/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py +++ b/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py @@ -36,11 +36,6 @@ class TurboDiffusionPipeline(LoRAPipeline, ComposedPipelineBase): logger.info("Initializing RCM scheduler for TurboDiffusion") self.modules["scheduler"] = RCMScheduler(sigma_max=80.0) - # Store checkpoint path for later loading - self._turbodiffusion_checkpoint = getattr(fastvideo_args, - 'turbodiffusion_checkpoint', - None) - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: """Set up pipeline stages with proper dependency injection.""" diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index 208a283aa..4728fa8dd 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -24,6 +24,7 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = { "WanVideoToVideoPipeline": "wan", "WanCausalDMDPipeline": "wan", "TurboDiffusionPipeline": "turbodiffusion", + "TurboDiffusionI2VPipeline": "turbodiffusion", "StepVideoPipeline": "stepvideo", "HunyuanVideoPipeline": "hunyuan", "HunyuanVideo15Pipeline": "hunyuan15", diff --git a/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 index ff53f78e8..e5734644d 100644 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 and b/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 differ diff --git a/fastvideo/tests/ssim/generated_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/steps4_Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of _ssim.json b/fastvideo/tests/ssim/generated_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/steps4_Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of _ssim.json index c2d82a787..77dfdbd5a 100644 --- a/fastvideo/tests/ssim/generated_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/steps4_Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of _ssim.json +++ b/fastvideo/tests/ssim/generated_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/steps4_Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of _ssim.json @@ -1,7 +1,7 @@ { - "mean_ssim": 0.9917645101194028, - "min_ssim": 0.9908460974693298, - "max_ssim": 0.9925968050956726, + "mean_ssim": 1.0, + "min_ssim": 1.0, + "max_ssim": 1.0, "reference_video": "/FastVideo/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4", "generated_video": "/FastVideo/fastvideo/tests/ssim/generated_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4", "parameters": { diff --git a/fastvideo/tests/ssim/generated_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/steps4_An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali_ssim.json b/fastvideo/tests/ssim/generated_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/steps4_An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali_ssim.json new file mode 100644 index 000000000..0a91ad365 --- /dev/null +++ b/fastvideo/tests/ssim/generated_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/steps4_An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali_ssim.json @@ -0,0 +1,11 @@ +{ + "mean_ssim": 1.0, + "min_ssim": 1.0, + "max_ssim": 1.0, + "reference_video": "/FastVideo/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4", + "generated_video": "/FastVideo/fastvideo/tests/ssim/generated_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4", + "parameters": { + "num_inference_steps": 4, + "prompt": "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot." + } +} \ No newline at end of file diff --git a/fastvideo/tests/ssim/test_turbodiffusion_similarity.py b/fastvideo/tests/ssim/test_turbodiffusion_similarity.py index b1ae13ef8..b71ace032 100644 --- a/fastvideo/tests/ssim/test_turbodiffusion_similarity.py +++ b/fastvideo/tests/ssim/test_turbodiffusion_similarity.py @@ -152,7 +152,145 @@ def test_turbodiffusion_inference_similarity(prompt, model_id): logger.error("Failed to write SSIM results to file") # TurboDiffusion uses fewer steps, may have slightly lower SSIM - min_acceptable_ssim = 0.90 + min_acceptable_ssim = 0.95 + assert mean_ssim >= min_acceptable_ssim, ( + f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " + f"for {model_id} with backend {ATTENTION_BACKEND}" + ) + + +# TurboDiffusion I2V parameters (dual-model with RCM scheduler + SLA attention) +TURBODIFFUSION_I2V_PARAMS = { + "num_gpus": 2, + "model_path": "loayrashid/TurboWan2.2-I2V-A14B-Diffusers", + "height": 480, + "width": 832, + "num_frames": 45, + "num_inference_steps": 4, # TurboDiffusion uses 1-4 steps + "guidance_scale": 1.0, # No CFG for TurboDiffusion + "seed": 42, + "sp_size": 2, + "tp_size": 1, + "fps": 24, +} + +TURBODIFFUSION_I2V_MODEL_TO_PARAMS = { + "TurboWan2.2-I2V-A14B-Diffusers": TURBODIFFUSION_I2V_PARAMS, +} + +TURBODIFFUSION_I2V_TEST_PROMPTS = [ + "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.", +] + +TURBODIFFUSION_I2V_IMAGE_PATHS = [ + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg", +] + + +@pytest.mark.parametrize("prompt", TURBODIFFUSION_I2V_TEST_PROMPTS) +@pytest.mark.parametrize("model_id", list(TURBODIFFUSION_I2V_MODEL_TO_PARAMS.keys())) +def test_turbodiffusion_i2v_inference_similarity(prompt, model_id): + """ + Test that runs TurboDiffusion I2V inference with dual-model switching, + then compares the output to reference videos using SSIM. + """ + # TurboDiffusion requires SLA attention backend + ATTENTION_BACKEND = "SLA_ATTN" + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + + assert len(TURBODIFFUSION_I2V_TEST_PROMPTS) == len(TURBODIFFUSION_I2V_IMAGE_PATHS), \ + "Expect number of prompts equal to number of images" + + script_dir = os.path.dirname(os.path.abspath(__file__)) + + base_output_dir = os.path.join(script_dir, 'generated_videos', model_id) + output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND) + output_video_name = f"{prompt[:100].strip()}.mp4" + + os.makedirs(output_dir, exist_ok=True) + + BASE_PARAMS = TURBODIFFUSION_I2V_MODEL_TO_PARAMS[model_id] + num_inference_steps = BASE_PARAMS["num_inference_steps"] + image_path = TURBODIFFUSION_I2V_IMAGE_PATHS[TURBODIFFUSION_I2V_TEST_PROMPTS.index(prompt)] + + init_kwargs = { + "num_gpus": BASE_PARAMS["num_gpus"], + "sp_size": BASE_PARAMS["sp_size"], + "tp_size": BASE_PARAMS["tp_size"], + "override_pipeline_cls_name": "TurboDiffusionI2VPipeline", + } + + generation_kwargs = { + "num_inference_steps": num_inference_steps, + "output_path": output_dir, + "image_path": image_path, + "height": BASE_PARAMS["height"], + "width": BASE_PARAMS["width"], + "num_frames": BASE_PARAMS["num_frames"], + "guidance_scale": BASE_PARAMS["guidance_scale"], + "seed": BASE_PARAMS["seed"], + "fps": BASE_PARAMS["fps"], + } + + generator = VideoGenerator.from_pretrained( + model_path=BASE_PARAMS["model_path"], + **init_kwargs + ) + generator.generate_video(prompt, **generation_kwargs) + + if isinstance(generator.executor, MultiprocExecutor): + generator.executor.shutdown() + + assert os.path.exists(output_dir), f"Output video was not generated at {output_dir}" + + reference_folder = os.path.join( + script_dir, device_reference_folder, model_id, ATTENTION_BACKEND + ) + + if not os.path.exists(reference_folder): + logger.error("Reference folder missing") + raise FileNotFoundError( + f"Reference video folder does not exist: {reference_folder}" + ) + + # Find the matching reference video based on the prompt + reference_video_name = None + + for filename in os.listdir(reference_folder): + if filename.endswith('.mp4') and prompt[:100].strip() in filename: + reference_video_name = filename + break + + if not reference_video_name: + logger.error( + f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}" + ) + raise FileNotFoundError(f"Reference video missing") + + reference_video_path = os.path.join(reference_folder, reference_video_name) + generated_video_path = os.path.join(output_dir, output_video_name) + + logger.info( + f"Computing SSIM between {reference_video_path} and {generated_video_path}" + ) + ssim_values = compute_video_ssim_torchvision( + reference_video_path, generated_video_path, use_ms_ssim=True + ) + + mean_ssim = ssim_values[0] + logger.info(f"SSIM mean value: {mean_ssim}") + logger.info(f"Writing SSIM results to directory: {output_dir}") + + success = write_ssim_results( + output_dir, ssim_values, reference_video_path, + generated_video_path, num_inference_steps, prompt + ) + + if not success: + logger.error("Failed to write SSIM results to file") + + # TurboDiffusion I2V uses fewer steps, may have slightly lower SSIM + min_acceptable_ssim = 0.95 assert mean_ssim >= min_acceptable_ssim, ( f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " f"for {model_id} with backend {ATTENTION_BACKEND}" diff --git a/scripts/checkpoint_conversion/convert_turbodiffusion_i2v_to_diffusers.py b/scripts/checkpoint_conversion/convert_turbodiffusion_i2v_to_diffusers.py new file mode 100644 index 000000000..3ed230b82 --- /dev/null +++ b/scripts/checkpoint_conversion/convert_turbodiffusion_i2v_to_diffusers.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python3 +""" +Convert TurboDiffusion I2V .pth checkpoint to Diffusers safetensors format. + +TurboDiffusion I2V uses two models: high-noise and low-noise. +This script converts both checkpoints to Diffusers format. + +Usage: + python convert_turbodiffusion_i2v_to_diffusers.py \ + --high_noise_path /path/to/TurboWan2.2-I2V-A14B-high-720P.pth \ + --low_noise_path /path/to/TurboWan2.2-I2V-A14B-low-720P.pth \ + --output_dir /path/to/output \ + --reference_repo Wan-AI/Wan2.1-I2V-14B-720P-Diffusers +""" + +import argparse +import os +import re +import json +import torch +import shutil +import glob +from safetensors import safe_open +from safetensors.torch import save_file +from huggingface_hub import snapshot_download + +# Weight mapping from TurboDiffusion -> Diffusers/FastVideo format +# Same as T2V but may need additional I2V-specific mappings +TURBODIFFUSION_WEIGHT_MAPPING = { + # Self attention + r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.to_q.\2", + r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.to_k.\2", + r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.to_v.\2", + r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.to_out.\2", + r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2", + r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2", + # Cross attention + r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2", + r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2", + r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2", + r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2", + # I2V-specific cross attention (add_k/v_proj for image context) + r"^blocks\.(\d+)\.cross_attn\.add_k\.(.*)$": r"blocks.\1.attn2.add_k_proj.\2", + r"^blocks\.(\d+)\.cross_attn\.add_v\.(.*)$": r"blocks.\1.attn2.add_v_proj.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_add_k\.(.*)$": r"blocks.\1.attn2.norm_added_k.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_add_q\.(.*)$": r"blocks.\1.attn2.norm_added_q.\2", + # Norms and FFN + r"^blocks\.(\d+)\.norm1\.(.*)$": r"blocks.\1.norm1.\2", + r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2", + r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.norm3.\2", + r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", + r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table", + # Embeddings + r"^text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1", + r"^text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1", + r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1", + r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1", + r"^time_projection\.1\.(.*)$": r"condition_embedder.time_modulation.linear.\1", + # Head + r"^head\.head\.(.*)$": r"proj_out.\1", + r"^head\.norm\.(.*)$": r"norm_out.\1", + r"^head\.modulation$": r"scale_shift_table", + # SLA proj_l weights + r"^blocks\.(\d+)\.self_attn\.attn_op\.local_attn\.proj_l\.(.*)$": r"blocks.\1.attn1.attn_impl.proj_l.\2", +} + +SKIP_PATTERNS = [] + + +def should_skip_key(key: str) -> bool: + for pattern in SKIP_PATTERNS: + if re.match(pattern, key): + return True + return False + + +def convert_key(turbo_key: str) -> str: + for pattern, replacement in TURBODIFFUSION_WEIGHT_MAPPING.items(): + if re.match(pattern, turbo_key): + return re.sub(pattern, replacement, turbo_key) + return turbo_key + + +def reshape_patch_embedding(tensor: torch.Tensor, target_shape: tuple) -> torch.Tensor: + if len(tensor.shape) == 2 and len(target_shape) == 5: + return tensor.view(target_shape) + return tensor + + +def get_reference_shapes(reference_repo: str) -> dict: + print(f"Downloading reference model shapes from {reference_repo}...") + + local_dir = snapshot_download( + repo_id=reference_repo, + allow_patterns=["transformer/config.json", "transformer/diffusion_pytorch_model*.safetensors"], + local_dir_use_symlinks=False + ) + + weight_files = glob.glob(os.path.join(local_dir, "transformer", "*.safetensors")) + shapes = {} + + for wf in weight_files: + with safe_open(wf, framework="pt") as f: + for key in f.keys(): + shapes[key] = f.get_tensor(key).shape + + return shapes, local_dir + + +def convert_checkpoint(input_path: str, output_dir: str, ref_shapes: dict, model_name: str) -> None: + """Convert a single TurboDiffusion checkpoint to Diffusers format.""" + print(f"\n{'='*60}") + print(f"Converting {model_name}: {input_path}") + print(f"{'='*60}") + + turbo_state_dict = torch.load(input_path, map_location="cpu", weights_only=True) + print(f"Loaded {len(turbo_state_dict)} keys") + + converted_state_dict = {} + skipped_keys = [] + + for turbo_key, tensor in turbo_state_dict.items(): + if should_skip_key(turbo_key): + skipped_keys.append(turbo_key) + continue + + new_key = convert_key(turbo_key) + + if "patch_embedding" in new_key and new_key in ref_shapes: + target_shape = ref_shapes[new_key] + if tensor.shape != target_shape: + print(f"Reshaping {new_key}: {tensor.shape} -> {target_shape}") + tensor = reshape_patch_embedding(tensor, target_shape) + + if new_key in ref_shapes: + if tensor.shape != ref_shapes[new_key]: + print(f"WARNING: Shape mismatch for {new_key}: got {tensor.shape}, expected {ref_shapes[new_key]}") + + converted_state_dict[new_key] = tensor + + print(f"Converted: {len(converted_state_dict)} keys, Skipped: {len(skipped_keys)} keys") + + os.makedirs(output_dir, exist_ok=True) + output_path = os.path.join(output_dir, "diffusion_pytorch_model.safetensors") + print(f"Saving to {output_path}...") + save_file(converted_state_dict, output_path) + + return converted_state_dict + + +def main(): + parser = argparse.ArgumentParser(description="Convert TurboDiffusion I2V checkpoints to Diffusers format") + parser.add_argument("--high_noise_path", type=str, required=True, + help="Path to high-noise TurboDiffusion .pth checkpoint") + parser.add_argument("--low_noise_path", type=str, required=True, + help="Path to low-noise TurboDiffusion .pth checkpoint") + parser.add_argument("--output_dir", type=str, required=True, + help="Output directory for converted safetensors") + parser.add_argument("--reference_repo", type=str, default="Wan-AI/Wan2.2-I2V-A14B-Diffusers", + help="Reference HF repo to get expected tensor shapes") + + args = parser.parse_args() + + # Get reference shapes + ref_shapes, ref_local_dir = get_reference_shapes(args.reference_repo) + print(f"Got {len(ref_shapes)} reference shapes") + + # Convert high-noise model + high_noise_output = os.path.join(args.output_dir, "transformer_high") + convert_checkpoint(args.high_noise_path, high_noise_output, ref_shapes, "high-noise") + + # Convert low-noise model + low_noise_output = os.path.join(args.output_dir, "transformer_low") + convert_checkpoint(args.low_noise_path, low_noise_output, ref_shapes, "low-noise") + + # Copy config.json to both + src_config = os.path.join(ref_local_dir, "transformer", "config.json") + shutil.copy(src_config, os.path.join(high_noise_output, "config.json")) + shutil.copy(src_config, os.path.join(low_noise_output, "config.json")) + print("Copied config.json to both transformer directories") + + print(f"\n{'='*60}") + print(f"Conversion complete!") + print(f"High-noise model: {high_noise_output}") + print(f"Low-noise model: {low_noise_output}") + print(f"{'='*60}") + + +if __name__ == "__main__": + main()