|
import gradio as gr |
|
from diffusers import AutoPipelineForInpainting, AutoencoderKL |
|
import torch |
|
from PIL import Image, ImageOps |
|
|
|
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) |
|
pipeline = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", vae=vae, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda") |
|
|
|
def squarify_image(img): |
|
if(img.height > img.width): bg_size = img.height |
|
else: bg_size = img.width |
|
bg = Image.new(mode="RGB", size=(bg_size,bg_size), color="white") |
|
bg.paste(img, ( int((bg.width - bg.width)/2), 0) ) |
|
|
|
return bg |
|
|
|
def divisible_by_8(image): |
|
width, height = image.size |
|
|
|
|
|
new_width = (width // 8) * 8 |
|
new_height = (height // 8) * 8 |
|
|
|
|
|
resized_image = image.resize((new_width, new_height)) |
|
|
|
return resized_image |
|
|
|
|
|
def generate(image_editor, prompt, neg_prompt): |
|
image = image_editor['background'].convert('RGB') |
|
image.thumbnail((1024, 1024)) |
|
image = divisible_by_8(image) |
|
original_image_size = image.size |
|
|
|
layer = image_editor["layers"][0].resize(image.size) |
|
|
|
image = squarify_image(image) |
|
|
|
mask = Image.new("RGBA", image.size, "WHITE") |
|
mask.paste(layer, (0, 0), layer) |
|
mask = ImageOps.invert(mask.convert('L')) |
|
|
|
final_image = pipeline(prompt=prompt, |
|
image=image, |
|
mask_image=mask).images[0] |
|
|
|
final_image = final_image.crop((0, 0, original_image_size[0], original_image_size[1])) |
|
|
|
return image_editor, image, mask, final_image |
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown(""" |
|
# Inpainting Sketch Pad |
|
by [Tony Assi](https://www.tonyassi.com/) |
|
""") |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
sketch_pad = gr.ImageMask(type='pil', label='Inpaint') |
|
prompt = gr.Textbox() |
|
generate_button = gr.Button("Generate") |
|
with gr.Accordion("Advanced Settings", open=False): |
|
neg_prompt = gr.Textbox(label='Negative Prompt', value='ugly, deformed, nsfw') |
|
with gr.Column(): |
|
version_gallery = gr.Gallery(label="Versions") |
|
restore_button = gr.Button("Restore Version") |
|
|
|
|
|
with gr.Row(): |
|
out1 = gr.Image(format='png') |
|
out2 = gr.Image(format='png') |
|
out3 = gr.Image(format='png') |
|
|
|
generate_button.click(fn=generate, inputs=[sketch_pad,prompt, neg_prompt], outputs=[sketch_pad, out1, out2, out3]) |
|
|
|
demo.launch() |
