Upload 10 files

src/backend/annotators/canny_control.py

@@ -0,0 +1,15 @@
+import numpy as np
+from backend.annotators.control_interface import ControlInterface
+from cv2 import Canny
+from PIL import Image
+
+
+class CannyControl(ControlInterface):
+    def get_control_image(self, image: Image) -> Image:
+        low_threshold = 100
+        high_threshold = 200
+        image = np.array(image)
+        image = Canny(image, low_threshold, high_threshold)
+        image = image[:, :, None]
+        image = np.concatenate([image, image, image], axis=2)
+        return Image.fromarray(image)

src/backend/annotators/control_interface.py

@@ -0,0 +1,12 @@
+from abc import ABC, abstractmethod
+
+from PIL import Image
+
+
+class ControlInterface(ABC):
+    @abstractmethod
+    def get_control_image(
+        self,
+        image: Image,
+    ) -> Image:
+        pass

src/backend/annotators/depth_control.py

@@ -0,0 +1,15 @@
+import numpy as np
+from backend.annotators.control_interface import ControlInterface
+from PIL import Image
+from transformers import pipeline
+
+
+class DepthControl(ControlInterface):
+    def get_control_image(self, image: Image) -> Image:
+        depth_estimator = pipeline("depth-estimation")
+        image = depth_estimator(image)["depth"]
+        image = np.array(image)
+        image = image[:, :, None]
+        image = np.concatenate([image, image, image], axis=2)
+        image = Image.fromarray(image)
+        return image

src/backend/annotators/image_control_factory.py

@@ -0,0 +1,31 @@
+from backend.annotators.canny_control import CannyControl
+from backend.annotators.depth_control import DepthControl
+from backend.annotators.lineart_control import LineArtControl
+from backend.annotators.mlsd_control import MlsdControl
+from backend.annotators.normal_control import NormalControl
+from backend.annotators.pose_control import PoseControl
+from backend.annotators.shuffle_control import ShuffleControl
+from backend.annotators.softedge_control import SoftEdgeControl
+
+
+class ImageControlFactory:
+    def create_control(self, controlnet_type: str):
+        if controlnet_type == "Canny":
+            return CannyControl()
+        elif controlnet_type == "Pose":
+            return PoseControl()
+        elif controlnet_type == "MLSD":
+            return MlsdControl()
+        elif controlnet_type == "Depth":
+            return DepthControl()
+        elif controlnet_type == "LineArt":
+            return LineArtControl()
+        elif controlnet_type == "Shuffle":
+            return ShuffleControl()
+        elif controlnet_type == "NormalBAE":
+            return NormalControl()
+        elif controlnet_type == "SoftEdge":
+            return SoftEdgeControl()
+        else:
+            print("Error: Control type not implemented!")
+            raise Exception("Error: Control type not implemented!")

src/backend/annotators/lineart_control.py

@@ -0,0 +1,11 @@
+import numpy as np
+from backend.annotators.control_interface import ControlInterface
+from controlnet_aux import LineartDetector
+from PIL import Image
+
+
+class LineArtControl(ControlInterface):
+    def get_control_image(self, image: Image) -> Image:
+        processor = LineartDetector.from_pretrained("lllyasviel/Annotators")
+        control_image = processor(image)
+        return control_image

src/backend/annotators/mlsd_control.py

@@ -0,0 +1,10 @@
+from backend.annotators.control_interface import ControlInterface
+from controlnet_aux import MLSDdetector
+from PIL import Image
+
+
+class MlsdControl(ControlInterface):
+    def get_control_image(self, image: Image) -> Image:
+        mlsd = MLSDdetector.from_pretrained("lllyasviel/ControlNet")
+        image = mlsd(image)
+        return image

src/backend/annotators/normal_control.py

@@ -0,0 +1,10 @@
+from backend.annotators.control_interface import ControlInterface
+from controlnet_aux import NormalBaeDetector
+from PIL import Image
+
+
+class NormalControl(ControlInterface):
+    def get_control_image(self, image: Image) -> Image:
+        processor = NormalBaeDetector.from_pretrained("lllyasviel/Annotators")
+        control_image = processor(image)
+        return control_image

src/backend/annotators/pose_control.py

@@ -0,0 +1,10 @@
+from backend.annotators.control_interface import ControlInterface
+from controlnet_aux import OpenposeDetector
+from PIL import Image
+
+
+class PoseControl(ControlInterface):
+    def get_control_image(self, image: Image) -> Image:
+        openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
+        image = openpose(image)
+        return image

src/backend/annotators/shuffle_control.py

@@ -0,0 +1,10 @@
+from backend.annotators.control_interface import ControlInterface
+from controlnet_aux import ContentShuffleDetector
+from PIL import Image
+
+
+class ShuffleControl(ControlInterface):
+    def get_control_image(self, image: Image) -> Image:
+        shuffle_processor = ContentShuffleDetector()
+        image = shuffle_processor(image)
+        return image

src/backend/annotators/softedge_control.py

@@ -0,0 +1,10 @@
+from backend.annotators.control_interface import ControlInterface
+from controlnet_aux import PidiNetDetector
+from PIL import Image
+
+
+class SoftEdgeControl(ControlInterface):
+    def get_control_image(self, image: Image) -> Image:
+        processor = PidiNetDetector.from_pretrained("lllyasviel/Annotators")
+        control_image = processor(image)
+        return control_image