82 lines
2.4 KiB
Python
82 lines
2.4 KiB
Python
from itertools import chain
|
|
from math import ceil
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from coreml_suite.latents import chunk_batch
|
|
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
|
|
|
|
|
|
def chunk_control(cn, target_size):
|
|
if cn is None:
|
|
return [None] * target_size
|
|
|
|
num_chunks = ceil(cn["output"][0].shape[0] / target_size)
|
|
|
|
out = [{"output": [], "middle": []} for _ in range(num_chunks)]
|
|
|
|
for k, v in cn.items():
|
|
for i, x in enumerate(v):
|
|
chunks = chunk_batch(x, (target_size, *x.shape[1:]))
|
|
for j, chunk in enumerate(chunks):
|
|
out[j][k].append(chunk)
|
|
|
|
return out
|