aszc
d0cca3c3f4
* fix(conversion): load .mlpackage instead of unloadable .mlmodelc The native ct.models.MLModel runtime added in the diffusers conversion path cannot load compiled .mlmodelc directories (no Manifest.json), so the converter's compiled output failed at load with "A valid manifest does not exist". Both converters now return the .mlpackage directly and the loader lists only .mlpackage. Removes the now-dead coremlcompiler wrappers. * chore(release): bump version to 2.0.1
71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
import os
|
|
|
|
from coremltools import ComputeUnit
|
|
|
|
from coreml_suite import COREML_NODE
|
|
from coreml_suite.coreml_model import CoreMLModel
|
|
|
|
|
|
class COREML_CONVERT_LCM(COREML_NODE):
|
|
"""Converts a LCM model to Core ML."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"height": ("INT", {"default": 512, "min": 512, "max": 768, "step": 8}),
|
|
"width": ("INT", {"default": 512, "min": 512, "max": 768, "step": 8}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
|
|
"compute_unit": (
|
|
[
|
|
ComputeUnit.CPU_AND_NE.name,
|
|
ComputeUnit.CPU_AND_GPU.name,
|
|
ComputeUnit.ALL.name,
|
|
ComputeUnit.CPU_ONLY.name,
|
|
],
|
|
),
|
|
"controlnet_support": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("COREML_UNET",)
|
|
RETURN_NAMES = ("coreml_model",)
|
|
FUNCTION = "convert"
|
|
|
|
def convert(self, height, width, batch_size, compute_unit, controlnet_support):
|
|
"""Converts a LCM model to Core ML.
|
|
|
|
Args:
|
|
height (int): Height of the target image.
|
|
width (int): Width of the target image.
|
|
batch_size (int): Batch size.
|
|
compute_unit (str): Compute unit to use when loading the model.
|
|
|
|
Returns:
|
|
coreml_model: The converted Core ML model.
|
|
|
|
The converted model is also saved to "models/unet" directory and
|
|
can be loaded with the "LCMCoreMLLoaderUNet" node.
|
|
"""
|
|
from coreml_suite.lcm import converter as lcm_converter
|
|
|
|
h = height
|
|
w = width
|
|
sample_size = (h // 8, w // 8)
|
|
batch_size = batch_size
|
|
cn_support_str = "_cn" if controlnet_support else ""
|
|
|
|
out_name = f"{lcm_converter.MODEL_NAME}_{batch_size}x{w}x{h}{cn_support_str}"
|
|
|
|
out_path = lcm_converter.get_out_path("unet", f"{out_name}")
|
|
|
|
if not os.path.exists(out_path):
|
|
lcm_converter.convert(
|
|
out_path=out_path,
|
|
sample_size=sample_size,
|
|
batch_size=batch_size,
|
|
controlnet_support=controlnet_support,
|
|
)
|
|
|
|
return (CoreMLModel(out_path, compute_unit),)
|