support video-to-video-translation

diffsynth/controlnets/processors.py

@@ -0,0 +1,50 @@
+from typing_extensions import Literal, TypeAlias
+import warnings
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
+    from controlnet_aux.processor import (
+        CannyDetector, MidasDetector, HEDdetector, LineartDetector, LineartAnimeDetector, OpenposeDetector
+    )
+
+
+Processor_id: TypeAlias = Literal[
+    "canny", "depth", "softedge", "lineart", "lineart_anime", "openpose", "tile"
+]
+
+class Annotator:
+    def __init__(self, processor_id: Processor_id, model_path="models/Annotators", detect_resolution=512):
+        if processor_id == "canny":
+            self.processor = CannyDetector()
+        elif processor_id == "depth":
+            self.processor = MidasDetector.from_pretrained(model_path)
+        elif processor_id == "softedge":
+            self.processor = HEDdetector.from_pretrained(model_path)
+        elif processor_id == "lineart":
+            self.processor = LineartDetector.from_pretrained(model_path)
+        elif processor_id == "lineart_anime":
+            self.processor = LineartAnimeDetector.from_pretrained(model_path)
+        elif processor_id == "openpose":
+            self.processor = OpenposeDetector.from_pretrained(model_path)
+        elif processor_id == "tile":
+            self.processor = None
+        else:
+            raise ValueError(f"Unsupported processor_id: {processor_id}")
+        
+        self.processor_id = processor_id
+        self.detect_resolution = detect_resolution
+
+    def __call__(self, image):
+        width, height = image.size
+        if self.processor_id == "openpose":
+            kwargs = {
+                "include_body": True,
+                "include_hand": True,
+                "include_face": True
+            }
+        else:
+            kwargs = {}
+        if self.processor is not None:
+            image = self.processor(image, detect_resolution=self.detect_resolution, image_resolution=min(width, height), **kwargs)
+        image = image.resize((width, height))
+        return image
+