from diffusers import StableDiffusionUpscalePipeline
import torch
from PIL import Image

# Cargar el pipeline
model_id = "stabilityai/stable-diffusion-x4-upscaler"
pipe = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cpu")

# Cargar la imagen a mejorar
input_image = Image.open("mi_imagen.jpg").convert("RGB")

# Prompt opcional: describe cómo quieres mejorar la imagen
prompt = "A high quality, detailed, and visually clear image"

# Ejecutar upscaling
upscaled_image = pipe(prompt=prompt, image=input_image, scale=4).images[0]

# Guardar resultado
upscaled_image.save("mi_imagen_upscaled.png")

print("Imagen upscalada guardada como mi_imagen_upscaled.png")