Added controlnet seg support,added schedulers - KDPM2DiscreteScheduler,HeunDiscreteScheduler,KDPM2AncestralDiscreteScheduler,DPMSolverSinglestepScheduler

Author: rupeshs

Readme.md

@@ -22,6 +22,7 @@ DiffusionMagic focused on the following areas:
    - Pose
    - Depth
    - Scribble
+   - Segmentation
 - Pytorch 2.0 support
 - Supports all stable diffusion Hugging Face models 
 - Supports Stable diffusion v1 and v2 models, derived models
@@ -109,6 +110,10 @@ E.g `https://huggingface.co/dreamlike-art/dreamlike-diffusion-1.0`
 Here model id is `dreamlike-art/dreamlike-diffusion-1.0`
 Or we can clone the model use the local folder path as model id.
 - Adding locally copied model path to configs/stable_diffusion_models.txt file
+## Linting (Development)
+Run the following commands from src folder
+`mypy --ignore-missing-imports --explicit-package-bases .`
+`flake8 --max-line-length=100 .`
 ## Contribute
 Contributions are welcomed.
 

configs/stable_diffusion_models.txt

@@ -15,4 +15,5 @@ lllyasviel/sd-controlnet-normal
 lllyasviel/sd-controlnet-hed
 lllyasviel/sd-controlnet-openpose
 lllyasviel/sd-controlnet-depth
-lllyasviel/sd-controlnet-scribble
+lllyasviel/sd-controlnet-scribble
+lllyasviel/sd-controlnet-seg

src/backend/controlnet/ControlContext.py

@@ -1,23 +1,19 @@
 from time import time
-from typing import Any
-import numpy as np
-from cv2 import Canny, bitwise_not
+
 import torch
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
-from PIL import Image, ImageOps
+from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
+from PIL import ImageOps
 
 from backend.computing import Computing
+from backend.controlnet.controls.image_control_factory import ImageControlFactory
+from backend.image_ops import resize_pil_image
 from backend.stablediffusion.models.scheduler_types import SchedulerType
-from backend.stablediffusion.models.setting import (
-    StableDiffusionControlnetSetting,
-)
+from backend.stablediffusion.models.setting import StableDiffusionControlnetSetting
 from backend.stablediffusion.scheduler_mixin import SamplerMixin
-from backend.image_ops import resize_pil_image
 from backend.stablediffusion.stable_diffusion_types import (
-    get_diffusion_type,
     StableDiffusionType,
+    get_diffusion_type,
 )
-from backend.controlnet.controls.image_control_factory import ImageControlFactory
 
 
 class ControlnetContext(SamplerMixin):
@@ -134,16 +130,6 @@ def _load_model(self):
         else:
             self._load_full_precision_model()
 
-    def get_canny_image(self, image: Image) -> Any:
-        low_threshold = 100
-        high_threshold = 200
-        image = np.array(image)
-        image = Canny(image, low_threshold, high_threshold)
-        image_inv = bitwise_not(image)
-        image = image[:, :, None]
-        image = np.concatenate([image, image, image], axis=2)
-        return Image.fromarray(image), Image.fromarray(image_inv)
-
     def _enable_slicing(self, setting: StableDiffusionControlnetSetting):
         if setting.attention_slicing:
             self.controlnet_pipeline.enable_attention_slicing()

src/backend/controlnet/controls/image_control_factory.py

@@ -1,13 +1,12 @@
-from backend.stablediffusion.stable_diffusion_types import (
-    StableDiffusionType,
-)
 from backend.controlnet.controls.canny_control import CannyControl
+from backend.controlnet.controls.depth_control import DepthControl
+from backend.controlnet.controls.hed_control import HedControl
 from backend.controlnet.controls.line_control import LineControl
 from backend.controlnet.controls.normal_control import NormalControl
-from backend.controlnet.controls.hed_control import HedControl
 from backend.controlnet.controls.pose_control import PoseControl
-from backend.controlnet.controls.depth_control import DepthControl
 from backend.controlnet.controls.scribble_control import ScribbleControl
+from backend.controlnet.controls.seg_control import SegControl
+from backend.stablediffusion.stable_diffusion_types import StableDiffusionType
 
 
 class ImageControlFactory:
@@ -26,6 +25,8 @@ def create_control(self, controlnet_type: StableDiffusionType):
             return DepthControl()
         elif controlnet_type == StableDiffusionType.controlnet_scribble:
             return ScribbleControl()
+        elif controlnet_type == StableDiffusionType.controlnet_seg:
+            return SegControl()
         else:
             print("Error: Control type not implemented!")
             raise Exception("Error: Control type not implemented!")

src/backend/controlnet/controls/seg_control.py

