Upload 27 files

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_canny.py

@@ -8,7 +8,7 @@ from PIL import Image
 
 from diffusion_webui.utils.model_list import (
     controlnet_canny_model_list,
-    stable_model_list,
+    stable_inpiant_model_list,
 )
 from diffusion_webui.utils.scheduler_list import (
     SCHEDULER_LIST,
@@ -40,9 +40,9 @@ class StableDiffusionControlNetInpaintCannyGenerator:
 
         return self.pipe
     
-    def load_img(self, image_path):
+    def load_image(self, image_path):
         image = np.array(image_path)
-        image = Image.fromarray(image)
+        image = Image.fromarray(image_path)
         return image
 
     def controlnet_canny_inpaint(
@@ -77,8 +77,8 @@ class StableDiffusionControlNetInpaintCannyGenerator:
         normal_image = image_path["image"].convert("RGB").resize((512, 512))
         mask_image = image_path["mask"].convert("RGB").resize((512, 512))
         
-        normal_image = self.load_img(image_path=normal_image)
-        mask_image = self.load_img(image_path=mask_image)
+        normal_image = self.load_image(image_path=normal_image)
+        mask_image = self.load_image(image_path=mask_image)
         
         control_image = self.controlnet_canny_inpaint(image_path=image_path)
         pipe = self.load_model(
@@ -133,8 +133,8 @@ class StableDiffusionControlNetInpaintCannyGenerator:
                         with gr.Column():
                             controlnet_canny_inpaint_stable_model_id = (
                                 gr.Dropdown(
-                                    choices=stable_model_list,
-                                    value=stable_model_list[0],
+                                    choices=stable_inpiant_model_list,
+                                    value=stable_inpiant_model_list[0],
                                     label="Stable Model Id",
                                 )
                             )

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_depth.py

@@ -1,13 +1,14 @@
 import gradio as gr
 import numpy as np
 import torch
-from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers import ControlNetModel
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
 from PIL import Image
 from transformers import pipeline
 
 from diffusion_webui.utils.model_list import (
     controlnet_depth_model_list,
-    stable_model_list,
+    stable_inpiant_model_list,
 )
 from diffusion_webui.utils.scheduler_list import (
     SCHEDULER_LIST,
@@ -26,7 +27,7 @@ class StableDiffusionControlInpaintNetDepthGenerator:
             controlnet = ControlNetModel.from_pretrained(
                 controlnet_model_path, torch_dtype=torch.float16
             )
-            self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
                 pretrained_model_name_or_path=stable_model_path,
                 controlnet=controlnet,
                 safety_checker=None,
@@ -39,6 +40,12 @@ class StableDiffusionControlInpaintNetDepthGenerator:
 
         return self.pipe
 
+    def load_image(self, image_path):
+        image = np.array(image_path)
+        image = Image.fromarray(image_path)
+        return image
+
+
     def controlnet_inpaint_depth(self, image_path: str):
         depth_estimator = pipeline("depth-estimation")
         image = image_path["image"].convert("RGB").resize((512, 512))
@@ -64,8 +71,13 @@ class StableDiffusionControlInpaintNetDepthGenerator:
         scheduler: str,
         seed_generator: int,
     ):
-
-        image = self.controlnet_inpaint_depth(image_path=image_path)
+        normal_image = image_path["image"].convert("RGB").resize((512, 512))
+        mask_image = image_path["mask"].convert("RGB").resize((512, 512))
+        
+        normal_image = self.load_image(image_path=normal_image)
+        mask_image = self.load_image(image_path=mask_image)
+        
+        control_image = self.controlnet_inpaint_depth(image_path=image_path)
 
         pipe = self.load_model(
             stable_model_path=stable_model_path,
@@ -81,7 +93,10 @@ class StableDiffusionControlInpaintNetDepthGenerator:
 
         output = pipe(
             prompt=prompt,
-            image=image,
+            
+            image=normal_image,
+            mask_image=mask_image,
+            control_image=control_image,
             negative_prompt=negative_prompt,
             num_images_per_prompt=num_images_per_prompt,
             num_inference_steps=num_inference_step,
@@ -117,8 +132,8 @@ class StableDiffusionControlInpaintNetDepthGenerator:
                         with gr.Column():
                             controlnet_depth_inpaint_stable_model_id = (
                                 gr.Dropdown(
-                                    choices=stable_model_list,
-                                    value=stable_model_list[0],
+                                    choices=stable_inpiant_model_list,
+                                    value=stable_inpiant_model_list[0],
                                     label="Stable Model Id",
                                 )
                             )

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_hed.py

@@ -2,16 +2,18 @@ import gradio as gr
 import numpy as np
 import torch
 from controlnet_aux import HEDdetector
-from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers import ControlNetModel
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
 
 from diffusion_webui.utils.model_list import (
     controlnet_hed_model_list,
-    stable_model_list,
+    stable_inpiant_model_list,
 )
 from diffusion_webui.utils.scheduler_list import (
     SCHEDULER_LIST,
     get_scheduler_list,
 )
+from PIL import Image
 
 # https://github.com/mikonvergence/ControlNetInpaint
 
@@ -25,7 +27,7 @@ class StableDiffusionControlNetInpaintHedGenerator:
             controlnet = ControlNetModel.from_pretrained(
                 controlnet_model_path, torch_dtype=torch.float16
             )
-            self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
                 pretrained_model_name_or_path=stable_model_path,
                 controlnet=controlnet,
                 safety_checker=None,
@@ -38,6 +40,12 @@ class StableDiffusionControlNetInpaintHedGenerator:
 
         return self.pipe
 
+    def load_image(self, image_path):
+        image = np.array(image_path)
+        image = Image.fromarray(image_path)
+        return image
+
+
     def controlnet_inpaint_hed(self, image_path: str):
         hed = HEDdetector.from_pretrained("lllyasviel/ControlNet")
         image = image_path["image"].convert("RGB").resize((512, 512))
@@ -60,8 +68,13 @@ class StableDiffusionControlNetInpaintHedGenerator:
         scheduler: str,
         seed_generator: int,
     ):
-
-        image = self.controlnet_inpaint_hed(image_path=image_path)
+        normal_image = image_path["image"].convert("RGB").resize((512, 512))
+        mask_image = image_path["mask"].convert("RGB").resize((512, 512))
+        
+        normal_image = self.load_image(image_path=normal_image)
+        mask_image = self.load_image(image_path=mask_image)
+        
+        control_image = self.controlnet_inpaint_hed(image_path=image_path)
 
         pipe = self.load_model(
             stable_model_path=stable_model_path,
@@ -77,7 +90,9 @@ class StableDiffusionControlNetInpaintHedGenerator:
 
         output = pipe(
             prompt=prompt,
-            image=image,
+            image=normal_image,
+            mask_image=mask_image,
+            control_image=control_image,
             negative_prompt=negative_prompt,
             num_images_per_prompt=num_images_per_prompt,
             num_inference_steps=num_inference_step,
@@ -113,8 +128,8 @@ class StableDiffusionControlNetInpaintHedGenerator:
                         with gr.Column():
                             controlnet_hed_inpaint_stable_model_id = (
                                 gr.Dropdown(
-                                    choices=stable_model_list,
-                                    value=stable_model_list[0],
+                                    choices=stable_inpiant_model_list,
+                                    value=stable_inpiant_model_list[0],
                                     label="Stable Model Id",
                                 )
                             )

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_mlsd.py

@@ -2,16 +2,18 @@ import gradio as gr
 import numpy as np
 import torch
 from controlnet_aux import MLSDdetector
-from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers import ControlNetModel
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
 
 from diffusion_webui.utils.model_list import (
     controlnet_mlsd_model_list,
-    stable_model_list,
+    stable_inpiant_model_list,
 )
 from diffusion_webui.utils.scheduler_list import (
     SCHEDULER_LIST,
     get_scheduler_list,
 )
+from PIL import Image
 
 # https://github.com/mikonvergence/ControlNetInpaint
 
@@ -25,7 +27,7 @@ class StableDiffusionControlNetInpaintMlsdGenerator:
             controlnet = ControlNetModel.from_pretrained(
                 controlnet_model_path, torch_dtype=torch.float16
             )
-            self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
                 pretrained_model_name_or_path=stable_model_path,
                 controlnet=controlnet,
                 safety_checker=None,
@@ -37,6 +39,12 @@ class StableDiffusionControlNetInpaintMlsdGenerator:
         self.pipe.enable_xformers_memory_efficient_attention()
 
         return self.pipe
+    
+    def load_image(self, image_path):
+        image = np.array(image_path)
+        image = Image.fromarray(image_path)
+        return image
+
 
     def controlnet_inpaint_mlsd(self, image_path: str):
         mlsd = MLSDdetector.from_pretrained("lllyasviel/ControlNet")
@@ -61,7 +69,13 @@ class StableDiffusionControlNetInpaintMlsdGenerator:
         seed_generator: int,
     ):
 
-        image = self.controlnet_inpaint_mlsd(image_path=image_path)
+        normal_image = image_path["image"].convert("RGB").resize((512, 512))
+        mask_image = image_path["mask"].convert("RGB").resize((512, 512))
+        
+        normal_image = self.load_image(image_path=normal_image)
+        mask_image = self.load_image(image_path=mask_image)
+        
+        control_image = self.controlnet_inpaint_mlsd(image_path=image_path)
 
         pipe = self.load_model(
             stable_model_path=stable_model_path,
@@ -77,7 +91,9 @@ class StableDiffusionControlNetInpaintMlsdGenerator:
 
         output = pipe(
             prompt=prompt,
-            image=image,
+            image=normal_image,
+            mask_image=mask_image,
+            control_image=control_image,
             negative_prompt=negative_prompt,
             num_images_per_prompt=num_images_per_prompt,
             num_inference_steps=num_inference_step,
@@ -113,8 +129,8 @@ class StableDiffusionControlNetInpaintMlsdGenerator:
                         with gr.Column():
                             controlnet_mlsd_inpaint_stable_model_id = (
                                 gr.Dropdown(
-                                    choices=stable_model_list,
-                                    value=stable_model_list[0],
+                                    choices=stable_inpiant_model_list,
+                                    value=stable_inpiant_model_list[0],
                                     label="Stable Model Id",
                                 )
                             )

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_pose.py

@@ -2,11 +2,13 @@ import gradio as gr
 import numpy as np
 import torch
 from controlnet_aux import OpenposeDetector
-from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers import ControlNetModel
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
+from PIL import Image
 
 from diffusion_webui.utils.model_list import (
     controlnet_pose_model_list,
-    stable_model_list,
+    stable_inpiant_model_list,
 )
 from diffusion_webui.utils.scheduler_list import (
     SCHEDULER_LIST,
@@ -26,7 +28,7 @@ class StableDiffusionControlNetInpaintPoseGenerator:
                 controlnet_model_path, torch_dtype=torch.float16
             )
 
-            self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
                 pretrained_model_name_or_path=stable_model_path,
                 controlnet=controlnet,
                 safety_checker=None,
@@ -38,7 +40,12 @@ class StableDiffusionControlNetInpaintPoseGenerator:
         self.pipe.enable_xformers_memory_efficient_attention()
 
         return self.pipe
-
+    
+    def load_image(self, image_path):
+        image = np.array(image_path)
+        image = Image.fromarray(image_path)
+        return image
+    
     def controlnet_pose_inpaint(self, image_path: str):
         openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
 
@@ -62,8 +69,13 @@ class StableDiffusionControlNetInpaintPoseGenerator:
         scheduler: str,
         seed_generator: int,
     ):
-
-        image = self.controlnet_pose_inpaint(image_path=image_path)
+        normal_image = image_path["image"].convert("RGB").resize((512, 512))
+        mask_image = image_path["mask"].convert("RGB").resize((512, 512))
+        
+        normal_image = self.load_image(image_path=normal_image)
+        mask_image = self.load_image(image_path=mask_image)
+        
+        controlnet_image = self.controlnet_pose_inpaint(image_path=image_path)
 
         pipe = self.load_model(
             stable_model_path=stable_model_path,
@@ -79,7 +91,9 @@ class StableDiffusionControlNetInpaintPoseGenerator:
 
         output = pipe(
             prompt=prompt,
-            image=image,
+            image=normal_image,
+            mask_image=mask_image,
+            controlnet_image=controlnet_image,
             negative_prompt=negative_prompt,
             num_images_per_prompt=num_images_per_prompt,
             num_inference_steps=num_inference_step,
@@ -115,8 +129,8 @@ class StableDiffusionControlNetInpaintPoseGenerator:
                         with gr.Column():
                             controlnet_pose_inpaint_stable_model_id = (
                                 gr.Dropdown(
-                                    choices=stable_model_list,
-                                    value=stable_model_list[0],
+                                    choices=stable_inpiant_model_list,
+                                    value=stable_inpiant_model_list[0],
                                     label="Stable Model Id",
                                 )
                             )

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_scribble.py

@@ -2,11 +2,13 @@ import gradio as gr
 import numpy as np
 import torch
 from controlnet_aux import HEDdetector
-from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers import ControlNetModel
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
+from PIL import Image
 
 from diffusion_webui.utils.model_list import (
     controlnet_scribble_model_list,
-    stable_model_list,
+    stable_inpiant_model_list,
 )
 from diffusion_webui.utils.scheduler_list import (
     SCHEDULER_LIST,
@@ -25,7 +27,7 @@ class StableDiffusionControlNetInpaintScribbleGenerator:
                 controlnet_model_path, torch_dtype=torch.float16
             )
 
-            self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
                 pretrained_model_name_or_path=stable_model_path,
                 controlnet=controlnet,
                 safety_checker=None,
@@ -37,6 +39,11 @@ class StableDiffusionControlNetInpaintScribbleGenerator:
         self.pipe.enable_xformers_memory_efficient_attention()
 
         return self.pipe
+    
+    def load_image(self, image_path):
+        image = np.array(image_path)
+        image = Image.fromarray(image_path)
+        return image
 
     def controlnet_inpaint_scribble(self, image_path: str):
         hed = HEDdetector.from_pretrained("lllyasviel/ControlNet")
@@ -61,8 +68,13 @@ class StableDiffusionControlNetInpaintScribbleGenerator:
         scheduler: str,
         seed_generator: int,
     ):
-
-        image = self.controlnet_inpaint_scribble(image_path=image_path)
+        normal_image = image_path["image"].convert("RGB").resize((512, 512))
+        mask_image = image_path["mask"].convert("RGB").resize((512, 512))
+        
+        normal_image = self.load_image(image_path=normal_image)
+        mask_image = self.load_image(image_path=mask_image)
+        
+        controlnet_image = self.controlnet_inpaint_scribble(image_path=image_path)
 
         pipe = self.load_model(
             stable_model_path=stable_model_path,
@@ -78,7 +90,9 @@ class StableDiffusionControlNetInpaintScribbleGenerator:
 
         output = pipe(
             prompt=prompt,
-            image=image,
+            image=normal_image,
+            mask_image=mask_image,
+            controlnet_image=controlnet_image,
             negative_prompt=negative_prompt,
             num_images_per_prompt=num_images_per_prompt,
             num_inference_steps=num_inference_step,
@@ -114,8 +128,8 @@ class StableDiffusionControlNetInpaintScribbleGenerator:
                         with gr.Column():
                             controlnet_scribble_inpaint_stable_model_id = (
                                 gr.Dropdown(
-                                    choices=stable_model_list,
-                                    value=stable_model_list[0],
+                                    choices=stable_inpiant_model_list,
+                                    value=stable_inpiant_model_list[0],
                                     label="Stable Model Id",
                                 )
                             )

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_seg.py

@@ -7,7 +7,7 @@ from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
 
 from diffusion_webui.utils.model_list import (
     controlnet_seg_model_list,
-    stable_model_list,
+    stable_inpiant_model_list,
 )
 from diffusion_webui.utils.scheduler_list import (
     SCHEDULER_LIST,
@@ -200,6 +200,11 @@ class StableDiffusionControlNetInpaintSegGenerator:
         self.pipe.enable_xformers_memory_efficient_attention()
 
         return self.pipe
+    
+    def load_image(self, image_path):
+        image = np.array(image_path)
+        image = Image.fromarray(image_path)
+        return image
 
     def controlnet_seg_inpaint(self, image_path: str):
         image_processor = AutoImageProcessor.from_pretrained(
@@ -246,7 +251,13 @@ class StableDiffusionControlNetInpaintSegGenerator:
         seed_generator: int,
     ):
 
-        image = self.controlnet_seg_inpaint(image_path=image_path)
+        normal_image = image_path["image"].convert("RGB").resize((512, 512))
+        mask_image = image_path["mask"].convert("RGB").resize((512, 512))
+        
+        normal_image = self.load_image(image_path=normal_image)
+        mask_image = self.load_image(image_path=mask_image)
+        
+        controlnet_image = self.controlnet_seg_inpaint(image_path=image_path)
 
         pipe = self.load_model(
             stable_model_path=stable_model_path,
@@ -262,7 +273,9 @@ class StableDiffusionControlNetInpaintSegGenerator:
 
         output = pipe(
             prompt=prompt,
-            image=image,
+            image=normal_image,
+            mask_image=mask_image,
+            controlnet_image=controlnet_image,
             negative_prompt=negative_prompt,
             num_images_per_prompt=num_images_per_prompt,
             num_inference_steps=num_inference_step,
@@ -298,8 +311,8 @@ class StableDiffusionControlNetInpaintSegGenerator:
                         with gr.Column():
                             controlnet_seg_inpaint_stable_model_id = (
                                 gr.Dropdown(
-                                    choices=stable_model_list,
-                                    value=stable_model_list[0],
+                                    choices=stable_inpiant_model_list,
+                                    value=stable_inpiant_model_list[0],
                                     label="Stable Model Id",
                                 )
                             )