update

diffusion_webui/controlnet_inpaint/canny_inpaint.py

@@ -0,0 +1,176 @@
+import cv2
+import gradio as gr
+import numpy as np
+import torch
+from diffusers import (
+    ControlNetModel,
+    StableDiffusionControlNetPipeline,
+    UniPCMultistepScheduler,
+)
+from PIL import Image
+
+stable_model_list = [
+    "runwayml/stable-diffusion-v1-5",
+    "stabilityai/stable-diffusion-2-1",
+]
+
+controlnet_canny_model_list = [
+    "lllyasviel/sd-controlnet-canny",
+    "thibaud/controlnet-sd21-canny-diffusers",
+]
+
+
+stable_prompt_list = ["a photo of a man.", "a photo of a girl."]
+
+stable_negative_prompt_list = ["bad, ugly", "deformed"]
+
+data_list = [
+    "data/test.png",
+]
+
+
+def controlnet_canny(
+    dict_image: str,
+    controlnet_model_path: str,
+):
+    image = dict_image["image"].convert("RGB").resize((512, 512))
+    image = np.array(image)
+
+    image = cv2.Canny(image, 100, 200)
+    image = image[:, :, None]
+    image = np.concatenate([image, image, image], axis=2)
+    image = Image.fromarray(image)
+
+    controlnet = ControlNetModel.from_pretrained(
+        controlnet_model_path, torch_dtype=torch.float16
+    )
+    return controlnet, image
+
+
+def stable_diffusion_controlnet_canny(
+    image_path: str,
+    stable_model_path: str,
+    controlnet_model_path: str,
+    prompt: str,
+    negative_prompt: str,
+    guidance_scale: int,
+    num_inference_step: int,
+):
+
+    controlnet, image = controlnet_canny(
+        image_path=image_path, controlnet_model_path=controlnet_model_path
+    )
+
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(
+        pretrained_model_name_or_path=stable_model_path,
+        controlnet=controlnet,
+        safety_checker=None,
+        torch_dtype=torch.float16,
+    )
+    pipe.to("cuda")
+    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+    pipe.enable_xformers_memory_efficient_attention()
+
+    output = pipe(
+        prompt=prompt,
+        image=image,
+        negative_prompt=negative_prompt,
+        num_inference_steps=num_inference_step,
+        guidance_scale=guidance_scale,
+    ).images
+
+    return output[0]
+
+
+def stable_diffusion_controlnet_canny_app():
+    with gr.Blocks():
+        with gr.Row():
+            with gr.Column():
+                controlnet_canny_image_file = gr.Image(
+                    type="filepath", label="Image"
+                )
+
+                controlnet_canny_stable_model_id = gr.Dropdown(
+                    choices=stable_model_list,
+                    value=stable_model_list[0],
+                    label="Stable Model Id",
+                )
+
+                controlnet_canny_model_id = gr.Dropdown(
+                    choices=controlnet_canny_model_list,
+                    value=controlnet_canny_model_list[0],
+                    label="Controlnet Model Id",
+                )
+
+                controlnet_canny_prompt = gr.Textbox(
+                    lines=1, value=stable_prompt_list[0], label="Prompt"
+                )
+
+                controlnet_canny_negative_prompt = gr.Textbox(
+                    lines=1,
+                    value=stable_negative_prompt_list[0],
+                    label="Negative Prompt",
+                )
+
+                with gr.Accordion("Advanced Options", open=False):
+                    controlnet_canny_guidance_scale = gr.Slider(
+                        minimum=0.1,
+                        maximum=15,
+                        step=0.1,
+                        value=7.5,
+                        label="Guidance Scale",
+                    )
+
+                    controlnet_canny_num_inference_step = gr.Slider(
+                        minimum=1,
+                        maximum=100,
+                        step=1,
+                        value=50,
+                        label="Num Inference Step",
+                    )
+
+                controlnet_canny_predict = gr.Button(value="Generator")
+
+            with gr.Column():
+                output_image = gr.Image(label="Output")
+
+        gr.Examples(
+            fn=stable_diffusion_controlnet_canny,
+            examples=[
+                [
+                    data_list[0],
+                    stable_model_list[0],
+                    controlnet_canny_model_list[0],
+                    stable_prompt_list[0],
+                    stable_negative_prompt_list[0],
+                    7.5,
+                    50,
+                ]
+            ],
+            inputs=[
+                controlnet_canny_image_file,
+                controlnet_canny_stable_model_id,
+                controlnet_canny_model_id,
+                controlnet_canny_prompt,
+                controlnet_canny_negative_prompt,
+                controlnet_canny_guidance_scale,
+                controlnet_canny_num_inference_step,
+            ],
+            outputs=[output_image],
+            cache_examples=False,
+            label="Controlnet Canny Example",
+        )
+
+        controlnet_canny_predict.click(
+            fn=stable_diffusion_controlnet_canny,
+            inputs=[
+                controlnet_canny_image_file,
+                controlnet_canny_stable_model_id,
+                controlnet_canny_model_id,
+                controlnet_canny_prompt,
+                controlnet_canny_negative_prompt,
+                controlnet_canny_guidance_scale,
+                controlnet_canny_num_inference_step,
+            ],
+            outputs=[output_image],
+        )

diffusion_webui/controlnet_inpaint/controlnet_inpaint_app.py

@@ -4,14 +4,14 @@ import torch
 from diffusers import UniPCMultistepScheduler
 from PIL import Image
 
-from diffusion_webui.controlnet.controlnet_canny import controlnet_canny
+from diffusion_webui.controlnet_inpaint.canny_inpaint import controlnet_canny
 from diffusion_webui.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import (
     StableDiffusionControlNetInpaintPipeline,
 )
 
 stable_inpaint_model_list = [
-    "stabilityai/stable-diffusion-2-inpainting",
     "runwayml/stable-diffusion-inpainting",
+    "stabilityai/stable-diffusion-2-inpainting",
 ]
 
 controlnet_model_list = [
@@ -36,7 +36,7 @@ def load_img(image_path: str):
 
 
 def stable_diffusion_inpiant_controlnet_canny(
-    normal_image_path: str,
+    dict_image: str,
     stable_model_path: str,
     controlnet_model_path: str,
     prompt: str,
@@ -45,15 +45,11 @@ def stable_diffusion_inpiant_controlnet_canny(
     guidance_scale: int,
     num_inference_steps: int,
 ):
-    pil_image = Image.open(normal_image_path)
-    normal_image = pil_image["image"].convert("RGB").resize((512, 512))
-    mask_image = pil_image["mask"].convert("RGB").resize((512, 512))
-
-    # normal_image = load_img(normal_image_path)
-    # mask_image = load_img(mask_image_path)
+    normal_image = dict_image["image"].convert("RGB").resize((512, 512))
+    mask_image = dict_image["mask"].convert("RGB").resize((512, 512))
 
     controlnet, control_image = controlnet_canny(
-        image_path=normal_image_path,
+        dict_image=dict_image,
         controlnet_model_path=controlnet_model_path,
     )
 
@@ -91,7 +87,7 @@ def stable_diffusion_inpiant_controlnet_canny_app():
                     source="upload",
                     tool="sketch",
                     elem_id="image_upload",
-                    type="filepath",
+                    type="pil",
                     label="Upload",
                 )