@@ -0,0 +1,186 @@
+import numpy as np
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
+
+from backend.controlnet.controls.control_interface import ControlInterface
+
+palette = np.asarray(
+    [
+        [0, 0, 0],
+        [120, 120, 120],
+        [180, 120, 120],
+        [6, 230, 230],
+        [80, 50, 50],
+        [4, 200, 3],
+        [120, 120, 80],
+        [140, 140, 140],
+        [204, 5, 255],
+        [230, 230, 230],
+        [4, 250, 7],
+        [224, 5, 255],
+        [235, 255, 7],
+        [150, 5, 61],
+        [120, 120, 70],
+        [8, 255, 51],
+        [255, 6, 82],
+        [143, 255, 140],
+        [204, 255, 4],
+        [255, 51, 7],
+        [204, 70, 3],
+        [0, 102, 200],
+        [61, 230, 250],
+        [255, 6, 51],
+        [11, 102, 255],
+        [255, 7, 71],
+        [255, 9, 224],
+        [9, 7, 230],
+        [220, 220, 220],
+        [255, 9, 92],
+        [112, 9, 255],
+        [8, 255, 214],
+        [7, 255, 224],
+        [255, 184, 6],
+        [10, 255, 71],
+        [255, 41, 10],
+        [7, 255, 255],
+        [224, 255, 8],
+        [102, 8, 255],
+        [255, 61, 6],
+        [255, 194, 7],
+        [255, 122, 8],
+        [0, 255, 20],
+        [255, 8, 41],
+        [255, 5, 153],
+        [6, 51, 255],
+        [235, 12, 255],
+        [160, 150, 20],
+        [0, 163, 255],
+        [140, 140, 140],
+        [250, 10, 15],
+        [20, 255, 0],
+        [31, 255, 0],
+        [255, 31, 0],
+        [255, 224, 0],
+        [153, 255, 0],
+        [0, 0, 255],
+        [255, 71, 0],
+        [0, 235, 255],
+        [0, 173, 255],
+        [31, 0, 255],
+        [11, 200, 200],
+        [255, 82, 0],
+        [0, 255, 245],
+        [0, 61, 255],
+        [0, 255, 112],
+        [0, 255, 133],
+        [255, 0, 0],
+        [255, 163, 0],
+        [255, 102, 0],
+        [194, 255, 0],
+        [0, 143, 255],
+        [51, 255, 0],
+        [0, 82, 255],
+        [0, 255, 41],
+        [0, 255, 173],
+        [10, 0, 255],
+        [173, 255, 0],
+        [0, 255, 153],
+        [255, 92, 0],
+        [255, 0, 255],
+        [255, 0, 245],
+        [255, 0, 102],
+        [255, 173, 0],
+        [255, 0, 20],
+        [255, 184, 184],
+        [0, 31, 255],
+        [0, 255, 61],
+        [0, 71, 255],
+        [255, 0, 204],
+        [0, 255, 194],
+        [0, 255, 82],
+        [0, 10, 255],
+        [0, 112, 255],
+        [51, 0, 255],
+        [0, 194, 255],
+        [0, 122, 255],
+        [0, 255, 163],
+        [255, 153, 0],
+        [0, 255, 10],
+        [255, 112, 0],
+        [143, 255, 0],
+        [82, 0, 255],
+        [163, 255, 0],
+        [255, 235, 0],
+        [8, 184, 170],
+        [133, 0, 255],
+        [0, 255, 92],
+        [184, 0, 255],
+        [255, 0, 31],
+        [0, 184, 255],
+        [0, 214, 255],
+        [255, 0, 112],
+        [92, 255, 0],
+        [0, 224, 255],
+        [112, 224, 255],
+        [70, 184, 160],
+        [163, 0, 255],
+        [153, 0, 255],
+        [71, 255, 0],
+        [255, 0, 163],
+        [255, 204, 0],
+        [255, 0, 143],
+        [0, 255, 235],
+        [133, 255, 0],
+        [255, 0, 235],
+        [245, 0, 255],
+        [255, 0, 122],
+        [255, 245, 0],
+        [10, 190, 212],
+        [214, 255, 0],
+        [0, 204, 255],
+        [20, 0, 255],
+        [255, 255, 0],
+        [0, 153, 255],
+        [0, 41, 255],
+        [0, 255, 204],
+        [41, 0, 255],
+        [41, 255, 0],
+        [173, 0, 255],
+        [0, 245, 255],
+        [71, 0, 255],
+        [122, 0, 255],
+        [0, 255, 184],
+        [0, 92, 255],
+        [184, 255, 0],
+        [0, 133, 255],
+        [255, 214, 0],
+        [25, 194, 194],
+        [102, 255, 0],
+        [92, 0, 255],
+    ]
+)
+
+
+class SegControl(ControlInterface):
+    def get_control_image(self, image: Image) -> Image:
+        image_processor = AutoImageProcessor.from_pretrained(
+            "openmmlab/upernet-convnext-small"
+        )
+        image_segmentor = UperNetForSemanticSegmentation.from_pretrained(
+            "openmmlab/upernet-convnext-small"
+        )
+        pixel_values = image_processor(image, return_tensors="pt").pixel_values
+        with torch.no_grad():
+            outputs = image_segmentor(pixel_values)
+        seg = image_processor.post_process_semantic_segmentation(
+            outputs, target_sizes=[image.size[::-1]]
+        )[0]
+
+        color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
+
+        for label, color in enumerate(palette):
+            color_seg[seg == label, :] = color
+        color_seg = color_seg.astype(np.uint8)
+        image = Image.fromarray(color_seg)
+        return image

