merge: modules/control/proc/dsine/__init__.py

modules/control/proc/dsine/__init__.py

@@ -0,0 +1,44 @@
+import os
+import torch
+import numpy as np
+from PIL import Image
+from modules import devices
+from modules.shared import opts
+
+
+class DSINEDetector:
+    def __init__(self, predictor):
+        self.predictor = predictor
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_or_path="hugoycj/DSINE-hub", cache_dir=None, local_files_only=False):
+        from installer import install
+        install('geffnet', quiet=True)
+        hub_dir = os.path.join(cache_dir, 'torch_hub') if cache_dir else None
+        old_hub_dir = torch.hub.get_dir()
+        if hub_dir:
+            os.makedirs(hub_dir, exist_ok=True)
+            torch.hub.set_dir(hub_dir)
+        try:
+            predictor = torch.hub.load(pretrained_model_or_path, "DSINE", trust_repo=True, source="github")
+        finally:
+            torch.hub.set_dir(old_hub_dir)
+        # Override hardcoded cuda device with project device
+        predictor.device = devices.device
+        predictor.model = predictor.model.to(devices.device).eval()
+        return cls(predictor)
+
+    def __call__(self, image, output_type="pil", **kwargs):
+        self.predictor.device = devices.device
+        self.predictor.model.to(devices.device)
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        with devices.inference_context():
+            normals = self.predictor.infer_pil(image)
+        if opts.control_move_processor:
+            self.predictor.model.to("cpu")
+        normals = normals[0].permute(1, 2, 0).cpu().numpy()
+        normals = ((normals + 1.0) * 0.5 * 255.0).clip(0, 255).astype(np.uint8)
+        if output_type == "pil":
+            normals = Image.fromarray(normals)
+        return normals