# converter_lora.py
import torch
import safetensors.torch
from diffusers.loaders import LoraLoaderMixin

def convert_comfyui_lora_to_diffusers(input_path, output_path):
    """
    Converte LoRA do formato ComfyUI para formato diffusers
    """
    # Carrega o LoRA original
    state_dict = safetensors.torch.load_file(input_path)
    
    # Novo dicionário para formato diffusers
    diffusers_dict = {}
    
    for key, value in state_dict.items():
        # Remove chaves problemáticas
        if "img_in.alpha" in key:
            print(f"Skipping incompatible key: {key}")
            continue
            
        # Converte nomenclatura ComfyUI para diffusers
        new_key = key
        
        # Ajusta as chaves para formato diffusers
        if "model.diffusion_model" in key:
            new_key = key.replace("model.diffusion_model", "transformer")
        elif "lora_unet" in key:
            new_key = key.replace("lora_unet", "transformer")
        
        # Adiciona prefixo se necessário
        if not new_key.startswith(("transformer.", "text_encoder.", "vae.")):
            if "lora" in new_key:
                new_key = f"transformer.{new_key}"
        
        diffusers_dict[new_key] = value
    
    # Salva no formato diffusers
    safetensors.torch.save_file(
        diffusers_dict, 
        output_path,
        metadata={"format": "diffusers"}
    )
    print(f"✅ Converted LoRA saved to {output_path}")
    return diffusers_dict

# Uso no seu código principal
def load_and_convert_lora():
    from huggingface_hub import hf_hub_download
    import os
    
    # Baixa o arquivo original
    original_lora = hf_hub_download(
        repo_id=BFS_LORA,
        filename=BFS_LORA_FILENAME
    )
    
    # Define caminho para versão convertida
    converted_path = "/tmp/bfs_lora_diffusers.safetensors"
    
    # Converte
    converted_dict = convert_comfyui_lora_to_diffusers(
        original_lora, 
        converted_path
    )
    
    return converted_path