import gradio as gr
from diffusers import StableDiffusionUpscalePipeline
from PIL import Image

# Cargar pipeline (float32 para CPU)
model_id = "stabilityai/stable-diffusion-x4-upscaler"
pipe = StableDiffusionUpscalePipeline.from_pretrained(model_id)
pipe = pipe.to("cpu")  # o "cuda" si hay GPU

# 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, scale=4).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()