63 lines
1.9 KiB
Python
63 lines
1.9 KiB
Python
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
|