first commit

app.py

@@ -0,0 +1,59 @@
+import gradio as gr
+import cv2
+import requests
+import numpy as np
+from PIL import Image
+from io import BytesIO
+from rembg import remove
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+import torch
+
+def remove_background(input_image: Image.Image, to_grayscale: bool) -> Image.Image:
+    output_image = remove(input_image)
+    if to_grayscale:
+        output_image = convert_to_grayscale(output_image)
+    return output_image
+
+def convert_to_grayscale(image: Image.Image) -> Image.Image:
+    return image.convert("L")
+
+def canny_image(image: Image.Image) -> Image.Image:
+    np_image = np.array(image)
+    low_threshold = 100
+    high_threshold = 200
+    np_image = cv2.Canny(np_image, low_threshold, high_threshold)
+    np_image = np_image[:, :, None]
+    np_image = np.concatenate([np_image, np_image, np_image], axis=2)
+    return Image.fromarray(np_image)
+
+def process_image(input_image: Image.Image, to_grayscale: bool, prompt: str) -> Image.Image:
+    output_image = remove_background(input_image, to_grayscale)
+    canny_output = canny_image(output_image)
+
+    controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
+    )
+
+    output = pipe(
+        [prompt],
+        canny_output,
+        negative_prompt=["monochrome, lowres, bad anatomy, worst quality, low quality"],
+        generator=[torch.Generator(device="cpu").manual_seed(2)],
+        num_inference_steps=20,
+    )
+
+    return image_grid(output.images, 1, 1)
+
+image_input = gr.components.Image(label="Input Image")
+grayscale_checkbox = gr.components.Checkbox(label="Convert output to grayscale", default=False)
+prompt_input = gr.components.Textbox(lines=1, label="Prompt")
+image_output = gr.components.Image(label="Output Image", type="pil")
+
+gr.Interface(
+    fn=process_image,
+    inputs=[image_input, grayscale_checkbox, prompt_input],
+    outputs=image_output,
+    title="Background Removal",
+    description="Upload an image to remove its background and optionally convert it to grayscale",
+).launch()