[feat] add Turbodiffusion I2V pipeline (#984)
This commit is contained in:
Loay Rashid
@@ -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()
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
BIN
Binary file not shown.
+3
-3
@@ -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": {
|
||||
|
||||
+11
@@ -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."
|
||||
}
|
||||
}
|
||||
@@ -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}"
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user