import gradio as gr
from diffusers import StableDiffusionImg2ImgPipeline
import torch
from PIL import Image
import io
import base64

# Carica il modello Stable Diffusion
pipe = StableDiffusionImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-2-1")
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 Stable Diffusion",
    description="Carica un'immagine e usa Stable Diffusion per migliorarla."
)

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