Create app.py

app.py

@@ -0,0 +1,89 @@
+import gradio as gr
+import cv2
+import numpy as np
+import os
+import torch
+from PIL import Image
+from diffusers import StableDiffusionInpaintPipeline
+
+auth_token = os.getenv("auth_token")
+
+
+def preview(image, state):
+    h, w = image.shape[:2]
+    scale_percent = 512 / max([w, h])
+
+    width = int(w * scale_percent)
+    height = int(h * scale_percent)
+    dim = (width, height)
+    resized = cv2.resize(image, dim, interpolation=cv2.INTER_AREA)
+    yoff = round((512-height)/2)
+    xoff = round((512-width)/2)
+
+    final_image = np.zeros((512, 512, 3), dtype=np.uint8)
+    final_image.fill(120)
+    final_image[yoff:yoff+height, xoff:xoff+width, :] = resized
+
+    mask_image = np.zeros((512, 512, 3), dtype=np.uint8)
+    mask_image.fill(255)
+    mask_image[yoff:yoff+height, xoff:xoff+width, :] = 0
+    state.clear()
+    state.append(mask_image)
+    state.append([yoff, xoff, height, width])
+    state.append(image)
+
+    return final_image, state
+
+
+def sd_inpaint(image, prompt, state):
+    mask = state[0]
+    yoff, xoff, height, width = state[1]
+    orig_image = state[2]
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+        "runwayml/stable-diffusion-inpainting",
+        revision="fp16",
+        torch_dtype=torch.float16,
+        use_auth_token=auth_token
+    ).to(device)
+
+    output = pipe(prompt=prompt, image=Image.fromarray(image), mask_image=Image.fromarray(mask)).images[0]
+    result = np.array(output)
+    result[yoff:yoff+height, xoff:xoff+width, :] = orig_image
+    result = Image.fromarray(result)
+
+    return result
+
+
+with gr.Blocks(title='Dreambooth Image Editing and Stable Diffusion Inpainting') as demo:
+    state = gr.State([])
+    gr.Markdown("# Dreambooth Image Editing and Stable Diffusion Inpainting")
+    gr.Markdown("It's difficult to get a good image to use for dreambooth, I do not have many photograhps of myself alone and it's very slow to edit the images (crop the selection, scale it to 512x512 and solve the problem of the background somehow)")
+    gr.Markdown("This app uses a combination of image selection, automatic scaling, and  stable diffusion inpainting to speed that process. Follow the next instructions:")
+    gr.Markdown("""- Upload an image
+- Use the select tool to select the area you want to use for dreambooth
+- The image will be resized to 512x512 and fill the rest of with a gray background
+- Then click the Inpaint button to use stable diffusion to inpaint the background
+- Save the image and use it for dreambooth
+    """)
+    with gr.Row():
+        with gr.Column():
+            img_ctr = gr.Image(tool='select')
+        with gr.Column():
+            output = gr.Image()
+    with gr.Row():
+        greet_btn = gr.Button("Selection")
+    with gr.Row():
+        sd_prompt = gr.Textbox(lines=2, label="Stable diffusion prompt")
+    with gr.Row():
+        final_image = gr.Image()
+    with gr.Row():
+        stab_btn = gr.Button("Inpaint")
+
+    greet_btn.click(fn=preview, inputs=[img_ctr, state], outputs=[output, state])
+    stab_btn.click(fn=sd_inpaint, inputs=[output, sd_prompt, state], outputs=final_image)
+
+
+demo.launch()