import gradio as gr
from diffusers import DiffusionPipeline
from PIL import Image
import torch

# Cargar pipeline (float32 para CPU)
model_id = "stabilityai/stable-diffusion-x4-upscaler"
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", dtype=torch.bfloat16, device_map="cuda")

# Función de Gradio
def upscaler_gradio(imagen):
    if imagen is None:
        return None

    img = Image.fromarray(imagen)  # convertir de NumPy a PIL
    prompt = "A high quality, detailed, and visually clear image"
    result = pipe(prompt=prompt, image=img).images[0]
    return result

# Interfaz Gradio
demo = gr.Interface(
    fn=upscaler_gradio,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Image(),
    title="Upscaler Difusión",
    description="Sube una imagen y el modelo la mejorará y aumentará resolución 4×."
)

demo.launch()