Spaces:
Running
Running
File size: 2,194 Bytes
8763fae 3d47c97 14f8699 8763fae ed3982d 8763fae d5b5eda 4f2f39b 237b649 14f8699 d5b5eda 8763fae db78d1c d56c06c 4f2f39b db78d1c 4f2f39b db78d1c 8763fae d8490c6 8763fae 4551a9d ab6b2b3 98a58d9 237b649 d8490c6 |
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 = StableDiffusionLatentUpscalePipeline.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, prompt, guidance):
# 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.1 4x Upscaler":
low_res_img = input_image.resize((128, 128))
upscaled_image = sd_2_1_4x(prompt, image=low_res_img, num_inference_steps=5, guidance_scale=guidance).images[0]
else:
upscaled_image = sd_2_0_2x(prompt, image=input_image, num_inference_steps=5, guidance_scale=guidance).images[0]
# Return the upscaled image
return upscaled_image
# Define the Gradio interface
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"),
gr.Textbox(label='Guide the AI Upscaling'),
gr.Slider(minimum=0, value=0, maximum=3, label='Guidance Scale')],
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(debug=True)
|