LongCat-Video

Checkpoints for LoRA training with musubi-tuner (relevant PR)
Converted from shards https://huggingface.co/meituan-longcat/LongCat-Video/tree/main/dit using the following script

import argparse
import itertools
import os
from musubi_tuner.utils.safetensors_utils import load_split_weights, MemoryEfficientSafeOpen
from safetensors.torch import save_file
import torch


def detect_dtype(path: str) -> torch.dtype:
    """Detect the dtype of the first floating point tensor in a safetensors file."""
    if not os.path.isfile(path):
        raise FileNotFoundError(f"File not found: {path}")

    with MemoryEfficientSafeOpen(path) as handle:
        keys = list(handle.keys())
        if not keys:
            raise ValueError(f"No tensors found in {path}")

        # Try to find a floating point tensor
        for key in keys:
            tensor = handle.get_tensor(key)
            if tensor.is_floating_point():
                dtype = tensor.dtype
                return dtype

        # If no floating point tensor, return dtype of first tensor
        return handle.get_tensor(keys[0]).dtype


def list_keys(state_dict, num_keys=20):
    """Display the first N keys from the state dict."""
    print(f"\nTotal tensors: {len(state_dict)}")
    print(f"First {num_keys} keys:")
    for key in itertools.islice(state_dict.keys(), num_keys):
        print(f"  {key}")
    print()


def convert_dtype(input_path: str, output_path: str, target_dtype: torch.dtype):
    """Convert safetensors file to target dtype."""
    print(f"Loading from: {input_path}")

    # Detect current dtype
    current_dtype = detect_dtype(input_path)
    print(f"Detected input dtype: {current_dtype}")
    print(f"Target dtype: {target_dtype}")

    # Load the model
    state_dict = load_split_weights(input_path)

    # List keys before conversion
    list_keys(state_dict)

    # Convert tensors
    print(f"Converting floating point tensors to {target_dtype}...")
    converted_count = 0
    for key, tensor in state_dict.items():
        if tensor.is_floating_point() and tensor.dtype != target_dtype:
            state_dict[key] = tensor.to(dtype=target_dtype)
            converted_count += 1

    print(f"Converted {converted_count} tensors")

    # Save the output
    print(f"Saving to: {output_path}")
    save_file(state_dict, output_path)
    print("Done!")


def main():
    parser = argparse.ArgumentParser(
        description="Convert safetensors file dtype with inspection and detection"
    )
    parser.add_argument(
        "input_path",
        type=str,
        help="Path to input safetensors file"
    )
    parser.add_argument(
        "output_path",
        type=str,
        help="Path to output safetensors file"
    )
    parser.add_argument(
        "--target-dtype",
        type=str,
        default="float16",
        choices=["float32", "float16", "bfloat16", "float8_e4m3fn", "float8_e5m2"],
        help="Target dtype for conversion (default: float16)"
    )

    args = parser.parse_args()

    # Map string dtype to torch dtype
    dtype_map = {
        "float32": torch.float32,
        "float16": torch.float16,
        "bfloat16": torch.bfloat16,
        "float8_e4m3fn": torch.float8_e4m3fn,
        "float8_e5m2": torch.float8_e5m2,
    }

    target_dtype = dtype_map[args.target_dtype]

    convert_dtype(args.input_path, args.output_path, target_dtype)


if name == "main":
    main()