Upload 36 files

app.py

@@ -12,14 +12,19 @@ from diffusion_webui import (
     StableDiffusionControlNetInpaintPoseGenerator,
     StableDiffusionControlNetInpaintScribbleGenerator,
     StableDiffusionControlNetInpaintSegGenerator,
+    StableDiffusionControlNetLineArtAnimeGenerator,
+    StableDiffusionControlNetLineArtGenerator,
     StableDiffusionControlNetMLSDGenerator,
+    StableDiffusionControlNetNormalGenerator,
+    StableDiffusionControlNetPix2PixGenerator,
     StableDiffusionControlNetPoseGenerator,
     StableDiffusionControlNetScribbleGenerator,
     StableDiffusionControlNetSegGenerator,
+    StableDiffusionControlNetShuffleGenerator,
+    StableDiffusionControlNetSoftEdgeGenerator,
     StableDiffusionImage2ImageGenerator,
     StableDiffusionInpaintGenerator,
     StableDiffusionText2ImageGenerator,
-    StableDiffusionControlNetNormalGenerator,
 )
 
 
@@ -51,6 +56,16 @@ def diffusion_app():
                         StableDiffusionControlNetNormalGenerator.app()
                     with gr.Tab("Seg"):
                         StableDiffusionControlNetSegGenerator.app()
+                    with gr.Tab("Shuffle"):
+                        StableDiffusionControlNetShuffleGenerator.app()
+                    with gr.Tab("Pix2Pix"):
+                        StableDiffusionControlNetPix2PixGenerator.app()
+                    with gr.Tab("LineArt"):
+                        StableDiffusionControlNetLineArtGenerator.app()
+                    with gr.Tab("LineArtAnime"):
+                        StableDiffusionControlNetLineArtAnimeGenerator.app()
+                    with gr.Tab("SoftEdge"):
+                        StableDiffusionControlNetSoftEdgeGenerator.app()
                 with gr.Tab("ControlNet Inpaint"):
                     with gr.Tab("Canny"):
                         StableDiffusionControlNetInpaintCannyGenerator.app()

diffusion_webui/__init__.py

@@ -1,29 +1,32 @@
-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,
+    StableDiffusionControlNetLineArtAnimeGenerator,
+    StableDiffusionControlNetLineArtGenerator,
     StableDiffusionControlNetMLSDGenerator,
     StableDiffusionControlNetNormalGenerator,
+    StableDiffusionControlNetPix2PixGenerator,
     StableDiffusionControlNetPoseGenerator,
     StableDiffusionControlNetScribbleGenerator,
     StableDiffusionControlNetSegGenerator,
+    StableDiffusionControlNetShuffleGenerator,
+    StableDiffusionControlNetSoftEdgeGenerator,
+)
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint import (
+    StableDiffusionControlInpaintNetDepthGenerator,
+    StableDiffusionControlNetInpaintCannyGenerator,
+    StableDiffusionControlNetInpaintHedGenerator,
+    StableDiffusionControlNetInpaintMlsdGenerator,
+    StableDiffusionControlNetInpaintPoseGenerator,
+    StableDiffusionControlNetInpaintScribbleGenerator,
+    StableDiffusionControlNetInpaintSegGenerator,
+)
+from diffusion_webui.diffusion_models.stable_diffusion import (
+    StableDiffusionImage2ImageGenerator,
+    StableDiffusionInpaintGenerator,
+    StableDiffusionText2ImageGenerator,
 )
 from diffusion_webui.upscaler_models import CodeformerUpscalerGenerator
 
-__version__ = "2.2.0"
+__version__ = "2.4.0"

diffusion_webui/diffusion_models/controlnet/__init__.py

