Spaces:
Running
Running
File size: 1,899 Bytes
8763fae 16e37d1 8763fae d4fd888 8763fae 1281d06 8763fae |
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 39 40 41 42 43 44 45 46 47 |
import gradio as gr
import torch
from torchvision.transforms import ToPILImage, ToTensor
from diffusers import StableDiffusionLatentUpscalePipeline, StableDiffusionUpscalePipeline
device = "cuda" if torch.cuda.is_available() else "cpu"
# Define the models
model_2x = "stabilityai/sd-x2-latent-upscaler"
model_4x = "stabilityai/stable-diffusion-x4-upscaler"
# Load the models
sd_2_0_2x = StableDiffusionLatentUpscalePipelin.from_pretrained(model_2x, torch_dtype=torch.float16, revision="fp16") if torch.cuda.is_available() else StableDiffusionLatentUpscalePipeline.from_pretrained(model_2x)
sd_2_1_4x = StableDiffusionUpscalePipeline.from_pretrained(model_4x, torch_dtype=torch.float16, revision="fp16") if torch.cuda.is_available() else StableDiffusionUpscalePipeline.from_pretrained(model_4x)
# Define the input and output components for the Gradio interface
input_image = gr.inputs.Image(type="filepath")
output_image = gr.outputs.Image(type="filepath")
# Define the function that will be called when the interface is used
def upscale_image(model, image):
# Convert the image to a PyTorch tensor
image_tensor = ToTensor()(image)
# Upscale the image using the selected model
if model == "SD 2.0 2x Latent Upscaler":
upscaled_tensor = sd_2_0_2x(image_tensor)
else:
upscaled_tensor = sd_2_1_4x(image_tensor)
# Convert the upscaled tensor back to a PIL image
upscaled_image = ToPILImage()(upscaled_tensor)
# Return the upscaled image
return upscaled_image
# Define the Gradio interface
iface = gr.Interface(
fn=upscale_image,
inputs=[gr.Radio(["SD 2.0 2x Latent Upscaler", "SD 2.1 4x Upscaler"]), input_image],
outputs=output_image,
title="Image Upscaler",
description="Upscale an image using either the SD 2.0 2x Latent Upscaler or the SD 2.1 4x Upscaler."
)
# Launch the interface
iface.launch(debug=True)
|