import gradio as gr
import torch
from PIL import Image
from diffusers import StableDiffusionImg2ImgPipeline
from huggingface_hub import hf_hub_download
from safetensors import safe_open

# Scarica il modello .safetensors dalla tua model card
model_path = hf_hub_download(
    repo_id="PietroC01/ImgEnhancerModels",
    filename="juggernaut_reborn.safetensors"
)

# Carica il modello Stable Diffusion
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",  # Modello base compatibile
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    use_safetensors=True,
    safety_checker=None
)

# Determina se è un LoRA o un modello completo
try:
    # Opzione 1: Carica come LoRA
    pipe.load_lora_weights(model_path)
except Exception as lora_error:
    try:
        # Opzione 2: Carica direttamente nell'UNet
        with safe_open(model_path, framework="pt", device="cpu") as f:
            unet_params = {k: f.get_tensor(k) for k in f.keys() if k.startswith("unet.")}
            
        # Rimuovi il prefisso "unet." dalle chiavi se presente
        clean_unet_params = {}
        for k, v in unet_params.items():
            if k.startswith("unet."):
                clean_unet_params[k[5:]] = v  # Rimuovi il prefisso "unet."
            else:
                clean_unet_params[k] = v
                
        # Carica i pesi nell'UNet
        if clean_unet_params:
            pipe.unet.load_state_dict(clean_unet_params, strict=False)
        else:
            # Opzione 3: Prova a caricare come checkpoint completo
            pipe = StableDiffusionImg2ImgPipeline.from_single_file(
                model_path,
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                use_safetensors=True,
                load_safety_checker=False
            )
    except Exception as e:
        print(f"Errore nel caricamento del modello: {e}")
        # Fallback: usa il modello base
        print("Utilizzo del modello base come fallback")

# Sposta il modello su GPU se disponibile
pipe.to("cuda" if torch.cuda.is_available() else "cpu")

def enhance_image(image, prompt, negative_prompt, cfg_scale, denoising_strength):
    """
    Migliora un'immagine utilizzando Stable Diffusion Img2Img.
    """
    image = Image.open(image).convert("RGB")
    enhanced_image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=image,
        strength=denoising_strength,
        guidance_scale=cfg_scale
    ).images[0]
    
    return enhanced_image

# Interfaccia utente con Gradio
demo = gr.Interface(
    fn=enhance_image,
    inputs=[
        gr.Image(type="filepath", label="Carica l'immagine:"),
        gr.Textbox(label="Prompt", value="highly detailed, ultra high resolution"),
        gr.Textbox(label="Negative Prompt", value="low quality, blurry, artifacts"),
        gr.Slider(1, 20, value=7, label="CFG Scale"),
        gr.Slider(0.1, 1.0, value=0.35, label="Denoising Strength")
    ],
    outputs=gr.Image(type="pil", label="Immagine migliorata"),
    title="Image Enhancer con Juggernaut Reborn",
    description="Carica un'immagine e usa il modello Juggernaut Reborn per migliorarla."
)

if __name__ == "__main__":
    demo.launch()