@@ -1,8 +1,39 @@
-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
+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_lineart import (
+    StableDiffusionControlNetLineArtGenerator,
+)
+from diffusion_webui.diffusion_models.controlnet.controlnet_lineart_anime import (
+    StableDiffusionControlNetLineArtAnimeGenerator,
+)
+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_pix2pix import (
+    StableDiffusionControlNetPix2PixGenerator,
+)
+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.controlnet.controlnet_shuffle import (
+    StableDiffusionControlNetShuffleGenerator,
+)
+from diffusion_webui.diffusion_models.controlnet.controlnet_softedge import (
+    StableDiffusionControlNetSoftEdgeGenerator,
+)

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/__init__.py

@@ -1,7 +1,21 @@
-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
+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_lineart.py

@@ -0,0 +1,178 @@
+import gradio as gr
+import torch
+from controlnet_aux import LineartDetector
+from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers.utils import load_image
+
+from diffusion_webui.utils.model_list import (
+    controlnet_lineart_model_list,
+    stable_model_list,
+)
+from diffusion_webui.utils.scheduler_list import (
+    SCHEDULER_LIST,
+    get_scheduler_list,
+)
+
+
+class StableDiffusionControlNetLineArtGenerator:
+    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_lineart(
+        self,
+        image_path: str,
+    ):
+        image = load_image(image_path)
+        image = image.resize((512, 512))
+        processor = LineartDetector.from_pretrained("lllyasviel/Annotators")
+        control_image = processor(image)
+        return control_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=stable_model_path,
+            controlnet_model_path=controlnet_model_path,
+            scheduler=scheduler,
+        )
+
+        image = self.controlnet_lineart(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_canny_image_file = gr.Image(
+                        type="filepath", label="Image"
+                    )
+
+                    controlnet_canny_prompt = gr.Textbox(
+                        lines=1,
+                        placeholder="Prompt",
+                        show_label=False,
+                    )
+
+                    controlnet_canny_negative_prompt = gr.Textbox(
+                        lines=1,
+                        placeholder="Negative Prompt",
+                        show_label=False,
+                    )
+                    with gr.Row():
+                        with gr.Column():
+                            controlnet_canny_stable_model_id = gr.Dropdown(
+                                choices=stable_model_list,
+                                value=stable_model_list[0],
+                                label="Stable Model Id",
+                            )
+
+                            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_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_canny_model_id = gr.Dropdown(
+                                    choices=controlnet_lineart_model_list,
+                                    value=controlnet_lineart_model_list[0],
+                                    label="ControlNet Model Id",
+                                )
+
+                                controlnet_canny_scheduler = gr.Dropdown(
+                                    choices=SCHEDULER_LIST,
+                                    value=SCHEDULER_LIST[0],
+                                    label="Scheduler",
+                                )
+
+                                controlnet_canny_seed_generator = gr.Number(
+                                    value=0,
+                                    label="Seed Generator",
+                                )
+                    controlnet_canny_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_canny_predict.click(
+                fn=StableDiffusionControlNetLineArtGenerator().generate_image,
+                inputs=[
+                    controlnet_canny_image_file,
+                    controlnet_canny_stable_model_id,
+                    controlnet_canny_model_id,
+                    controlnet_canny_prompt,
+                    controlnet_canny_negative_prompt,
+                    controlnet_canny_num_images_per_prompt,
+                    controlnet_canny_guidance_scale,
+                    controlnet_canny_num_inference_step,
+                    controlnet_canny_scheduler,
+                    controlnet_canny_seed_generator,
+                ],
+                outputs=[output_image],
+            )

diffusion_webui/diffusion_models/controlnet/controlnet_lineart_anime.py

@@ -0,0 +1,191 @@
+import gradio as gr
+import torch
+from controlnet_aux import LineartAnimeDetector
+from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers.utils import load_image
+from transformers import CLIPTextModel
+
+from diffusion_webui.utils.model_list import (
+    controlnet_lineart_anime_model_list,
+    stable_model_list,
+)
+from diffusion_webui.utils.scheduler_list import (
+    SCHEDULER_LIST,
+    get_scheduler_list,
+)
+
+
+class StableDiffusionControlNetLineArtAnimeGenerator:
+    def __init__(self):
+        self.pipe = None
+
+    def load_model(self, stable_model_path, controlnet_model_path, scheduler):
+        if self.pipe is None:
+            text_encoder = CLIPTextModel.from_pretrained(
+                stable_model_path,
+                subfolder="text_encoder",
+                num_hidden_layers=11,
+                torch_dtype=torch.float16,
+            )
+
+            controlnet = ControlNetModel.from_pretrained(
+                controlnet_model_path, torch_dtype=torch.float16
+            )
+            self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+                pretrained_model_name_or_path=stable_model_path,
+                text_encoder=text_encoder,
+                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_lineart_anime(
+        self,
+        image_path: str,
+    ):
+        image = load_image(image_path)
+        image = image.resize((512, 512))
+        processor = LineartAnimeDetector.from_pretrained(
+            "lllyasviel/Annotators"
+        )
+        control_image = processor(image)
+        return control_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=stable_model_path,
+            controlnet_model_path=controlnet_model_path,
+            scheduler=scheduler,
+        )
+
+        image = self.controlnet_lineart_anime(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_canny_image_file = gr.Image(
+                        type="filepath", label="Image"
+                    )
+
+                    controlnet_canny_prompt = gr.Textbox(
+                        lines=1,
+                        placeholder="Prompt",
+                        show_label=False,
+                    )
+
+                    controlnet_canny_negative_prompt = gr.Textbox(
+                        lines=1,
+                        placeholder="Negative Prompt",
+                        show_label=False,
+                    )
+                    with gr.Row():
+                        with gr.Column():
+                            controlnet_canny_stable_model_id = gr.Dropdown(
+                                choices=stable_model_list,
+                                value=stable_model_list[0],
+                                label="Stable Model Id",
+                            )
+
+                            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_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_canny_model_id = gr.Dropdown(
+                                    choices=controlnet_lineart_anime_model_list,
+                                    value=controlnet_lineart_anime_model_list[
+                                        0
+                                    ],
+                                    label="ControlNet Model Id",
+                                )
+
+                                controlnet_canny_scheduler = gr.Dropdown(
+                                    choices=SCHEDULER_LIST,
+                                    value=SCHEDULER_LIST[0],
+                                    label="Scheduler",
+                                )
+
+                                controlnet_canny_seed_generator = gr.Number(
+                                    value=0,
+                                    label="Seed Generator",
+                                )
+                    controlnet_canny_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_canny_predict.click(
+                fn=StableDiffusionControlNetLineArtAnimeGenerator().generate_image,
+                inputs=[
+                    controlnet_canny_image_file,
+                    controlnet_canny_stable_model_id,
+                    controlnet_canny_model_id,
+                    controlnet_canny_prompt,
+                    controlnet_canny_negative_prompt,
+                    controlnet_canny_num_images_per_prompt,
+                    controlnet_canny_guidance_scale,
+                    controlnet_canny_num_inference_step,
+                    controlnet_canny_scheduler,
+                    controlnet_canny_seed_generator,
+                ],
+                outputs=[output_image],
+            )

diffusion_webui/diffusion_models/controlnet/controlnet_mlsd.py

@@ -4,7 +4,10 @@ from controlnet_aux import MLSDdetector
 from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
 from PIL import Image
 
-from diffusion_webui.utils.model_list import stable_model_list, controlnet_mlsd_model_list
+from diffusion_webui.utils.model_list import (
+    controlnet_mlsd_model_list,
+    stable_model_list,
+)
 from diffusion_webui.utils.scheduler_list import (
     SCHEDULER_LIST,
     get_scheduler_list,
@@ -125,11 +128,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_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,

diffusion_webui/diffusion_models/controlnet/controlnet_normal.py

@@ -1,12 +1,11 @@
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
-from diffusers.utils import load_image
-from transformers import pipeline
-from PIL import Image
+import cv2
 import gradio as gr
 import numpy as np
 import torch
-import cv2
-
+from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers.utils import load_image
+from PIL import Image
+from transformers import pipeline
 
 from diffusion_webui.utils.model_list import (
     controlnet_normal_model_list,
@@ -45,8 +44,10 @@ class StableDiffusionControlNetNormalGenerator:
         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]
+        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)
@@ -76,7 +77,9 @@ class StableDiffusionControlNetNormalGenerator:
         scheduler: str,
         seed_generator: int,
     ):