src/backend/stablediffusion/models/scheduler_types.py

@@ -9,7 +9,11 @@
     EulerDiscreteScheduler,
     LMSDiscreteScheduler,
     PNDMScheduler,
-    UniPCMultistepScheduler
+    UniPCMultistepScheduler,
+    KDPM2DiscreteScheduler,
+    HeunDiscreteScheduler,
+    KDPM2AncestralDiscreteScheduler,
+    DPMSolverSinglestepScheduler,
 )
 
 
@@ -25,6 +29,10 @@ class SchedulerType(Enum):
     PNDMScheduler = "PNDM"
     DEISScheduler = "DEISMultistep"
     UniPCMultistepScheduler = "UniPCMultistep"
+    KDPM2DiscreteScheduler = "KDPM2DiscreteScheduler"
+    HeunDiscreteScheduler = "HeunDiscreteScheduler"
+    KDPM2AncestralDiscreteScheduler = "KDPM2AncestralDiscreteScheduler"
+    DPMSolverSinglestepScheduler = "DPMSolverSinglestepScheduler"
 
 
 Scheduler = Union[
@@ -36,6 +44,10 @@ class SchedulerType(Enum):
     LMSDiscreteScheduler,
     PNDMScheduler,
     UniPCMultistepScheduler,
+    KDPM2DiscreteScheduler,
+    HeunDiscreteScheduler,
+    KDPM2AncestralDiscreteScheduler,
+    DPMSolverSinglestepScheduler,
 ]
 
 

src/backend/stablediffusion/scheduler_factory.py

@@ -7,6 +7,10 @@
     LMSDiscreteScheduler,
     PNDMScheduler,
     UniPCMultistepScheduler,
+    KDPM2DiscreteScheduler,
+    HeunDiscreteScheduler,
+    KDPM2AncestralDiscreteScheduler,
+    DPMSolverSinglestepScheduler,
 )
 
 from backend.stablediffusion.models.scheduler_types import SchedulerType, Scheduler
@@ -58,6 +62,26 @@ def get_scheduler(
                 repo_id,
                 subfolder="scheduler",
             )
+        elif scheduler_type == SchedulerType.KDPM2DiscreteScheduler.value:
+            return KDPM2DiscreteScheduler.from_pretrained(
+                repo_id,
+                subfolder="scheduler",
+            )
+        elif scheduler_type == SchedulerType.HeunDiscreteScheduler.value:
+            return HeunDiscreteScheduler.from_pretrained(
+                repo_id,
+                subfolder="scheduler",
+            )
+        elif scheduler_type == SchedulerType.KDPM2AncestralDiscreteScheduler.value:
+            return KDPM2AncestralDiscreteScheduler.from_pretrained(
+                repo_id,
+                subfolder="scheduler",
+            )
+        elif scheduler_type == SchedulerType.DPMSolverSinglestepScheduler.value:
+            return DPMSolverSinglestepScheduler.from_pretrained(
+                repo_id,
+                subfolder="scheduler",
+            )
         else:
             print(f"Scheduler {scheduler_type} not found")
             return None

src/backend/stablediffusion/stable_diffusion_types.py

@@ -15,6 +15,7 @@ class StableDiffusionType(str, Enum):
     controlnet_pose = "controlnet_pose"
     controlnet_depth = "controlnet_depth"
     controlnet_scribble = "controlnet_scribble"
+    controlnet_seg = "controlnet_seg"
 
 
 def get_diffusion_type(
@@ -39,6 +40,8 @@ def get_diffusion_type(
         stable_diffusion_type = StableDiffusionType.controlnet_depth
     elif "depth" in model_id:
         stable_diffusion_type = StableDiffusionType.depth2img
-    elif "scribble" in model_id:
+    elif "controlnet-scribble" in model_id:
         stable_diffusion_type = StableDiffusionType.controlnet_scribble
+    elif "controlnet-seg" in model_id:
+        stable_diffusion_type = StableDiffusionType.controlnet_seg
     return stable_diffusion_type

src/constants.py

@@ -1,4 +1,4 @@
-VERSION = "2.0.0-alpha.0"
+VERSION = "2.0.0-beta.0"
 STABLE_DIFFUSION_MODELS_FILE = "stable_diffusion_models.txt"
 APP_SETTINGS_FILE = "settings.yaml"
 CONFIG_DIRECTORY = "configs"