diff --git a/modules/control/proc/dsine/__init__.py b/modules/control/proc/dsine/__init__.py new file mode 100644 index 000000000..3704706f9 --- /dev/null +++ b/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