File size: 2,565 Bytes
f5d16a4 31089e0 f5d16a4 31089e0 53af09b 31089e0 53af09b 31089e0 f5d16a4 31089e0 f5d16a4 31089e0 |
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 |
import torch #needed only for GPU
from PIL import Image
from io import BytesIO
import numpy as np
from diffusers import StableDiffusionLatentUpscalePipeline, StableDiffusionUpscalePipeline
import gradio as gr
import modin.pandas as pd
# load model for CPU or GPU
model_2x = "stabilityai/sd-x2-latent-upscaler"
model_4x = "stabilityai/stable-diffusion-x4-upscaler"
device = "cuda" if torch.cuda.is_available() else "cpu"
upscaler2x = StableDiffusionLatentUpscalePipeline.from_pretrained(model_2x, torch_dtype=torch.float16, revision="fp16") if torch.cuda.is_available() else StableDiffusionLatentUpscalePipeline.from_pretrained(model_2x, safety_checker=None)
upscaler4x = StableDiffusionUpscalePipeline.from_pretrained(model_4x, torch_dtype=torch.float16, revision="fp16") if torch.cuda.is_available() else StableDiffusionUpscalePipeline.from_pretrained(model_4x)
#define interface
def upscale(raw_img, model, prompt, negative_prompt, scale, steps):
generator = torch.manual_seed(999999)
low_res_img = Image.open(raw_img).convert("RGB")
if model == "Upscaler 4x":
low_res_img = low_res_img.resize((128, 128))
else:
low_res_img
image = upscaler2x(prompt=prompt, negative_prompt=negative_prompt, image=low_res_img, guidance_scale=scale, num_inference_steps=steps).images[0] if model == "Upscaler 2x" else upscaler4x(prompt=prompt, negative_prompt=negative_prompt, image=low_res_img, guidance_scale=scale, num_inference_steps=steps).images[0]
return image
#launch interface
gr.Interface(fn=upscale, inputs=[
gr.Image(type="filepath", label='Lower Resolution Image'),
gr.Radio(['Upscaler 2x','Upscaler 4x'], label="Models"),
gr.Textbox(label="Optional: Enter a Prompt to Slightly Guide the AI's Enhancement"),
gr.Textbox(label='Experimental: Slightly influence What you do not want the AI to Enhance.'),
gr.Slider(2, 15, 7, step=1, label='Guidance Scale: How much the AI influences the Upscaling.'),
gr.Slider(5, 25, 10, step=1, label='Number of Iterations')],
outputs=gr.Image(type="filepath", label = 'Upscaled Image'),
title='SD Upscaler',
description='2x Latent Upscaler using SD 2.0 And 4x Upscaler using SD 2.1. This version runs on CPU or GPU and is currently running on a T4 GPU. For 4x Upscaling use images lower than 512x512. For 2x Upscaling use 512x512 to 768x768 images.<br><br><b>Notice: Largest Accepted Resolution is 768x768',
article = "Code Monkey: <a href=\"https://huggingface.co/Manjushri\">Manjushri</a>").launch(max_threads=True, debug=True) |