Rearrange stuff

coreml_suite/controlnet.py

@@ -0,0 +1,62 @@
+from itertools import chain
+
+import numpy as np
+import torch
+
+from coreml_suite.logger import logger
+
+
+def expand_inputs(inputs):
+    expanded = inputs.copy()
+    for k, v in inputs.items():
+        if isinstance(v, np.ndarray):
+            expanded[k] = np.concatenate([v] * 2) if v.shape[0] == 1 else v
+        elif isinstance(v, torch.Tensor):
+            expanded[k] = torch.cat([v] * 2) if v.shape[0] == 1 else v
+        elif isinstance(v, list):
+            expanded[k] = v * 2 if len(v) == 1 else v
+        elif isinstance(v, dict):
+            expand_inputs(v)
+    return expanded
+
+
+def extract_residual_kwargs(model, control):
+    if "additional_residual_0" not in model.expected_inputs.keys():
+        return {}
+    if control is None:
+        return no_control(model)
+
+    residual_kwargs = {
+        "additional_residual_{}".format(i): r.cpu().numpy().astype(np.float16)
+        for i, r in enumerate(chain(control["output"], control["middle"]))
+    }
+    return residual_kwargs
+
+
+def no_control(model):
+    # Dirty hack to get the expected input shape when doing partial ControlNet
+    # 0.18215 is the latent scale factor (IDK, it kinda works)
+    # TODO: Find a better way to do this or tweak the values
+
+    logger.warning(
+        "No ControlNet input, despite the model supports it. "
+        "Using random noise as ControlNet residuals. "
+        "For better results, please use a ControlNet or a model "
+        "that does not support ControlNet."
+    )
+    residuals_names = [
+        name
+        for name in model.expected_inputs.keys()
+        if name.startswith("additional_residual")
+    ]
+    residual_kwargs = {
+        "additional_residual_{}".format(i): 0.18215
+        * torch.randn(
+            *model.expected_inputs["additional_residual_{}".format(i)]["shape"]
+        )
+        .cpu()
+        .numpy()
+        .astype(dtype=np.float16)
+        for i in range(len(residuals_names))
+    }
+    return residual_kwargs