import torch
from safetensors.torch import load_file

# === Adjust these paths as needed ===
input_path = "nsfw_wan_umt5-xxl_bf16.safetensors"          # your original BF16 file
output_path = "nsfw_wan_umt5-xxl_fp16.pt"                  # will be saved as .pt

print(f"Loading model from: {input_path}")
state_dict = load_file(input_path, device='cpu')
print("Model loaded. Starting conversion to FP16...")

# Convert all tensors to FP16
converted_state_dict = {}
for key, value in state_dict.items():
    if isinstance(value, torch.Tensor):
        converted_state_dict[key] = value.to(torch.float16)
    else:
        converted_state_dict[key] = value  # keep non-tensor items as-is (rare)

# Optional: clean up any NaN / inf values that sometimes appear after casting
# Uncomment if you get warnings or bad generations later
# for key in converted_state_dict:
#     if isinstance(converted_state_dict[key], torch.Tensor):
#         converted_state_dict[key] = torch.nan_to_num(
#             converted_state_dict[key],
#             nan=0.0,
#             posinf=1e4,
#             neginf=-1e4
#         )

print(f"Saving converted model to: {output_path}")
torch.save(converted_state_dict, output_path)
print(f"Done! Converted FP16 model saved as: {output_path}")

# Quick size check (optional - helps confirm it didn't explode in memory)
total_params = sum(p.numel() for p in converted_state_dict.values() if isinstance(p, torch.Tensor))
print(f"Total parameters: {total_params:,}")
print(f"Approximate FP16 size on disk: ~{total_params * 2 / 1024**3:.1f} GB")