Add LongCat T2V (Base, Distillation and Refinement) Support to FastVideo (#883)

Co-authored-by: Shao Duan <shaoxiongduan@gmail.com>

scripts/checkpoint_conversion/longcat_to_fastvideo.py

@@ -0,0 +1,708 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""
+Convert LongCat weights to FastVideo native format.
+
+This script performs a complete conversion from original LongCat weights
+to FastVideo native implementation in a single step:
+
+1. Converts transformer weights (with QKV/KV splitting)
+2. Copies other components (VAE, text encoder, tokenizer, scheduler)
+3. Converts LoRA weights (cfg_step_lora, refinement_lora)
+4. Updates config files to point to native model
+
+Usage:
+    python scripts/checkpoint_conversion/longcat_to_fastvideo.py \
+        --source /path/to/LongCat-Video/weights/LongCat-Video \
+        --output weights/longcat-native \
+        --validate
+"""
+
+import argparse
+import glob
+import json
+import re
+import shutil
+from pathlib import Path
+from collections import OrderedDict
+
+import torch
+from safetensors.torch import load_file, save_file
+from tqdm import tqdm
+
+
+def split_qkv(qkv_weight: torch.Tensor, qkv_bias: torch.Tensor | None = None):
+    """Split fused QKV projection into separate Q, K, V."""
+    dim = qkv_weight.shape[0] // 3
+    q, k, v = torch.chunk(qkv_weight, 3, dim=0)
+    
+    if qkv_bias is not None:
+        q_bias, k_bias, v_bias = torch.chunk(qkv_bias, 3, dim=0)
+    else:
+        q_bias = k_bias = v_bias = None
+    
+    return (q, k, v), (q_bias, k_bias, v_bias)
+
+
+def split_kv(kv_weight: torch.Tensor, kv_bias: torch.Tensor | None = None):
+    """Split fused KV projection into separate K, V."""
+    dim = kv_weight.shape[0] // 2
+    k, v = torch.chunk(kv_weight, 2, dim=0)
+    
+    if kv_bias is not None:
+        k_bias, v_bias = torch.chunk(kv_bias, 2, dim=0)
+    else:
+        k_bias = v_bias = None
+    
+    return (k, v), (k_bias, v_bias)
+
+
+def convert_transformer_weights(source_weights: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """
+    Convert LongCat transformer weights to native FastVideo format.
+    
+    Main transformations:
+    1. Split fused QKV projections (self-attention)
+    2. Split fused KV projections (cross-attention)
+    3. Rename parameters according to mapping
+    """
+    converted = OrderedDict()
+    processed_keys = set()
+    
+    print("  Converting transformer weights...")
+    
+    for key, value in tqdm(source_weights.items(), desc="  Processing parameters"):
+        if key in processed_keys:
+            continue
+        
+        # === Embedders ===
+        if key.startswith("x_embedder."):
+            new_key = key.replace("x_embedder.", "patch_embed.")
+            converted[new_key] = value
+        
+        elif key.startswith("t_embedder.mlp.0."):
+            new_key = key.replace("t_embedder.mlp.0.", "time_embedder.linear_1.")
+            converted[new_key] = value
+        
+        elif key.startswith("t_embedder.mlp.2."):
+            new_key = key.replace("t_embedder.mlp.2.", "time_embedder.linear_2.")
+            converted[new_key] = value
+        
+        elif key.startswith("y_embedder.y_proj.0."):
+            new_key = key.replace("y_embedder.y_proj.0.", "caption_embedder.linear_1.")
+            converted[new_key] = value
+        
+        elif key.startswith("y_embedder.y_proj.2."):
+            new_key = key.replace("y_embedder.y_proj.2.", "caption_embedder.linear_2.")
+            converted[new_key] = value
+        
+        # === Self-Attention QKV Splitting ===
+        elif ".attn.qkv.weight" in key:
+            block_idx = key.split(".")[1]
+            qkv_weight = value
+            qkv_bias_key = key.replace(".weight", ".bias")
+            qkv_bias = source_weights.get(qkv_bias_key)
+            
+            (q, k, v), (q_bias, k_bias, v_bias) = split_qkv(qkv_weight, qkv_bias)
+            
+            converted[f"blocks.{block_idx}.self_attn.to_q.weight"] = q
+            converted[f"blocks.{block_idx}.self_attn.to_k.weight"] = k
+            converted[f"blocks.{block_idx}.self_attn.to_v.weight"] = v
+            
+            if q_bias is not None:
+                converted[f"blocks.{block_idx}.self_attn.to_q.bias"] = q_bias
+                converted[f"blocks.{block_idx}.self_attn.to_k.bias"] = k_bias
+                converted[f"blocks.{block_idx}.self_attn.to_v.bias"] = v_bias
+            
+            processed_keys.add(key)
+            if qkv_bias is not None:
+                processed_keys.add(qkv_bias_key)
+        
+        elif ".attn.qkv.bias" in key:
+            continue
+        
+        elif ".attn.proj." in key:
+            new_key = key.replace(".attn.proj.", ".self_attn.to_out.")
+            converted[new_key] = value
+        
+        elif ".attn.q_norm." in key or ".attn.k_norm." in key:
+            new_key = key.replace(".attn.", ".self_attn.")
+            converted[new_key] = value
+        
+        # === Cross-Attention ===
+        elif ".cross_attn.q_linear." in key:
+            new_key = key.replace(".cross_attn.q_linear.", ".cross_attn.to_q.")
+            converted[new_key] = value
+        
+        elif ".cross_attn.kv_linear.weight" in key:
+            block_idx = key.split(".")[1]
+            kv_weight = value
+            kv_bias_key = key.replace(".weight", ".bias")
+            kv_bias = source_weights.get(kv_bias_key)
+            
+            (k, v), (k_bias, v_bias) = split_kv(kv_weight, kv_bias)
+            
+            converted[f"blocks.{block_idx}.cross_attn.to_k.weight"] = k
+            converted[f"blocks.{block_idx}.cross_attn.to_v.weight"] = v
+            
+            if k_bias is not None:
+                converted[f"blocks.{block_idx}.cross_attn.to_k.bias"] = k_bias
+                converted[f"blocks.{block_idx}.cross_attn.to_v.bias"] = v_bias
+            
+            processed_keys.add(key)
+            if kv_bias is not None:
+                processed_keys.add(kv_bias_key)
+        
+        elif ".cross_attn.kv_linear.bias" in key:
+            continue
+        
+        elif ".cross_attn.proj." in key:
+            new_key = key.replace(".cross_attn.proj.", ".cross_attn.to_out.")
+            converted[new_key] = value
+        
+        elif ".cross_attn.q_norm." in key or ".cross_attn.k_norm." in key:
+            converted[key] = value
+        
+        # === Final Layer (must come BEFORE general transformer block patterns) ===
+        elif key.startswith("final_layer.adaLN_modulation.1."):
+            new_key = key.replace("final_layer.adaLN_modulation.1.", "final_layer.adaln_linear.")
+            converted[new_key] = value
+        
+        # === Transformer Block AdaLN ===
+        elif ".adaLN_modulation.1." in key:
+            new_key = key.replace(".adaLN_modulation.1.", ".adaln_linear_1.")
+            converted[new_key] = value
+        
+        # === Transformer Block Normalization ===
+        elif ".mod_norm_attn." in key or ".mod_norm_ffn." in key:
+            continue
+        
+        elif ".pre_crs_attn_norm.weight" in key:
+            new_key = key.replace(".pre_crs_attn_norm.", ".norm_cross.")
+            converted[new_key] = value
+        
+        elif ".pre_crs_attn_norm.bias" in key:
+            new_key = key.replace(".pre_crs_attn_norm.", ".norm_cross.")
+            converted[new_key] = value
+        
+        # === FFN (SwiGLU) ===
+        elif ".ffn.w1." in key or ".ffn.w2." in key or ".ffn.w3." in key:
+            converted[key] = value
+        
+        elif key.startswith("final_layer.norm_final."):
+            continue
+        
+        elif key.startswith("final_layer.linear."):
+            new_key = key.replace("final_layer.linear.", "final_layer.proj.")
+            converted[new_key] = value
+        
+        else:
+            print(f"    ⚠️  Unknown key: {key}")
+            converted[key] = value
+    
+    return converted
+
+
+def validate_conversion(original: dict, converted: dict) -> bool:
+    """Validate that conversion preserved all parameters correctly."""
+    print("\n  Validating conversion...")
+    
+    orig_count = sum(p.numel() for p in original.values())
+    conv_count = sum(p.numel() for p in converted.values())
+    
+    dropped_count = 0
+    for key, value in original.items():
+        if ".mod_norm_attn." in key or ".mod_norm_ffn." in key:
+            dropped_count += value.numel()
+        elif "final_layer.norm_final." in key:
+            dropped_count += value.numel()
+    
+    expected_conv_count = orig_count - dropped_count
+    
+    print(f"    Original parameters: {orig_count:,}")
+    print(f"    Converted parameters: {conv_count:,}")
+    print(f"    Dropped parameters (norms without params): {dropped_count:,}")
+    
+    if conv_count != expected_conv_count:
+        print(f"    ⚠️  Parameter count mismatch!")
+        return False
+    
+    print(f"    ✓ Parameter count matches")
+    
+    # Verify QKV/KV splits
+    print("\n  Verifying QKV/KV splits...")
+    num_blocks = 48
+    
+    for i in range(num_blocks):
+        orig_qkv_weight = original.get(f"blocks.{i}.attn.qkv.weight")
+        if orig_qkv_weight is not None:
+            conv_q = converted[f"blocks.{i}.self_attn.to_q.weight"]
+            conv_k = converted[f"blocks.{i}.self_attn.to_k.weight"]
+            conv_v = converted[f"blocks.{i}.self_attn.to_v.weight"]
+            reconstructed = torch.cat([conv_q, conv_k, conv_v], dim=0)
+            if not torch.allclose(orig_qkv_weight, reconstructed):
+                print(f"    ❌ QKV weight mismatch in block {i}")
+                return False
+        
+        orig_kv_weight = original.get(f"blocks.{i}.cross_attn.kv_linear.weight")
+        if orig_kv_weight is not None:
+            conv_k = converted[f"blocks.{i}.cross_attn.to_k.weight"]
+            conv_v = converted[f"blocks.{i}.cross_attn.to_v.weight"]
+            reconstructed = torch.cat([conv_k, conv_v], dim=0)
+            if not torch.allclose(orig_kv_weight, reconstructed):
+                print(f"    ❌ KV weight mismatch in block {i}")
+                return False
+    
+    print(f"    ✓ All splits verified successfully")
+    return True
+
+
+def copy_component(source_dir: Path, output_dir: Path, component: str, mapping: dict = None) -> bool:
+    """Copy a component directory, optionally with name mapping."""
+    source_name = mapping.get(component, component) if mapping else component
+    source_path = source_dir / source_name
+    
+    if source_path.exists():
+        output_path = output_dir / component
+        if output_path.exists():
+            shutil.rmtree(output_path)
+        shutil.copytree(source_path, output_path)
+        print(f"  ✓ {component} copied")
+        return True
+    else:
+        print(f"  ⚠️  {component} not found, skipping")
+        return False
+
+
+def create_model_index():
+    """Create model_index.json for FastVideo native model."""
+    return {
+        "_class_name": "LongCatPipeline",
+        "_diffusers_version": "0.32.0",
+        "workload_type": "video-generation",
+        "tokenizer": ["transformers", "AutoTokenizer"],
+        "text_encoder": ["transformers", "UMT5EncoderModel"],
+        "vae": ["diffusers", "AutoencoderKLWan"],
+        "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
+        "transformer": ["diffusers", "LongCatTransformer3DModel"]  # Native model
+    }
+
+
+def update_transformer_config(transformer_dir: Path):
+    """Update transformer config.json to point to native model."""
+    config_path = transformer_dir / "config.json"
+    if not config_path.exists():
+        print("  ⚠️  Transformer config not found, skipping")
+        return
+    
+    with open(config_path, 'r') as f:
+        config = json.load(f)
+    
+    if '_class_name' in config:
+        old_class = config['_class_name']
+        config['_class_name'] = 'LongCatTransformer3DModel'
+        print(f"  Updated _class_name: {old_class} → LongCatTransformer3DModel")
+    else:
+        config['_class_name'] = 'LongCatTransformer3DModel'
+        print(f"  Added _class_name: LongCatTransformer3DModel")
+    
+    # Fix num_heads -> num_attention_heads for FastVideo compatibility
+    if 'num_heads' in config and 'num_attention_heads' not in config:
+        config['num_attention_heads'] = config.pop('num_heads')
+        print(f"  Updated num_heads → num_attention_heads")
+    
+    with open(config_path, 'w') as f:
+        json.dump(config, f, indent=2)
+    
+    print("  ✓ Transformer config updated")
+
+
+# ============================================================================
+# LoRA Conversion Functions
+# ============================================================================
+
+def parse_lora_key(key: str) -> tuple[str, str]:
+    """Parse LongCat LoRA key into module path and weight type."""
+    if key.startswith("lora___lorahyphen___"):
+        key = key[len("lora___lorahyphen___"):]
+    
+    key = key.replace("___lorahyphen___", ".")
+    
+    if ".lora_down.weight" in key:
+        return key.replace(".lora_down.weight", ""), "lora_down.weight"
+    elif ".lora_up.weight" in key:
+        return key.replace(".lora_up.weight", ""), "lora_up.weight"
+    elif ".lora_up.blocks." in key:
+        match = re.match(r"(.+)\.lora_up\.blocks\.(\d+)\.weight", key)
+        if match:
+            return match.group(1), f"lora_up.blocks.{match.group(2)}.weight"
+    elif ".alpha_scale" in key:
+        return key.replace(".alpha_scale", ""), "alpha_scale"
+    elif ".lora_alpha" in key:
+        return key.replace(".lora_alpha", ""), "lora_alpha"
+    
+    raise ValueError(f"Unknown LoRA key format: {key}")
+
+
+def map_lora_module(module_path: str) -> list[tuple[str, str]]:
+    """Map LongCat module path to FastVideo paths. Returns [(path, component)]."""
+    # Self-attention QKV → Q, K, V
+    match = re.match(r"blocks\.(\d+)\.attn\.qkv", module_path)
+    if match:
+        b = match.group(1)
+        return [(f"blocks.{b}.self_attn.to_q", "q"),
+                (f"blocks.{b}.self_attn.to_k", "k"),
+                (f"blocks.{b}.self_attn.to_v", "v")]
+    
+    # Self-attention output
+    match = re.match(r"blocks\.(\d+)\.attn\.proj", module_path)
+    if match:
+        return [(f"blocks.{match.group(1)}.self_attn.to_out", "single")]
+    
+    # Cross-attention Q
+    match = re.match(r"blocks\.(\d+)\.cross_attn\.q_linear", module_path)
+    if match:
+        return [(f"blocks.{match.group(1)}.cross_attn.to_q", "single")]
+    
+    # Cross-attention KV → K, V
+    match = re.match(r"blocks\.(\d+)\.cross_attn\.kv_linear", module_path)
+    if match:
+        b = match.group(1)
+        return [(f"blocks.{b}.cross_attn.to_k", "k"),
+                (f"blocks.{b}.cross_attn.to_v", "v")]
+    
+    # FFN
+    match = re.match(r"blocks\.(\d+)\.ffn\.(w[123])", module_path)
+    if match:
+        return [(f"blocks.{match.group(1)}.ffn.{match.group(2)}", "single")]
+    
+    # AdaLN modulation
+    match = re.match(r"blocks\.(\d+)\.adaLN_modulation\.1", module_path)
+    if match:
+        return [(f"blocks.{match.group(1)}.adaln_linear_1", "single")]
+    
+    # Final layer
+    if module_path == "final_layer.adaLN_modulation.1":
+        return [("final_layer.adaln_linear", "single")]
+    if module_path == "final_layer.linear":
+        return [("final_layer.proj", "single")]
+    
+    raise ValueError(f"Unknown LoRA module: {module_path}")
+
+
+def convert_lora_weights(source_weights: dict[str, torch.Tensor], lora_name: str) -> dict[str, torch.Tensor]:
+    """Convert LongCat LoRA to FastVideo format."""
+    print(f"  Converting {lora_name}...")
+    print(f"    Source keys: {len(source_weights)}")
+    
+    converted = OrderedDict()
+    
+    # Group by module
+    modules = {}
+    for key in source_weights.keys():
+        try:
+            module_path, weight_type = parse_lora_key(key)
+            if module_path not in modules:
+                modules[module_path] = {}
+            modules[module_path][weight_type] = key
+        except ValueError:
+            continue
+    
+    # Process each module
+    for module_path, weight_keys in modules.items():
+        try:
+            targets = map_lora_module(module_path)
+        except ValueError:
+            continue
+        
+        # Get alpha_scale if present (defaults to 1.0 if missing)
+        alpha_scale = 1.0
+        if "alpha_scale" in weight_keys:
+            alpha_scale_tensor = source_weights[weight_keys["alpha_scale"]]
+            alpha_scale = alpha_scale_tensor.item() if alpha_scale_tensor.numel() == 1 else float(alpha_scale_tensor.mean())
+        
+        # Handle lora_down (lora_A)
+        if "lora_down.weight" in weight_keys:
+            lora_down = source_weights[weight_keys["lora_down.weight"]]
+            
+            if len(targets) == 1:
+                converted[f"{targets[0][0]}.lora_A"] = lora_down
+                # Compute alpha from alpha_scale and rank
+                rank = lora_down.shape[0]
+                alpha = alpha_scale * rank
+                converted[f"{targets[0][0]}.lora_alpha"] = torch.tensor(alpha, dtype=torch.float32)
+            else:
+                # Split for fused projections
+                n = len(targets)
+                rank = lora_down.shape[0] // n
+                for i, (path, _) in enumerate(targets):
+                    converted[f"{path}.lora_A"] = lora_down[i*rank:(i+1)*rank, :]
+                    # Compute alpha from alpha_scale and rank for each split
+                    alpha = alpha_scale * rank
+                    converted[f"{path}.lora_alpha"] = torch.tensor(alpha, dtype=torch.float32)
+        
+        # Handle lora_up (lora_B) - may have multiple blocks
+        lora_up_blocks = []
+        i = 0
+        while f"lora_up.blocks.{i}.weight" in weight_keys:
+            lora_up_blocks.append(source_weights[weight_keys[f"lora_up.blocks.{i}.weight"]])
+            i += 1
+        
+        if lora_up_blocks:
+            # Multi-block LoRA: construct block-diagonal lora_B
+            # This is equivalent to the multi-block computation without modifying fastvideo
+            n_blocks = len(lora_up_blocks)
+            out_per_block, rank_per_block = lora_up_blocks[0].shape  # e.g., [4096, 128]
+            
+            if len(targets) == 1:
+                # Single layer with multi-block: create block-diagonal matrix
+                total_out = out_per_block * n_blocks
+                total_rank = rank_per_block * n_blocks
+                lora_B_blockdiag = torch.zeros(total_out, total_rank, dtype=lora_up_blocks[0].dtype)
+                
+                for i in range(n_blocks):
+                    lora_B_blockdiag[i*out_per_block:(i+1)*out_per_block, 
+                                     i*rank_per_block:(i+1)*rank_per_block] = lora_up_blocks[i]
+                
+                converted[f"{targets[0][0]}.lora_B"] = lora_B_blockdiag
+                # Note: rank for alpha calculation should be total_rank (will be computed from lora_A.shape[0])
+            else:
+                # Multi-block with split targets (e.g., QKV split)
+                # Each target gets one block
+                for i, (path, _) in enumerate(targets):
+                    if i < n_blocks:
+                        converted[f"{path}.lora_B"] = lora_up_blocks[i]
+        elif "lora_up.weight" in weight_keys:
+            lora_up = source_weights[weight_keys["lora_up.weight"]]
+            # Split if needed
+            if len(targets) == 1:
+                converted[f"{targets[0][0]}.lora_B"] = lora_up
+            else:
+                n = len(targets)
+                out_dim = lora_up.shape[0] // n
+                for i, (path, _) in enumerate(targets):
+                    converted[f"{path}.lora_B"] = lora_up[i*out_dim:(i+1)*out_dim, :]
+        else:
+            continue
+    
+    print(f"    Output keys: {len(converted)} (including lora_alpha)")
+    # Count how many lora_alpha values were added
+    alpha_count = sum(1 for k in converted.keys() if "lora_alpha" in k)
+    print(f"    Alpha values saved: {alpha_count}")
+    return converted
+
+
+def convert_loras(source_dir: Path, output_dir: Path) -> bool:
+    """Convert all LoRA files in source directory."""
+    lora_source = source_dir / "lora"
+    if not lora_source.exists():
+        print("  No LoRA directory found, skipping")
+        return False
+    
+    lora_files = list(lora_source.glob("*.safetensors"))
+    if not lora_files:
+        print("  No LoRA files found, skipping")
+        return False
+    
+    print(f"  Found {len(lora_files)} LoRA file(s)")
+    
+    # Map LoRA filenames to subdirectory names for FastVideo compatibility
+    # Each LoRA gets its own directory under lora/
+    lora_subdir_mapping = {
+        "cfg_step_lora.safetensors": "distilled",
+        "refinement_lora.safetensors": "refinement",
+    }
+    
+    for lora_file in lora_files:
+        try:
+            # Determine output subdirectory - use mapping if available, otherwise generic name
+            if lora_file.name in lora_subdir_mapping:
+                lora_subdir_name = lora_subdir_mapping[lora_file.name]
+            else:
+                # For unknown LoRAs, create subdirectory based on filename
+                lora_subdir_name = lora_file.stem
+            
+            lora_output = output_dir / "lora" / lora_subdir_name
+            lora_output.mkdir(parents=True, exist_ok=True)
+            
+            # Load
+            source_weights = load_file(str(lora_file))
+            
+            # Convert
+            converted = convert_lora_weights(source_weights, lora_file.stem)
+            
+            # Save
+            output_file = lora_output / lora_file.name
+            save_file(converted, str(output_file))
+            
+            size_mb = output_file.stat().st_size / (1024**2)
+            print(f"    ✓ {lora_file.name} → lora/{lora_subdir_name}/ ({size_mb:.1f} MB)")
+        
+        except Exception as e:
+            print(f"    ❌ Failed to convert {lora_file.name}: {e}")
+            return False
+    
+    return True
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert LongCat weights to FastVideo native format"
+    )
+    parser.add_argument(
+        "--source",
+        type=str,
+        required=True,
+        help="Path to original LongCat weights (LongCat-Video/weights/LongCat-Video/)",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        required=True,
+        help="Path to output directory for native weights",
+    )
+    parser.add_argument(
+        "--validate",
+        action="store_true",
+        help="Run validation after conversion",
+    )
+    
+    args = parser.parse_args()
+    
+    source_dir = Path(args.source)
+    output_dir = Path(args.output)
+    
+    # Check source directory
+    if not source_dir.exists():
+        print(f"❌ Error: Source directory not found: {source_dir}")
+        return 1
+    
+    # Check for dit/transformer directory (original uses 'dit', we output to 'transformer')
+    transformer_source = source_dir / "dit"
+    if not transformer_source.exists():
+        print(f"❌ Error: DiT directory not found in source")
+        return 1
+    
+    print("=" * 60)
+    print("LongCat → FastVideo Native Conversion")
+    print("=" * 60)
+    print(f"Source: {source_dir}")
+    print(f"Output: {output_dir}")
+    print()
+    
+    # Step 1: Convert transformer weights
+    print("[Step 1/4] Converting transformer weights...")
+    
+    # Load source weights
+    shard_files = sorted(glob.glob(str(transformer_source / "*.safetensors")))
+    if not shard_files:
+        print(f"❌ Error: No safetensors files found in {transformer_source}")
+        return 1
+    
+    print(f"  Found {len(shard_files)} shard(s)")
+    source_weights = {}
+    for shard_file in shard_files:
+        print(f"  Loading {Path(shard_file).name}...")
+        source_weights.update(load_file(shard_file))
+    
+    print(f"  Loaded {len(source_weights)} parameters")
+    
+    # Convert
+    converted_weights = convert_transformer_weights(source_weights)
+    print(f"  Converted to {len(converted_weights)} parameters")
+    
+    # Validate if requested
+    if args.validate:
+        if not validate_conversion(source_weights, converted_weights):
+            print("\n❌ Validation failed!")
+            return 1
+        print("\n✓ Validation passed!")
+    
+    # Save
+    transformer_output = output_dir / "transformer"
+    transformer_output.mkdir(parents=True, exist_ok=True)
+    output_file = transformer_output / "model.safetensors"
+    print(f"\n  Saving to {output_file}...")
+    save_file(converted_weights, str(output_file))
+    size_gb = output_file.stat().st_size / (1024**3)
+    print(f"  ✓ Saved ({size_gb:.2f} GB)")
+    print()
+    
+    # Step 2: Copy other components
+    print("[Step 2/5] Copying other components...")
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    components = ["vae", "text_encoder", "tokenizer", "scheduler"]
+    for component in components:
+        copy_component(source_dir, output_dir, component)
+    
+    print()
+    
+    # Step 3: Convert LoRA weights
+    print("[Step 3/5] Converting LoRA weights...")
+    convert_loras(source_dir, output_dir)
+    print()
+    
+    # Step 4: Update transformer config
+    print("[Step 4/5] Updating transformer config...")
+    
+    # Copy config.json from source if exists
+    source_config = transformer_source / "config.json"
+    output_config = transformer_output / "config.json"
+    if source_config.exists():
+        shutil.copy(source_config, output_config)
+        print(f"  Copied config.json")
+    
+    update_transformer_config(transformer_output)
+    print()
+    
+    # Step 5: Create model_index.json
+    print("[Step 5/5] Creating model_index.json...")
+    model_index_path = output_dir / "model_index.json"
+    with open(model_index_path, 'w') as f:
+        json.dump(create_model_index(), f, indent=2)
+    print(f"  ✓ Created {model_index_path}")
+    print()
+    
+    print("=" * 60)
+    print("✓ Conversion Complete!")
+    print("=" * 60)
+    print(f"Native weights ready at: {output_dir}")
+    print()
+    print("Converted components:")
+    print("  ✓ Transformer (native FastVideo implementation)")
+    print("  ✓ VAE, text encoder, tokenizer, scheduler")
+    if (output_dir / "lora").exists():
+        lora_dirs = [d for d in (output_dir / "lora").iterdir() if d.is_dir()]
+        if lora_dirs:
+            print(f"  ✓ LoRA weights ({len(lora_dirs)} adapters)")
+            for lora_dir in sorted(lora_dirs):
+                print(f"    - lora/{lora_dir.name}/")
+    print()
+    print("Next steps:")
+    print()
+    print("  1. Test basic generation:")
+    print("     from fastvideo import VideoGenerator")
+    print(f"     generator = VideoGenerator.from_pretrained('{output_dir}')")
+    print("     video = generator.generate_video(")
+    print("         prompt='A cat playing piano',")
+    print("         num_inference_steps=50")
+    print("     )")
+    print()
+    if (output_dir / "lora" / "distilled").exists():
+        print("  2. Test distilled generation (16 steps with LoRA):")
+        print(f"     generator = VideoGenerator.from_pretrained('{output_dir}',")
+        print(f"         lora_path='{output_dir}/lora/distilled',")
+        print("         lora_nickname='distilled')")
+        print("     video = generator.generate_video(")
+        print("         prompt='A cat playing piano',")
+        print("         num_inference_steps=16,")
+        print("         guidance_scale=1.0)")
+    print()
+    print()
+    
+    return 0
+
+
+if __name__ == "__main__":
+    exit(main())

scripts/checkpoint_conversion/validate_longcat_weights.py

@@ -0,0 +1,241 @@
+import os
+import json
+import argparse
+from pathlib import Path
+from safetensors import safe_open
+
+
+def validate_components(model_path):
+    """Validate all model components exist."""
+    print("=" * 60)
+    print("VALIDATING COMPONENTS")
+    print("=" * 60)
+    
+    components = {
+        "tokenizer": ["special_tokens_map.json", "tokenizer_config.json"],
+        "text_encoder": ["config.json", "model.safetensors.index.json"],
+        "vae": ["config.json", "diffusion_pytorch_model.safetensors"],
+        "scheduler": ["scheduler_config.json"],
+        "transformer": ["config.json", "diffusion_pytorch_model.safetensors.index.json"]
+    }
+    
+    all_valid = True
+    for component, required_files in components.items():
+        component_path = os.path.join(model_path, component)
+        print(f"\n{component}:")
+        
+        if not os.path.exists(component_path):
+            print(f"  ✗ Directory not found")
+            all_valid = False
+            continue
+        
+        for req_file in required_files:
+            file_path = os.path.join(component_path, req_file)
+            exists = os.path.exists(file_path)
+            symbol = "✓" if exists else "✗"
+            print(f"  {symbol} {req_file}")
+            if not exists:
+                all_valid = False
+    
+    return all_valid
+
+
+def validate_dit_weights(dit_path):
+    """Validate DiT weights structure."""
+    print("\n" + "=" * 60)
+    print("VALIDATING DiT WEIGHTS")
+    print("=" * 60)
+    
+    # Load config
+    config_path = os.path.join(dit_path, "config.json")
+    with open(config_path) as f:
+        config = json.load(f)
+    
+    hidden_size = config["hidden_size"]
+    depth = config["depth"]
+    num_heads = config["num_heads"]
+    
+    print(f"\nArchitecture:")
+    print(f"  - hidden_size: {hidden_size}")
+    print(f"  - depth: {depth}")
+    print(f"  - num_heads: {num_heads}")
+    
+    # Load weight index
+    index_path = os.path.join(dit_path, "diffusion_pytorch_model.safetensors.index.json")
+    with open(index_path) as f:
+        index = json.load(f)
+    
+    weight_map = index["weight_map"]
+    all_keys = list(weight_map.keys())
+    
+    print(f"\nWeight statistics:")
+    print(f"  - Total keys: {len(all_keys)}")
+    print(f"  - Total size: {index['metadata']['total_size'] / 1e9:.2f} GB")
+    
+    # Check structure
+    embedder_keys = [k for k in all_keys if 'embedder' in k]
+    block_keys = [k for k in all_keys if k.startswith('blocks.')]
+    final_keys = [k for k in all_keys if k.startswith('final_layer.')]
+    
+    print(f"\nKey distribution:")
+    print(f"  - Embedder layers: {len(embedder_keys)}")
+    print(f"  - Transformer blocks: {len(block_keys)}")
+    print(f"  - Final layer: {len(final_keys)}")
+    
+    # Verify all blocks present
+    block_nums = set()
+    for key in block_keys:
+        if key.startswith('blocks.'):
+            block_num = int(key.split('.')[1])
+            block_nums.add(block_num)
+    
+    expected_blocks = set(range(depth))
+    missing_blocks = expected_blocks - block_nums
+    
+    if missing_blocks:
+        print(f"\n✗ Missing blocks: {sorted(missing_blocks)}")
+        return False
+    else:
+        print(f"\n✓ All {depth} blocks present (0-{depth-1})")
+    
+    # Sample weights
+    first_shard = os.path.join(dit_path, "diffusion_pytorch_model-00001-of-00006.safetensors")
+    print(f"\nSampling weights from first shard:")
+    
+    with safe_open(first_shard, framework="pt", device="cpu") as f:
+        sample_keys = [k for k in f.keys() if k in all_keys][:5]
+        for key in sample_keys:
+            tensor = f.get_tensor(key)
+            print(f"  - {key}")
+            print(f"    Shape: {tuple(tensor.shape)}, Dtype: {tensor.dtype}")
+    
+    return True
+
+
+def validate_shapes(dit_path):
+    """Validate weight shapes match expected architecture."""
+    print("\n" + "=" * 60)
+    print("VALIDATING WEIGHT SHAPES")
+    print("=" * 60)
+    
+    # Load config
+    config_path = os.path.join(dit_path, "config.json")
+    with open(config_path) as f:
+        config = json.load(f)
+    
+    hidden_size = config["hidden_size"]
+    num_heads = config["num_heads"]
+    head_dim = hidden_size // num_heads
+    adaln_dim = config.get("adaln_tembed_dim", 512)
+    mlp_ratio = config.get("mlp_ratio", 4)
+    
+    # Calculate FFN hidden_dim using SwiGLU formula from blocks.py
+    # hidden_dim = int(2 * (hidden_size * mlp_ratio) / 3)
+    # rounded to multiple_of=256
+    multiple_of = 256
+    ffn_hidden = int(2 * hidden_size * mlp_ratio / 3)
+    ffn_hidden = multiple_of * ((ffn_hidden + multiple_of - 1) // multiple_of)
+    
+    # Expected shapes
+    expected = {
+        "x_embedder.proj.weight": (hidden_size, 16, 1, 2, 2),
+        "x_embedder.proj.bias": (hidden_size,),
+        "t_embedder.mlp.0.weight": (adaln_dim, 256),
+        "t_embedder.mlp.2.weight": (adaln_dim, adaln_dim),
+        "y_embedder.y_proj.0.weight": (hidden_size, 4096),
+        "blocks.0.attn.qkv.weight": (3 * hidden_size, hidden_size),
+        "blocks.0.attn.q_norm.weight": (head_dim,),
+        "blocks.0.attn.proj.weight": (hidden_size, hidden_size),
+        "blocks.0.cross_attn.q_linear.weight": (hidden_size, hidden_size),
+        "blocks.0.cross_attn.kv_linear.weight": (2 * hidden_size, hidden_size),
+        "blocks.0.ffn.w1.weight": (ffn_hidden, hidden_size),
+        "blocks.0.adaLN_modulation.1.weight": (6 * hidden_size, adaln_dim),
+        "final_layer.linear.weight": (64, hidden_size),
+    }
+    
+    # Load and check
+    first_shard = os.path.join(dit_path, "diffusion_pytorch_model-00001-of-00006.safetensors")
+    
+    all_valid = True
+    with safe_open(first_shard, framework="pt", device="cpu") as f:
+        for key, expected_shape in expected.items():
+            if key in f.keys():
+                tensor = f.get_tensor(key)
+                actual_shape = tuple(tensor.shape)
+                
+                if actual_shape == expected_shape:
+                    print(f"✓ {key}: {actual_shape}")
+                else:
+                    print(f"✗ {key}: expected {expected_shape}, got {actual_shape}")
+                    all_valid = False
+    
+    return all_valid
+
+
+def validate_model_index(model_path):
+    """Validate model_index.json exists and is correct."""
+    print("\n" + "=" * 60)
+    print("VALIDATING MODEL INDEX")
+    print("=" * 60)
+    
+    model_index_path = os.path.join(model_path, "model_index.json")
+    
+    if not os.path.exists(model_index_path):
+        print("✗ model_index.json not found")
+        return False
+    
+    with open(model_index_path) as f:
+        index = json.load(f)
+    
+    required_keys = ["_class_name", "workload_type", "tokenizer", "text_encoder", 
+                     "vae", "scheduler", "transformer"]
+    
+    all_valid = True
+    for key in required_keys:
+        if key in index:
+            print(f"✓ {key}: {index[key]}")
+        else:
+            print(f"✗ {key}: missing")
+            all_valid = False
+    
+    return all_valid
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Validate LongCat weights for FastVideo")
+    parser.add_argument("--model-path", type=str, required=True,
+                        help="Path to LongCat model directory")
+    parser.add_argument("--check-shapes", action="store_true",
+                        help="Also validate weight shapes (slower)")
+    
+    args = parser.parse_args()
+    
+    print(f"\nValidating: {args.model_path}\n")
+    
+    # Run validations
+    components_valid = validate_components(args.model_path)
+    model_index_valid = validate_model_index(args.model_path)
+    dit_valid = validate_dit_weights(os.path.join(args.model_path, "transformer"))
+    
+    if args.check_shapes:
+        shapes_valid = validate_shapes(os.path.join(args.model_path, "transformer"))
+    else:
+        shapes_valid = True
+        print("\nSkipping shape validation (use --check-shapes to enable)")
+    
+    # Summary
+    print("\n" + "=" * 60)
+    print("VALIDATION SUMMARY")
+    print("=" * 60)
+    
+    all_valid = components_valid and model_index_valid and dit_valid and shapes_valid
+    
+    if all_valid:
+        print("✓ All validations passed!")
+        print("✓ Model ready for FastVideo")
+    else:
+        print("✗ Some validations failed")
+        print("✗ Please check errors above")
+    
+    print("=" * 60)
+

scripts/inference/v1_inference_longcat.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+
+num_gpus=2
+
+export FASTVIDEO_ATTENTION_BACKEND=
+
+# For longcat, we must first convert the official weights to FastVideo native format
+# conversion method: python scripts/checkpoint_conversion/longcat_to_fastvideo.py
+# --source /path/to/LongCat-Video/weights/LongCat-Video
+# --output weights/longcat-native
+export MODEL_BASE=weights/longcat-native
+
+fastvideo generate \
+    --model-path $MODEL_BASE \
+    --sp-size $num_gpus \
+    --tp-size 1 \
+    --num-gpus $num_gpus \
+    --dit-cpu-offload False \
+    --vae-cpu-offload False \
+    --text-encoder-cpu-offload False \
+    --pin-cpu-memory False \
+    --enable-bsa False \
+    --height 480 \
+    --width 832 \
+    --num-frames 93 \
+    --num-inference-steps 50 \
+    --fps 15 \
+    --guidance-scale 4.0 \
+    --prompt-txt assets/prompt.txt \
+    --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
+    --seed 42 \
+    --output-path outputs_video/longcat_480p

scripts/inference/v1_inference_longcat_distill.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+
+num_gpus=1
+export FASTVIDEO_ATTENTION_BACKEND=
+# For longcat, we must first convert the official weights to FastVideo native format
+# conversion method: python scripts/checkpoint_conversion/longcat_to_fastvideo.py
+# --source /path/to/LongCat-Video/weights/LongCat-Video
+# --output weights/longcat-native
+export MODEL_BASE=weights/longcat-native
+
+fastvideo generate \
+    --model-path $MODEL_BASE \
+    --sp-size $num_gpus \
+    --tp-size 1 \
+    --num-gpus $num_gpus \
+    --dit-cpu-offload False \
+    --vae-cpu-offload False \
+    --text-encoder-cpu-offload False \
+    --pin-cpu-memory False \
+    --enable-bsa False \
+    --lora-path "$MODEL_BASE/lora/distilled" \
+    --lora-nickname "distilled" \
+    --height 480 \
+    --width 832 \
+    --num-frames 93 \
+    --num-inference-steps 16 \
+    --fps 15 \
+    --guidance-scale 1.0 \
+    --prompt "In a realistic photography style, an asian boy around seven or eight years old sits on a park bench, wearing a light yellow T-shirt, denim shorts, and white sneakers. He holds an ice cream cone with vanilla and chocolate flavors, and beside him is a medium-sized golden Labrador. Smiling, the boy offers the ice cream to the dog, who eagerly licks it with its tongue. The sun is shining brightly, and the background features a green lawn and several tall trees, creating a warm and loving scene." \
+    --seed 42 \
+    --output-path outputs_video/longcat_distill

scripts/inference/v1_inference_longcat_refine_fromvideo.sh

@@ -0,0 +1,73 @@
+#!/bin/bash
+
+num_gpus=1
+export FASTVIDEO_ATTENTION_BACKEND=
+# For longcat, we must first convert the official weights to FastVideo native format
+# conversion method: python scripts/checkpoint_conversion/longcat_to_fastvideo.py
+# --source /path/to/LongCat-Video/weights/LongCat-Video
+# --output weights/longcat-native
+export MODEL_BASE=weights/longcat-native
+
+INPUT_VIDEO="outputs_video/longcat_distill/In a realistic photography style, an asian boy around seven or eight years old sits on a park bench,.mp4"
+REFINE_OUTPUT="outputs_video/longcat_refine_720p"
+
+# Prompt used for base generation
+PROMPT="In a realistic photography style, an asian boy around seven or eight years old sits on a park bench, wearing a light yellow T-shirt, denim shorts, and white sneakers. He holds an ice cream cone with vanilla and chocolate flavors, and beside him is a medium-sized golden Labrador. Smiling, the boy offers the ice cream to the dog, who eagerly licks it with its tongue. The sun is shining brightly, and the background features a green lawn and several tall trees, creating a warm and loving scene."
+
+echo "=========================================="
+echo "LongCat 480p -> 720p Refinement"
+echo "=========================================="
+echo ""
+echo "Input:  $INPUT_VIDEO"
+echo "Output: $REFINE_OUTPUT"
+echo ""
+
+# Check if input video exists
+if [ ! -f "$INPUT_VIDEO" ]; then
+    echo "Error: Input video not found: $INPUT_VIDEO"
+    echo "Please set INPUT_VIDEO to your 480p video path"
+    exit 1
+fi
+
+echo "🔧 Configuring refinement (BSA enabled, refinement LoRA)..."
+echo "✅ Input video: $INPUT_VIDEO"
+echo "✅ BSA enabled with sparsity=0.875"
+echo "✅ Refinement LoRA loaded"
+echo ""
+
+fastvideo generate \
+    --model-path $MODEL_BASE \
+    --sp-size $num_gpus \
+    --tp-size 1 \
+    --num-gpus $num_gpus \
+    --dit-cpu-offload True \
+    --vae-cpu-offload False \
+    --text-encoder-cpu-offload True \
+    --pin-cpu-memory False \
+    --enable-bsa True \
+    --bsa-sparsity 0.875 \
+    --bsa-chunk-q 4 4 8 \
+    --bsa-chunk-k 4 4 8 \
+    --lora-path "$MODEL_BASE/lora/refinement" \
+    --lora-nickname "refinement" \
+    --refine-from "$INPUT_VIDEO" \
+    --t-thresh 0.5 \
+    --spatial-refine-only False \
+    --num-cond-frames 0 \
+    --height 720 \
+    --width 1280 \
+    --num-inference-steps 50 \
+    --fps 30 \
+    --guidance-scale 1.0 \
+    --prompt "$PROMPT" \
+    --seed 42 \
+    --output-path "$REFINE_OUTPUT"
+
+echo ""
+echo "=========================================="
+echo "✓ Refinement Complete!"
+echo "=========================================="
+echo ""
+echo "Output directory: $REFINE_OUTPUT"
+echo ""
+