Spaces:
Running
Running
File size: 2,025 Bytes
8763fae 3d47c97 14f8699 8763fae 16e37d1 8763fae d5b5eda 074954d 237b649 14f8699 d5b5eda 8763fae 8194643 8763fae d56c06c 8763fae 98a58d9 237b649 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 |
import gradio as gr
import torch
import modin.pandas as pd
from PIL import Image
from io import BytesIO
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 function that will be called when the interface is used
def upscale_image(model, input_image):
# Convert the image
generator = torch.manual_seed(999999)
input_image = Image.open(input_image).convert("RGB")
# Upscale the image using the selected model
if model == "SD 2.0 2x Latent Upscaler":
upscaled_image = sd_2_0_2x(prompt="", image=input_image, num_inference_steps=5).images[0]
else:
low_res_img = input_image.resize((128, 128))
upscaled_image = sd_2_1_4x(prompt="", image=low_res_img, num_inference_steps=5).images[0]
# 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"], label="Models:"), gr.Image(type="filepath", label = "Raw Image")],
outputs=gr.Image(type="filepath", label = "Upscaled Image"),
title="SD Image Upscaler",
description="Upscale an image using either the SD 2.0 2x Latent Upscaler or the SD 2.1 4x Upscaler. Use the 4x Upscaler for images lower than 512x512. Use the 2x Upscaler for images 512x512 to 768x768"
)
# Launch the interface
iface.launch(debug=True)
|