-        pipe = self.load_model(stable_model_path, controlnet_model_path, scheduler)
+        pipe = self.load_model(
+            stable_model_path, controlnet_model_path, scheduler
+        )
         image = self.controlnet_normal(image_path)
 
         if seed_generator == 0:
@@ -84,7 +87,7 @@ class StableDiffusionControlNetNormalGenerator:
             generator = torch.manual_seed(random_seed)
         else:
             generator = torch.manual_seed(seed_generator)
- 
+
         output = pipe(
             prompt=prompt,
             image=image,

diffusion_webui/diffusion_models/controlnet/controlnet_pix2pix.py

@@ -0,0 +1,174 @@
+import gradio as gr
+import torch
+from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers.utils import load_image
+
+from diffusion_webui.utils.model_list import (
+    controlnet_lineart_model_list,
+    stable_model_list,
+)
+from diffusion_webui.utils.scheduler_list import (
+    SCHEDULER_LIST,
+    get_scheduler_list,
+)
+
+
+class StableDiffusionControlNetPix2PixGenerator:
+    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_pix2pix(
+        self,
+        image_path: str,
+    ):
+        control_image = load_image(image_path).convert("RGB")
+        return control_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=stable_model_path,
+            controlnet_model_path=controlnet_model_path,
+            scheduler=scheduler,
+        )
+
+        image = self.controlnet_pix2pix(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_canny_image_file = gr.Image(
+                        type="filepath", label="Image"
+                    )
+
+                    controlnet_canny_prompt = gr.Textbox(
+                        lines=1,
+                        placeholder="Prompt",
+                        show_label=False,
+                    )
+
+                    controlnet_canny_negative_prompt = gr.Textbox(
+                        lines=1,
+                        placeholder="Negative Prompt",
+                        show_label=False,
+                    )
+                    with gr.Row():
+                        with gr.Column():
+                            controlnet_canny_stable_model_id = gr.Dropdown(
+                                choices=stable_model_list,
+                                value=stable_model_list[0],
+                                label="Stable Model Id",
+                            )
+
+                            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_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_canny_model_id = gr.Dropdown(
+                                    choices=controlnet_lineart_model_list,
+                                    value=controlnet_lineart_model_list[0],
+                                    label="ControlNet Model Id",
+                                )
+
+                                controlnet_canny_scheduler = gr.Dropdown(
+                                    choices=SCHEDULER_LIST,
+                                    value=SCHEDULER_LIST[0],
+                                    label="Scheduler",
+                                )
+
+                                controlnet_canny_seed_generator = gr.Number(
+                                    value=0,
+                                    label="Seed Generator",
+                                )
+                    controlnet_canny_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_canny_predict.click(
+                fn=StableDiffusionControlNetPix2PixGenerator().generate_image,
+                inputs=[
+                    controlnet_canny_image_file,
+                    controlnet_canny_stable_model_id,
+                    controlnet_canny_model_id,
+                    controlnet_canny_prompt,
+                    controlnet_canny_negative_prompt,
+                    controlnet_canny_num_images_per_prompt,
+                    controlnet_canny_guidance_scale,
+                    controlnet_canny_num_inference_step,
+                    controlnet_canny_scheduler,
+                    controlnet_canny_seed_generator,
+                ],
+                outputs=[output_image],
+            )

diffusion_webui/diffusion_models/controlnet/controlnet_shuffle.py

@@ -0,0 +1,176 @@
+import gradio as gr
+import torch
+from controlnet_aux import ContentShuffleDetector
+from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers.utils import load_image
+
+from diffusion_webui.utils.model_list import (
+    controlnet_shuffle_model_list,
+    stable_model_list,
+)
+from diffusion_webui.utils.scheduler_list import (
+    SCHEDULER_LIST,
+    get_scheduler_list,
+)
+
+
+class StableDiffusionControlNetShuffleGenerator:
+    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_shuffle(
+        self,
+        image_path: str,
+    ):
+        image = load_image(image_path)
+        control_image = ContentShuffleDetector()(image)
+        return control_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=stable_model_path,
+            controlnet_model_path=controlnet_model_path,
+            scheduler=scheduler,
+        )
+
+        image = self.controlnet_shuffle(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_canny_image_file = gr.Image(
+                        type="filepath", label="Image"
+                    )
+
+                    controlnet_canny_prompt = gr.Textbox(
+                        lines=1,
+                        placeholder="Prompt",
+                        show_label=False,
+                    )
+
+                    controlnet_canny_negative_prompt = gr.Textbox(
+                        lines=1,
+                        placeholder="Negative Prompt",
+                        show_label=False,
+                    )
+                    with gr.Row():
+                        with gr.Column():
+                            controlnet_canny_stable_model_id = gr.Dropdown(
+                                choices=stable_model_list,
+                                value=stable_model_list[0],
+                                label="Stable Model Id",
+                            )
+
+                            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_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_canny_model_id = gr.Dropdown(
+                                    choices=controlnet_shuffle_model_list,
+                                    value=controlnet_shuffle_model_list[0],
+                                    label="ControlNet Model Id",
+                                )
+
+                                controlnet_canny_scheduler = gr.Dropdown(
+                                    choices=SCHEDULER_LIST,
+                                    value=SCHEDULER_LIST[0],
+                                    label="Scheduler",
+                                )
+
+                                controlnet_canny_seed_generator = gr.Number(
+                                    value=0,
+                                    label="Seed Generator",
+                                )
+                    controlnet_canny_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_canny_predict.click(
+                fn=StableDiffusionControlNetShuffleGenerator().generate_image,
+                inputs=[
+                    controlnet_canny_image_file,
+                    controlnet_canny_stable_model_id,
+                    controlnet_canny_model_id,
+                    controlnet_canny_prompt,
+                    controlnet_canny_negative_prompt,
+                    controlnet_canny_num_images_per_prompt,
+                    controlnet_canny_guidance_scale,
+                    controlnet_canny_num_inference_step,
+                    controlnet_canny_scheduler,
+                    controlnet_canny_seed_generator,
+                ],
+                outputs=[output_image],
+            )

diffusion_webui/diffusion_models/controlnet/controlnet_softedge.py

@@ -0,0 +1,179 @@
+import gradio as gr
+import torch
+from controlnet_aux import HEDdetector, PidiNetDetector
+from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers.utils import load_image
+
+from diffusion_webui.utils.model_list import (
+    controlnet_softedge_model_list,
+    stable_model_list,
+)
+from diffusion_webui.utils.scheduler_list import (
+    SCHEDULER_LIST,
+    get_scheduler_list,
+)
+
+
+class StableDiffusionControlNetSoftEdgeGenerator:
+    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_softedge(
+        self,
+        image_path: str,
+    ):
+
+        image = load_image(image_path)
+        processor = HEDdetector.from_pretrained("lllyasviel/Annotators")
+        processor = PidiNetDetector.from_pretrained("lllyasviel/Annotators")
+        control_image = processor(image, safe=True)
+        return control_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=stable_model_path,
+            controlnet_model_path=controlnet_model_path,
+            scheduler=scheduler,
+        )
+
+        image = self.controlnet_softedge(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_canny_image_file = gr.Image(
+                        type="filepath", label="Image"
+                    )
+
+                    controlnet_canny_prompt = gr.Textbox(
+                        lines=1,
+                        placeholder="Prompt",
+                        show_label=False,
+                    )
+
+                    controlnet_canny_negative_prompt = gr.Textbox(
+                        lines=1,
+                        placeholder="Negative Prompt",
+                        show_label=False,
+                    )
+                    with gr.Row():
+                        with gr.Column():
+                            controlnet_canny_stable_model_id = gr.Dropdown(
+                                choices=stable_model_list,
+                                value=stable_model_list[0],
+                                label="Stable Model Id",
+                            )
+
+                            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_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_canny_model_id = gr.Dropdown(
+                                    choices=controlnet_softedge_model_list,
+                                    value=controlnet_softedge_model_list[0],
+                                    label="ControlNet Model Id",
+                                )
+
+                                controlnet_canny_scheduler = gr.Dropdown(
+                                    choices=SCHEDULER_LIST,
+                                    value=SCHEDULER_LIST[0],
+                                    label="Scheduler",
+                                )
+
+                                controlnet_canny_seed_generator = gr.Number(
+                                    value=0,
+                                    label="Seed Generator",
+                                )
+                    controlnet_canny_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_canny_predict.click(
+                fn=StableDiffusionControlNetSoftEdgeGenerator().generate_image,
+                inputs=[
+                    controlnet_canny_image_file,
+                    controlnet_canny_stable_model_id,
+                    controlnet_canny_model_id,
+                    controlnet_canny_prompt,
+                    controlnet_canny_negative_prompt,
+                    controlnet_canny_num_images_per_prompt,
+                    controlnet_canny_guidance_scale,
+                    controlnet_canny_num_inference_step,
+                    controlnet_canny_scheduler,
+                    controlnet_canny_seed_generator,
+                ],
+                outputs=[output_image],
+            )

diffusion_webui/diffusion_models/stable_diffusion/__init__.py

@@ -1,3 +1,9 @@
-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
+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,
+)

