[feat] add Turbodiffusion I2V pipeline (#984)
This commit is contained in:
Loay Rashid
@@ -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