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()
