99 lines
3.3 KiB
Python
99 lines
3.3 KiB
Python
import numpy as np
|
|
import torch
|
|
|
|
from coreml_suite.controlnet import extract_residual_kwargs, chunk_control
|
|
from coreml_suite.latents import chunk_batch, merge_chunks
|
|
|
|
|
|
class CoreMLModelWrapper:
|
|
def __init__(self, coreml_model):
|
|
self.coreml_model = coreml_model
|
|
self.dtype = torch.float16
|
|
|
|
def __call__(self, x, t, context, control, transformer_options, **kwargs):
|
|
inputs = CoreMLInputs(x, t, context, control, **kwargs)
|
|
input_list = inputs.chunks(self.expected_inputs)
|
|
|
|
chunked_out = [
|
|
self.get_torch_outputs(
|
|
self.coreml_model(**input_kwargs.coreml_kwargs(self.expected_inputs)),
|
|
x.device,
|
|
)
|
|
for input_kwargs in input_list
|
|
]
|
|
merged_out = merge_chunks(chunked_out, x.shape)
|
|
|
|
return merged_out
|
|
|
|
@staticmethod
|
|
def get_torch_outputs(model_output, device):
|
|
return torch.from_numpy(model_output["noise_pred"]).to(device)
|
|
|
|
@property
|
|
def expected_inputs(self):
|
|
return self.coreml_model.expected_inputs
|
|
|
|
|
|
class CoreMLModelWrapperLCM(CoreMLModelWrapper):
|
|
def __init__(self, coreml_model):
|
|
super().__init__(coreml_model)
|
|
self.config = None
|
|
|
|
|
|
class CoreMLInputs:
|
|
def __init__(self, x, t, context, control, **kwargs):
|
|
self.x = x
|
|
self.t = t
|
|
self.context = context
|
|
self.control = control
|
|
self.ts_cond = kwargs.get("timestep_cond")
|
|
|
|
def coreml_kwargs(self, expected_inputs):
|
|
sample = self.x.cpu().numpy().astype(np.float16)
|
|
|
|
context = self.context.cpu().numpy().astype(np.float16)
|
|
context = context.transpose(0, 2, 1)[:, :, None, :]
|
|
|
|
t = self.t.cpu().numpy().astype(np.float16)
|
|
|
|
model_input_kwargs = {
|
|
"sample": sample,
|
|
"encoder_hidden_states": context,
|
|
"timestep": t,
|
|
}
|
|
residual_kwargs = extract_residual_kwargs(expected_inputs, self.control)
|
|
model_input_kwargs |= residual_kwargs
|
|
|
|
if self.ts_cond is not None:
|
|
model_input_kwargs["timestep_cond"] = (
|
|
self.ts_cond.cpu().numpy().astype(np.float16)
|
|
)
|
|
|
|
return model_input_kwargs
|
|
|
|
def chunks(self, expected_inputs):
|
|
sample_shape = expected_inputs["sample"]["shape"]
|
|
timestep_shape = expected_inputs["timestep"]["shape"]
|
|
hidden_shape = expected_inputs["encoder_hidden_states"]["shape"]
|
|
context_shape = (hidden_shape[0], hidden_shape[3], hidden_shape[1])
|
|
|
|
chunked_x = chunk_batch(self.x, sample_shape)
|
|
ts = list(torch.full((len(chunked_x), timestep_shape[0]), self.t[0]))
|
|
chunked_context = chunk_batch(self.context, context_shape)
|
|
|
|
chunked_control = [None] * len(chunked_x)
|
|
if self.control is not None:
|
|
chunked_control = chunk_control(self.control, sample_shape[0])
|
|
|
|
chunked_ts_cond = [None] * len(chunked_x)
|
|
if self.ts_cond is not None:
|
|
ts_cond_shape = expected_inputs["timestep_cond"]["shape"]
|
|
chunked_ts_cond = chunk_batch(self.ts_cond, ts_cond_shape)
|
|
|
|
return [
|
|
CoreMLInputs(x, t, context, control, timestep_cond=ts_cond)
|
|
for x, t, context, control, ts_cond in zip(
|
|
chunked_x, ts, chunked_context, chunked_control, chunked_ts_cond
|
|
)
|
|
]
|