Upload 32 files

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 
-from diffusion_webui.helpers import (
+from diffusion_webui import (
     CodeformerUpscalerGenerator,
     StableDiffusionControlInpaintNetDepthGenerator,
     StableDiffusionControlNetCannyGenerator,
@@ -19,10 +19,11 @@ from diffusion_webui.helpers import (
     StableDiffusionImage2ImageGenerator,
     StableDiffusionInpaintGenerator,
     StableDiffusionText2ImageGenerator,
+    StableDiffusionControlNetNormalGenerator,
 )
 
 
-def main():
+def diffusion_app():
     app = gr.Blocks()
     with app:
         with gr.Row():
@@ -46,6 +47,8 @@ def main():
                         StableDiffusionControlNetPoseGenerator.app()
                     with gr.Tab("Scribble"):
                         StableDiffusionControlNetScribbleGenerator.app()
+                    with gr.Tab("Normal"):
+                        StableDiffusionControlNetNormalGenerator.app()
                     with gr.Tab("Seg"):
                         StableDiffusionControlNetSegGenerator.app()
                 with gr.Tab("ControlNet Inpaint"):
@@ -66,9 +69,9 @@ def main():
                 with gr.Tab("Upscaler"):
                     CodeformerUpscalerGenerator.app()
 
-    app.queue(concurrency_count=1)
+    app.queue(concurrency_count=2)
     app.launch(debug=True, enable_queue=True)
 
 
 if __name__ == "__main__":
-    main()
+    diffusion_app()

diffusion_webui/__init__.py

@@ -1 +1,30 @@
-__version__ = "2.0.1"
+from diffusion_webui.diffusion_models.stable_diffusion import (
+    StableDiffusionText2ImageGenerator,
+    StableDiffusionImage2ImageGenerator,
+    StableDiffusionInpaintGenerator,
+)
+
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint import (
+    StableDiffusionControlNetInpaintCannyGenerator,
+    StableDiffusionControlInpaintNetDepthGenerator,
+    StableDiffusionControlNetInpaintHedGenerator,
+    StableDiffusionControlNetInpaintMlsdGenerator,
+    StableDiffusionControlNetInpaintPoseGenerator,
+    StableDiffusionControlNetInpaintScribbleGenerator,
+    StableDiffusionControlNetInpaintSegGenerator,
+)   
+
+from diffusion_webui.diffusion_models.controlnet import (
+    StableDiffusionControlNetCannyGenerator,
+    StableDiffusionControlNetDepthGenerator,
+    StableDiffusionControlNetHEDGenerator,
+    StableDiffusionControlNetMLSDGenerator,
+    StableDiffusionControlNetNormalGenerator,
+    StableDiffusionControlNetPoseGenerator,
+    StableDiffusionControlNetScribbleGenerator,
+    StableDiffusionControlNetSegGenerator,
+)
+
+from diffusion_webui.app import diffusion_app as app
+
+__version__ = "2.2.0"

diffusion_webui/app.py

@@ -0,0 +1,77 @@
+import gradio as gr
+
+from diffusion_webui import (
+    CodeformerUpscalerGenerator,
+    StableDiffusionControlInpaintNetDepthGenerator,
+    StableDiffusionControlNetCannyGenerator,
+    StableDiffusionControlNetDepthGenerator,
+    StableDiffusionControlNetHEDGenerator,
+    StableDiffusionControlNetInpaintCannyGenerator,
+    StableDiffusionControlNetInpaintHedGenerator,
+    StableDiffusionControlNetInpaintMlsdGenerator,
+    StableDiffusionControlNetInpaintPoseGenerator,
+    StableDiffusionControlNetInpaintScribbleGenerator,
+    StableDiffusionControlNetInpaintSegGenerator,
+    StableDiffusionControlNetMLSDGenerator,
+    StableDiffusionControlNetPoseGenerator,
+    StableDiffusionControlNetScribbleGenerator,
+    StableDiffusionControlNetSegGenerator,
+    StableDiffusionImage2ImageGenerator,
+    StableDiffusionInpaintGenerator,
+    StableDiffusionText2ImageGenerator,
+    StableDiffusionControlNetNormalGenerator,
+)
+
+
+def diffusion_app():
+    app = gr.Blocks()
+    with app:
+        with gr.Row():
+            with gr.Column():
+                with gr.Tab("Text2Img"):
+                    StableDiffusionText2ImageGenerator.app()
+                with gr.Tab("Img2Img"):
+                    StableDiffusionImage2ImageGenerator.app()
+                with gr.Tab("Inpaint"):
+                    StableDiffusionInpaintGenerator.app()
+                with gr.Tab("ControlNet"):
+                    with gr.Tab("Canny"):
+                        StableDiffusionControlNetCannyGenerator.app()
+                    with gr.Tab("Depth"):
+                        StableDiffusionControlNetDepthGenerator.app()
+                    with gr.Tab("HED"):
+                        StableDiffusionControlNetHEDGenerator.app()
+                    with gr.Tab("MLSD"):
+                        StableDiffusionControlNetMLSDGenerator.app()
+                    with gr.Tab("Pose"):
+                        StableDiffusionControlNetPoseGenerator.app()
+                    with gr.Tab("Scribble"):
+                        StableDiffusionControlNetScribbleGenerator.app()
+                    with gr.Tab("Normal"):
+                        StableDiffusionControlNetNormalGenerator.app()
+                    with gr.Tab("Seg"):
+                        StableDiffusionControlNetSegGenerator.app()
+                with gr.Tab("ControlNet Inpaint"):
+                    with gr.Tab("Canny"):
+                        StableDiffusionControlNetInpaintCannyGenerator.app()
+                    with gr.Tab("Depth"):
+                        StableDiffusionControlInpaintNetDepthGenerator.app()
+                    with gr.Tab("HED"):
+                        StableDiffusionControlNetInpaintHedGenerator.app()
+                    with gr.Tab("MLSD"):
+                        StableDiffusionControlNetInpaintMlsdGenerator.app()
+                    with gr.Tab("Pose"):
+                        StableDiffusionControlNetInpaintPoseGenerator.app()
+                    with gr.Tab("Scribble"):
+                        StableDiffusionControlNetInpaintScribbleGenerator.app()
+                    with gr.Tab("Seg"):
+                        StableDiffusionControlNetInpaintSegGenerator.app()
+                with gr.Tab("Upscaler"):
+                    CodeformerUpscalerGenerator.app()
+
+    app.queue(concurrency_count=2)
+    app.launch(debug=True, enable_queue=True)
+
+
+if __name__ == "__main__":
+    diffusion_app()

diffusion_webui/diffusion_models/controlnet/__init__.py

@@ -0,0 +1,8 @@
+from diffusion_webui.diffusion_models.controlnet.controlnet_canny import StableDiffusionControlNetCannyGenerator
+from diffusion_webui.diffusion_models.controlnet.controlnet_depth import StableDiffusionControlNetDepthGenerator
+from diffusion_webui.diffusion_models.controlnet.controlnet_hed import StableDiffusionControlNetHEDGenerator
+from diffusion_webui.diffusion_models.controlnet.controlnet_mlsd import StableDiffusionControlNetMLSDGenerator
+from diffusion_webui.diffusion_models.controlnet.controlnet_normal import StableDiffusionControlNetNormalGenerator
+from diffusion_webui.diffusion_models.controlnet.controlnet_pose import StableDiffusionControlNetPoseGenerator
+from diffusion_webui.diffusion_models.controlnet.controlnet_scribble import StableDiffusionControlNetScribbleGenerator
+from diffusion_webui.diffusion_models.controlnet.controlnet_seg import StableDiffusionControlNetSegGenerator

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/__init__.py

@@ -0,0 +1,7 @@
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_canny import StableDiffusionControlNetInpaintCannyGenerator
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_depth import StableDiffusionControlInpaintNetDepthGenerator
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_hed import StableDiffusionControlNetInpaintHedGenerator
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_mlsd import StableDiffusionControlNetInpaintMlsdGenerator
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_pose import StableDiffusionControlNetInpaintPoseGenerator
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_scribble import StableDiffusionControlNetInpaintScribbleGenerator
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_seg import StableDiffusionControlNetInpaintSegGenerator

diffusion_webui/diffusion_models/controlnet/controlnet_mlsd.py

@@ -4,7 +4,7 @@ from controlnet_aux import MLSDdetector
 from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
 from PIL import Image
 
-from diffusion_webui.utils.model_list import stable_model_list
+from diffusion_webui.utils.model_list import stable_model_list, controlnet_mlsd_model_list
 from diffusion_webui.utils.scheduler_list import (
     SCHEDULER_LIST,
     get_scheduler_list,
@@ -45,7 +45,8 @@ class StableDiffusionControlNetMLSDGenerator:
     def generate_image(
         self,
         image_path: str,
-        model_path: str,
+        stable_model_path: str,
+        controlnet_model_path: str,
         prompt: str,
         negative_prompt: str,
         num_images_per_prompt: int,
@@ -57,8 +58,8 @@ class StableDiffusionControlNetMLSDGenerator:
         image = self.controlnet_mlsd(image_path=image_path)
 
         pipe = self.load_model(
-            stable_model_path=model_path,
-            controlnet_model_path="lllyasviel/sd-controlnet-mlsd",
+            stable_model_path=stable_model_path,
+            controlnet_model_path=controlnet_model_path,
             scheduler=scheduler,
         )
 
@@ -124,6 +125,12 @@ class StableDiffusionControlNetMLSDGenerator:
 
                         with gr.Row():
                             with gr.Column():
+                                controlnet_mlsd_controlnet_model_id = gr.Dropdown(
+                                    choices=controlnet_mlsd_model_list,
+                                    value=controlnet_mlsd_model_list[0],
+                                    label="ControlNet Model Id",
+                                
+                                )
                                 controlnet_mlsd_scheduler = gr.Dropdown(
                                     choices=SCHEDULER_LIST,
                                     value=SCHEDULER_LIST[0],

diffusion_webui/diffusion_models/controlnet/controlnet_normal.py

@@ -0,0 +1,190 @@
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers.utils import load_image
+from transformers import pipeline
+from PIL import Image
+import gradio as gr
+import numpy as np
+import torch
+import cv2
+
+
+from diffusion_webui.utils.model_list import (
+    controlnet_normal_model_list,
+    stable_model_list,
+)
+from diffusion_webui.utils.scheduler_list import (
+    SCHEDULER_LIST,
+    get_scheduler_list,
+)
+
+
+class StableDiffusionControlNetNormalGenerator:
+    def __init__(self):
+        self.pipe = None
+
+    def load_model(self, stable_model_path, controlnet_model_path, scheduler):
+        if self.pipe is None:
+            controlnet = ControlNetModel.from_pretrained(
+                controlnet_model_path, torch_dtype=torch.float16
+            )
+            self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+                pretrained_model_name_or_path=stable_model_path,
+                controlnet=controlnet,
+                safety_checker=None,
+                torch_dtype=torch.float16,
+            )
+
+        self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
+        self.pipe.to("cuda")
+        self.pipe.enable_xformers_memory_efficient_attention()
+
+        return self.pipe
+
+    def controlnet_normal(
+        self,
+        image_path: str,
+    ):
+        image = load_image(image_path).convert("RGB")
+        depth_estimator = pipeline("depth-estimation", model ="Intel/dpt-hybrid-midas" )
+        image = depth_estimator(image)['predicted_depth'][0]
+        image = image.numpy()
+        image_depth = image.copy()
+        image_depth -= np.min(image_depth)
+        image_depth /= np.max(image_depth)
+        bg_threhold = 0.4
+        x = cv2.Sobel(image, cv2.CV_32F, 1, 0, ksize=3)
+        x[image_depth < bg_threhold] = 0
+        y = cv2.Sobel(image, cv2.CV_32F, 0, 1, ksize=3)
+        y[image_depth < bg_threhold] = 0
+        z = np.ones_like(x) * np.pi * 2.0
+        image = np.stack([x, y, z], axis=2)
+        image /= np.sum(image ** 2.0, axis=2, keepdims=True) ** 0.5
+        image = (image * 127.5 + 127.5).clip(0, 255).astype(np.uint8)
+        image = Image.fromarray(image)
+        return image
+
+    def generate_image(
+        self,
+        image_path: str,
+        stable_model_path: str,
+        controlnet_model_path: str,
+        prompt: str,
+        negative_prompt: str,
+        num_images_per_prompt: int,
+        guidance_scale: int,
+        num_inference_step: int,
+        scheduler: str,
+        seed_generator: int,
+    ):
+        pipe = self.load_model(stable_model_path, controlnet_model_path, scheduler)
+        image = self.controlnet_normal(image_path)
+
+        if seed_generator == 0:
+            random_seed = torch.randint(0, 1000000, (1,))
+            generator = torch.manual_seed(random_seed)
+        else:
+            generator = torch.manual_seed(seed_generator)
+ 
+        output = pipe(
+            prompt=prompt,
+            image=image,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            num_inference_steps=num_inference_step,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        ).images
+
+        return output
+
+    def app():
+        with gr.Blocks():
+            with gr.Row():
+                with gr.Column():
+                    controlnet_normal_image_file = gr.Image(
+                        type="filepath", label="Image"
+                    )
+
+                    controlnet_normal_prompt = gr.Textbox(
+                        lines=1,
+                        placeholder="Prompt",
+                        show_label=False,
+                    )
+
+                    controlnet_normal_negative_prompt = gr.Textbox(
+                        lines=1,
+                        placeholder="Negative Prompt",
+                        show_label=False,
+                    )
+                    with gr.Row():
+                        with gr.Column():
+                            controlnet_normal_stable_model_id = gr.Dropdown(
+                                choices=stable_model_list,
+                                value=stable_model_list[0],
+                                label="Stable Model Id",
+                            )
+
+                            controlnet_normal_guidance_scale = gr.Slider(
+                                minimum=0.1,
+                                maximum=15,
+                                step=0.1,
+                                value=7.5,
+                                label="Guidance Scale",
+                            )
+                            controlnet_normal_num_inference_step = gr.Slider(
+                                minimum=1,
+                                maximum=100,
+                                step=1,
+                                value=50,
+                                label="Num Inference Step",
+                            )
+                            controlnet_normal_num_images_per_prompt = gr.Slider(
+                                minimum=1,
+                                maximum=10,
+                                step=1,
+                                value=1,
+                                label="Number Of Images",
+                            )
+                        with gr.Row():
+                            with gr.Column():
+                                controlnet_normal_model_id = gr.Dropdown(
+                                    choices=controlnet_normal_model_list,
+                                    value=controlnet_normal_model_list[0],
+                                    label="ControlNet Model Id",
+                                )
+
+                                controlnet_normal_scheduler = gr.Dropdown(
+                                    choices=SCHEDULER_LIST,
+                                    value=SCHEDULER_LIST[0],
+                                    label="Scheduler",
+                                )
+
+                                controlnet_normal_seed_generator = gr.Number(
+                                    value=0,
+                                    label="Seed Generator",
+                                )
+                    controlnet_normal_predict = gr.Button(value="Generator")
+
+                with gr.Column():
+                    output_image = gr.Gallery(
+                        label="Generated images",
+                        show_label=False,
+                        elem_id="gallery",
+                    ).style(grid=(1, 2))
+
+            controlnet_normal_predict.click(
+                fn=StableDiffusionControlNetCannyGenerator().generate_image,
+                inputs=[
+                    controlnet_normal_image_file,
+                    controlnet_normal_stable_model_id,
+                    controlnet_normal_model_id,
+                    controlnet_normal_prompt,
+                    controlnet_normal_negative_prompt,
+                    controlnet_normal_num_images_per_prompt,
+                    controlnet_normal_guidance_scale,
+                    controlnet_normal_num_inference_step,
+                    controlnet_normal_scheduler,
+                    controlnet_normal_seed_generator,
+                ],
+                outputs=[output_image],
+            )

diffusion_webui/diffusion_models/stable_diffusion/__init__.py

@@ -0,0 +1,3 @@
+from diffusion_webui.diffusion_models.stable_diffusion.text2img_app import StableDiffusionText2ImageGenerator
+from diffusion_webui.diffusion_models.stable_diffusion.img2img_app import StableDiffusionImage2ImageGenerator
+from diffusion_webui.diffusion_models.stable_diffusion.inpaint_app import StableDiffusionInpaintGenerator

diffusion_webui/helpers.py

@@ -1,54 +1,25 @@
-from diffusion_webui.diffusion_models.controlnet.controlnet_canny import (
-    StableDiffusionControlNetCannyGenerator,
-)
-from diffusion_webui.diffusion_models.controlnet.controlnet_depth import (
-    StableDiffusionControlNetDepthGenerator,
-)
-from diffusion_webui.diffusion_models.controlnet.controlnet_hed import (
-    StableDiffusionControlNetHEDGenerator,
+from diffusion_webui.diffusion_models.stable_diffusion import (
+    StableDiffusionText2ImageGenerator,
+    StableDiffusionImage2ImageGenerator,
+    StableDiffusionInpaintGenerator,
 )
-from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_canny import (
+
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint import (
     StableDiffusionControlNetInpaintCannyGenerator,
-)
-from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_depth import (
     StableDiffusionControlInpaintNetDepthGenerator,
-)
-from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_hed import (
     StableDiffusionControlNetInpaintHedGenerator,
-)
-from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_mlsd import (
     StableDiffusionControlNetInpaintMlsdGenerator,
-)
-from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_pose import (
     StableDiffusionControlNetInpaintPoseGenerator,
-)
-from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_scribble import (
     StableDiffusionControlNetInpaintScribbleGenerator,
-)
-from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.controlnet_inpaint_seg import (
     StableDiffusionControlNetInpaintSegGenerator,
-)
-from diffusion_webui.diffusion_models.controlnet.controlnet_mlsd import (
+)   
+
+from diffusion_webui.diffusion_models.controlnet import (
+    StableDiffusionControlNetCannyGenerator,
+    StableDiffusionControlNetDepthGenerator,
+    StableDiffusionControlNetHEDGenerator,
     StableDiffusionControlNetMLSDGenerator,
-)
-from diffusion_webui.diffusion_models.controlnet.controlnet_pose import (
     StableDiffusionControlNetPoseGenerator,
-)
-from diffusion_webui.diffusion_models.controlnet.controlnet_scribble import (
     StableDiffusionControlNetScribbleGenerator,
-)
-from diffusion_webui.diffusion_models.controlnet.controlnet_seg import (
     StableDiffusionControlNetSegGenerator,
 )
-from diffusion_webui.diffusion_models.stable_diffusion.img2img_app import (
-    StableDiffusionImage2ImageGenerator,
-)
-from diffusion_webui.diffusion_models.stable_diffusion.inpaint_app import (
-    StableDiffusionInpaintGenerator,
-)
-from diffusion_webui.diffusion_models.stable_diffusion.text2img_app import (
-    StableDiffusionText2ImageGenerator,
-)
-from diffusion_webui.upscaler_models.codeformer_upscaler import (
-    CodeformerUpscalerGenerator,
-)

diffusion_webui/utils/model_list.py

@@ -6,22 +6,22 @@ stable_model_list = [
     "dreamlike-art/dreamlike-diffusion-1.0",
     "gsdf/Counterfeit-V2.5",
     "dreamlike-art/dreamlike-photoreal-2.0"
-    
-    
 ]
 
 controlnet_canny_model_list = [
     "lllyasviel/sd-controlnet-canny",
     "lllyasviel/control_v11p_sd15_canny",
     "thibaud/controlnet-sd21-canny-diffusers",
+
 ]
 
 controlnet_depth_model_list = [
     "lllyasviel/sd-controlnet-depth",
-    "lllyasviel/control_v11p_sd15_depth",
+    "lllyasviel/control_v11f1p_sd15_depth",
     "thibaud/controlnet-sd21-depth-diffusers",
 ]
 
+
 controlnet_pose_model_list = [
     "lllyasviel/sd-controlnet-openpose",
     "lllyasviel/control_v11p_sd15_openpose",
@@ -30,12 +30,12 @@ controlnet_pose_model_list = [
 
 controlnet_hed_model_list = [
     "lllyasviel/sd-controlnet-hed",
-    "lllyasviel/control_v11p_sd15_scribble",
     "thibaud/controlnet-sd21-hed-diffusers",
 ]
 
 controlnet_scribble_model_list = [
     "lllyasviel/sd-controlnet-scribble",
+    "lllyasviel/control_v11p_sd15_scribble",
     "thibaud/controlnet-sd21-scribble-diffusers",
 ]
 stable_inpiant_model_list = [
@@ -48,6 +48,11 @@ controlnet_mlsd_model_list = [
     "lllyasviel/control_v11p_sd15_mlsd",
 ]
 
+controlnet_normal_model_list = [
+    "lllyasviel/sd-controlnet-normal",
+    "lllyasviel/control_v11p_sd15_normalbae",
+]
+
 controlnet_seg_model_list = [
     "lllyasviel/sd-controlnet-seg",
     "lllyasviel/control_v11p_sd15_seg",