import torch
from safetensors.torch import save_file

# Load the original .bin file
orig = torch.load("pytorch_model.bin", map_location="cpu")
converted = {}

# Convert ip_adapter weights
for key, value in orig["ip_adapter"].items():
    # e.g. '0.to_k_ip.weight' → 'ip_adapter.0.to_k_ip.weight'
    parts = key.split(".", 1)
    if len(parts) == 2:
        new_key = f"ip_adapter.{parts[0]}.{parts[1]}"
    else:
        new_key = f"ip_adapter.{key}"
    converted[new_key] = value

# Convert image_proj weights
for key, value in orig["image_proj"].items():
    # Do not rename! Just prepend with 'image_proj.'
    new_key = f"image_proj.{key}"
    converted[new_key] = value

# Save to safetensors
save_file(converted, "ip_adapter.safetensors")