File size: 1,272 Bytes
70f6074 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
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)
|