diffusion_webui/upscaler_models/__init__.py

@@ -1 +1,3 @@
-from diffusion_webui.upscaler_models.codeformer_upscaler import CodeformerUpscalerGenerator
+from diffusion_webui.upscaler_models.codeformer_upscaler import (
+    CodeformerUpscalerGenerator,
+)

diffusion_webui/upscaler_models/codeformer_upscaler.py

@@ -1,5 +1,4 @@
 import gradio as gr
-from codeformer.app import inference_app
 
 
 class CodeformerUpscalerGenerator:
@@ -11,6 +10,7 @@ class CodeformerUpscalerGenerator:
         upscale: int,
         codeformer_fidelity: int,
     ):
+        from codeformer.app import inference_app
 
         pipe = inference_app(
             image=image_path,

diffusion_webui/utils/model_list.py

@@ -5,14 +5,13 @@ stable_model_list = [
     "wavymulder/Analog-Diffusion",
     "dreamlike-art/dreamlike-diffusion-1.0",
     "gsdf/Counterfeit-V2.5",
-    "dreamlike-art/dreamlike-photoreal-2.0"
+    "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 = [
@@ -57,3 +56,23 @@ controlnet_seg_model_list = [
     "lllyasviel/sd-controlnet-seg",
     "lllyasviel/control_v11p_sd15_seg",
 ]
+
+controlnet_shuffle_model_list = [
+    "lllyasviel/control_v11e_sd15_shuffle",
+]
+
+controlnet_pix2pix_model_list = [
+    "lllyasviel/control_v11e_sd15_ip2p",
+]
+
+controlnet_lineart_model_list = [
+    "ControlNet-1-1-preview/control_v11p_sd15_lineart",
+]
+
+controlnet_lineart_anime_model_list = [
+    "lllyasviel/control_v11p_sd15s2_lineart_anime",
+]
+
+controlnet_softedge_model_list = [
+    "lllyasviel/control_v11p_sd15_softedge",
+]