aszc
02b6e8ece3
Modernizes ComfyUI-CoreMLSuite onto Python 3.12 / torch 2.7 / coremltools 9 with a characterization-test safety net. The default conversion path is unchanged; existing saved workflows produce identical output. - Toolchain bump (Python 3.12, torch 2.7, coremltools 9, numpy <2) with the blocking upstream pins overridden. - Framework-free logic moved into coreml_suite/core/ (no comfy/coremltools imports); old module paths re-export from there. - Opt-in quantize_nbits dropdown (none|8|6|4) for k-means weight palettization; default none is byte-for-byte identical to before. - Tiered CI: Tier 0 (Linux unit), Tier 1 (macOS-ARM smoke), Tier 2 (self-hosted Apple Silicon golden-image check on the ANE).
68 lines
2.0 KiB
Python
68 lines
2.0 KiB
Python
"""Pure helpers around the ControlNet residual inputs of the Core ML UNet.
|
|
|
|
Re-exported by coreml_suite.controlnet. Characterization tests cover
|
|
shapes, dtype (fp16), and zero-fill fallback.
|
|
"""
|
|
from itertools import chain
|
|
from math import ceil
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from coreml_suite.core.latents import chunk_batch
|
|
|
|
|
|
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(expected_inputs, control):
|
|
if "additional_residual_0" not in expected_inputs.keys():
|
|
return {}
|
|
if control is None:
|
|
return no_control(expected_inputs)
|
|
|
|
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(expected_inputs):
|
|
shapes_dict = {
|
|
k: v["shape"] for k, v in expected_inputs.items() if k.startswith("additional")
|
|
}
|
|
residual_kwargs = {
|
|
k: torch.zeros(*shape).cpu().numpy().astype(dtype=np.float16)
|
|
for k, shape in shapes_dict.items()
|
|
}
|
|
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
|