import torch
from safetensors.torch import load_file, save_file

# Paths - adjust these
input_path = "models_t5_umt5-xxl-enc-bf16_fully_uncensored.safetensors"   # your BF16 file
output_path = "models_t5_umt5-xxl-enc-fp16_fully_uncensored.safetensors"  # desired output

# Load the state dict (safetensors format)
state_dict = load_file(input_path, device='cpu')

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

# Optional: If you see any warnings about non-finite values, you can add clamping/cleanup
# 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)

# Save as new safetensors
save_file(converted_state_dict, output_path)

print(f"Converted and saved to: {output_path}")