Create block.py

block.py

@@ -0,0 +1,83 @@
+from diffusers.modular_pipelines import ModularPiplineBlocks, PipelineState, InputParam, OutputParam
+from diffusers.utils import load_image
+from PIL import Image
+import cv2
+import numpy as np
+
+class GetImageStep(ModularPiplineBlocks):
+
+    PROCESSOR_IDS = set([
+        "canny", "lineart_anime",
+    ])
+
+    def __init__(self):
+        super().__init__()
+        from controlnet_aux.processor import Processor
+        self.processor = Processor
+    
+    @staticmethod
+    def make_canny(image):
+        image = np.array(image)
+        image = cv2.Canny(image, 100, 200)
+        image = image[:, :, None]
+        image = np.concatenate([image, image, image], axis=2)
+        return Image.fromarray(image)
+    
+    def make_lineart_anime(self, image):
+        return self.processor("lineart_anime")(image)
+    
+    
+    def check_inputs(self, data) -> None:
+        """
+        Validates that `processor_id` is one of the supported processors.
+        Raises:
+            ValueError: if `processor_id` is not in PROCESSOR_IDS.
+        """
+
+        if data.image_url is None and data.image is None:
+            raise ValueError("Either `image_url` or `image` must be provided.")
+
+        if data.image_url is not None and data.image is not None:
+            raise ValueError("Only one of `image_url` or `image` must be provided.")
+        
+        if data.processor_id is not None and data.processor_id not in self.PROCESSOR_IDS:
+            raise ValueError(
+                f"Processor id '{data.processor_id}' not found. "
+                f"Please use one of the following: {self.PROCESSOR_IDS}"
+            )
+    
+    @property
+    def inputs(self):
+        return [
+            InputParam("image", type_hint=Image.Image),
+            InputParam("image_url", type_hint=str, description="The url of the image to load"),
+            InputParam("size", description="The size of the image"),
+            InputParam("processor_id", type_hint=str, description="The id of the processor to use for controlnet")
+        ]
+
+    @property
+    def intermediate_outputs(self):
+        return [
+          OutputParam("image", type_hint=Image.Image),
+        ]
+
+    def __call__(self, pipeline, state: PipelineState):
+
+        data = self.get_block_state(state)
+        self.check_inputs(data)
+
+        if data.image is None:
+            data.image = load_image(data.image_url).convert("RGB")
+        
+        if data.size is not None:
+            data.image = data.image.resize(data.size)
+
+        if data.processor_id is not None:
+            if data.processor_id == "canny":
+                data.image = self.make_canny(data.image)
+            elif data.processor_id == "lineart_anime":
+                data.image = self.make_lineart_anime(data.image)
+        
+        self.set_block_state(state, data)
+        
+        return pipeline, state