| import gradio as gr | |
| from PIL import Image | |
| from diffusers import StableDiffusionUpscalePipeline | |
| import torch | |
| import numpy as np | |
| # Load model and scheduler | |
| model_id = "stabilityai/stable-diffusion-x4-upscaler" | |
| pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16) | |
| pipeline = pipeline.to("cuda") | |
| def upscale_image(image, prompt): | |
| # Convert uploaded image to PIL | |
| low_res_img = Image.fromarray(image).convert("RGB") | |
| # Upscale the image | |
| upscaled_image = pipeline(prompt=prompt, image=low_res_img).images[0] | |
| # Convert upscaled PIL image back to numpy array for Gradio | |
| upscaled_image_np = np.array(upscaled_image) | |
| return upscaled_image_np | |
| # Create the Gradio interface | |
| interface = gr.Interface( | |
| fn=upscale_image, | |
| inputs=[ | |
| gr.Image(type="numpy", label="Upload Low-Resolution Image"), | |
| gr.Textbox(label="Upscaling Prompt", placeholder="Enter a prompt, e.g., 'a red box with glasses'") | |
| ], | |
| outputs=gr.Image(type="numpy", label="Upscaled Image"), | |
| title="Image Upscaler", | |
| description="Upload a low-resolution image and provide a prompt to upscale it using Stable Diffusion." | |
| ) | |
| # Launch the Gradio app | |
| interface.launch(share=True) | |