coreml-diffusion 0.1.3__tar.gz → 0.1.4__tar.gz

coreml_diffusion-0.1.4/.release-please-manifest.json

@@ -0,0 +1,3 @@
+{
+  ".": "0.1.4"
+}

{coreml_diffusion-0.1.3 → coreml_diffusion-0.1.4}/CHANGELOG.md

@@ -1,5 +1,13 @@
 # Changelog
 
+## [0.1.4](https://github.com/aszc-dev/coreml-diffusion/compare/v0.1.3...v0.1.4) (2026-06-13)
+
+
+### 🐛 Bug Fixes
+
+* **convert:** generic LCM conversion for arbitrary checkpoints ([bedeb49](https://github.com/aszc-dev/coreml-diffusion/commit/bedeb49a84a33b4530c6c7d4ed4343a914a400e2))
+* **inference:** add output_dtype to CoreMLTextEncoder ([eb6d3b5](https://github.com/aszc-dev/coreml-diffusion/commit/eb6d3b5b08b6ff34a4eb8683de3d1223fb447186))
+
 ## [0.1.3](https://github.com/aszc-dev/coreml-diffusion/compare/v0.1.2...v0.1.3) (2026-06-04)
 
 

{coreml_diffusion-0.1.3 → coreml_diffusion-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coreml-diffusion
-Version: 0.1.3
+Version: 0.1.4
 Summary: Convert diffusion-model checkpoints (SD1.5/SDXL) to Core ML for Apple Neural Engine — framework-free, ComfyUI-independent.
 Project-URL: Homepage, https://github.com/aszc-dev/coreml-diffusion
 Project-URL: Repository, https://github.com/aszc-dev/coreml-diffusion

coreml_diffusion-0.1.4/coreml_diffusion/conversion/state_dict.py

@@ -0,0 +1,20 @@
+"""State-dict layout predicates — framework-free (no coremltools/diffusers).
+
+Single-file checkpoints come in two layouts: original LDM (UNet keys under
+``model.diffusion_model.``) and diffusers-native UNet-only dumps (block keys at
+the top level — e.g. ``LCM_Dreamshaper_v7_4k.safetensors``, the canonical
+full-distill LCM artifact). diffusers' ``from_single_file`` only understands
+the former and raises ``SingleFileComponentError`` on the latter;
+``convert.load_unet`` routes on this predicate.
+"""
+
+DIFFUSERS_UNET_KEY_PREFIXES = ("down_blocks.", "up_blocks.", "mid_block.")
+LDM_UNET_KEY_PREFIX = "model.diffusion_model."
+
+
+def is_diffusers_unet_layout(keys) -> bool:
+    """True when ``keys`` form a diffusers-format UNet-only state dict."""
+    keys = list(keys)
+    has_diffusers_blocks = any(k.startswith(DIFFUSERS_UNET_KEY_PREFIXES) for k in keys)
+    has_ldm_prefix = any(k.startswith(LDM_UNET_KEY_PREFIX) for k in keys)
+    return has_diffusers_blocks and not has_ldm_prefix

{coreml_diffusion-0.1.3 → coreml_diffusion-0.1.4}/coreml_diffusion/convert.py

@@ -22,6 +22,7 @@ from coreml_diffusion.attention import ATTENTION_IMPLEMENTATIONS
 from coreml_diffusion.component import Component
 from coreml_diffusion.conversion.attention import apply_attention_implementation
 from coreml_diffusion.conversion.shapes import conv2d_output_shape
+from coreml_diffusion.conversion.state_dict import is_diffusers_unet_layout
 from coreml_diffusion.conversion.text_encoder import (
     CoreMLTextEncoderWrapper,
     static_causal_mask,
@@ -139,9 +140,17 @@ def get_sample_input(
     return sample_unet_inputs
 
 
-def lcm_inputs(sample_unet_inputs):
+def lcm_inputs(sample_unet_inputs, ref_unet=None):
+    """Build the LCM guidance-embedding (``timestep_cond``) trace input.
+
+    The embedding dim comes from ``ref_unet.config.time_cond_proj_dim``. The
+    256 fallback preserves the legacy ref-unet-less call shape (the Suite's LCM
+    converter calls this with one argument); 256 is correct for every known
+    SD1.5 full-distill LCM.
+    """
     batch_size = sample_unet_inputs["sample"].shape[0]
-    return {"timestep_cond": torch.randn(batch_size, 256).to(torch.float32)}
+    dim = 256 if ref_unet is None else ref_unet.config.time_cond_proj_dim
+    return {"timestep_cond": torch.randn(batch_size, dim).to(torch.float32)}
 
 
 def sdxl_inputs(sample_unet_inputs, ref_unet, model_version):
@@ -257,7 +266,15 @@ def convert_unet(
     )
 
     if model_version == ModelVersion.LCM:
-        sample_inputs |= lcm_inputs(sample_inputs)
+        if ref_unet.config.time_cond_proj_dim is None:
+            raise ValueError(
+                "model_version=LCM requires a UNet with a guidance embedding "
+                "(config.time_cond_proj_dim), but this checkpoint has none — "
+                "it is an LCM-LoRA merge with plain SD1.5 architecture. "
+                "Convert it with model_version=SD15 and use an LCM scheduler "
+                "at sampling time."
+            )
+        sample_inputs |= lcm_inputs(sample_inputs, ref_unet)
 
     if model_version in {ModelVersion.SDXL, ModelVersion.SDXL_REFINER}:
         sample_inputs |= sdxl_inputs(sample_inputs, ref_unet, model_version)
@@ -529,12 +546,63 @@ def convert(
 
 
 def load_unet(ckpt_path, config_path):
+    """Load the UNet from a single-file checkpoint, routing on state-dict layout.
+
+    LDM-layout files (the Civitai norm) go through ``from_single_file``.
+    Diffusers-layout UNet-only dumps (the canonical full-distill LCM artifacts,
+    e.g. ``LCM_Dreamshaper_v7_4k.safetensors``) are rejected by
+    ``from_single_file`` outright, so they get a direct state-dict load.
+    """
+    keys = _safetensors_keys(ckpt_path)
+    if keys is not None and is_diffusers_unet_layout(keys):
+        return load_unet_from_diffusers_state_dict(ckpt_path)
     return UNet2DConditionModel.from_single_file(
         ckpt_path,
         original_config=config_path,
     )
 
 
+def _safetensors_keys(ckpt_path):
+    """The file's key list when it is safetensors, else None.
+
+    Probes by content, not filename — a resolved checkpoint path may point at
+    an extension-less blob (e.g. inside the Hugging Face cache).
+    """
+    from safetensors import SafetensorError, safe_open
+
+    try:
+        with safe_open(ckpt_path, framework="pt") as f:
+            return list(f.keys())
+    except SafetensorError:
+        return None
+
+
+def load_unet_from_diffusers_state_dict(ckpt_path, **config_overrides):
+    """Load a diffusers-layout UNet-only safetensors dump (SD1.5-class).
+
+    SD1.5 architecture is assumed (``UNet2DConditionModel`` defaults); the
+    cross-attention and guidance-embedding dims are read from the weights, so
+    both full-distill LCM dumps (``time_embedding.cond_proj`` present) and
+    plain SD1.5 UNet dumps load correctly. A non-SD1.5-class dump fails the
+    strict ``load_state_dict`` with an explicit shape/key error.
+    ``config_overrides`` exists for tests exercising miniature architectures.
+    """
+    from safetensors.torch import load_file
+
+    state_dict = load_file(ckpt_path)
+    cond_proj = state_dict.get("time_embedding.cond_proj.weight")
+    config_kwargs = {
+        "sample_size": 64,
+        "cross_attention_dim": state_dict[
+            "down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_k.weight"
+        ].shape[1],
+        "time_cond_proj_dim": None if cond_proj is None else cond_proj.shape[1],
+    } | config_overrides
+    unet = UNet2DConditionModel(**config_kwargs)
+    unet.load_state_dict(state_dict)
+    return unet
+
+
 def load_vae(ckpt_path):
     """Load ``AutoencoderKL`` from a single-file checkpoint."""
     from diffusers import AutoencoderKL

{coreml_diffusion-0.1.3 → coreml_diffusion-0.1.4}/coreml_diffusion/inference.py

@@ -268,13 +268,24 @@ class CoreMLTextEncoder(torch.nn.Module):
     """
 
     def __init__(
-        self, mlpackage_path, ref_text_encoder, *, compute_unit=DEFAULT_COMPUTE_UNIT
+        self,
+        mlpackage_path,
+        ref_text_encoder,
+        *,
+        compute_unit=DEFAULT_COMPUTE_UNIT,
+        output_dtype=torch.float16,
     ):
         super().__init__()
         import coremltools as ct
 
         self.config = ref_text_encoder.config
-        self.dtype = torch.float16
+        # The package runs fp16 internally. ``output_dtype`` is the dtype of the
+        # embeddings handed back to the pipeline; keep fp16 (default) for an
+        # all-Core ML pipeline, or set fp32 when the embeddings feed a torch fp32
+        # component (e.g. an OAT config with a Core ML text encoder + torch UNet),
+        # where diffusers propagates ``prompt_embeds.dtype`` to the latents and time
+        # embedding and a Half/Float mismatch would otherwise crash the UNet.
+        self.dtype = output_dtype
         unit = _resolve_unit(compute_unit)
         logger.info(f"Loading text encoder {mlpackage_path} to {unit.name}")
         self.model = ct.models.MLModel(mlpackage_path, compute_units=unit)
@@ -294,14 +305,18 @@ class CoreMLTextEncoder(torch.nn.Module):
         **_ignored,
     ):
         prediction = self.model.predict({"input_ids": _i32(input_ids)})
-        embeds = torch.from_numpy(np.ascontiguousarray(prediction["hidden_states"])).to(
-            input_ids.device
+        embeds = (
+            torch.from_numpy(np.ascontiguousarray(prediction["hidden_states"]))
+            .to(input_ids.device)
+            .to(self.dtype)
         )
         pooled = None
         if self._pooled_name is not None:
-            pooled = torch.from_numpy(
-                np.ascontiguousarray(prediction[self._pooled_name])
-            ).to(input_ids.device)
+            pooled = (
+                torch.from_numpy(np.ascontiguousarray(prediction[self._pooled_name]))
+                .to(input_ids.device)
+                .to(self.dtype)
+            )
         return _CoreMLTextEncoderOutput(embeds, pooled)
 
 

{coreml_diffusion-0.1.3 → coreml_diffusion-0.1.4}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "coreml-diffusion"
 description = "Convert diffusion-model checkpoints (SD1.5/SDXL) to Core ML for Apple Neural Engine — framework-free, ComfyUI-independent."
-version = "0.1.3"
+version = "0.1.4"
 license = "MIT"
 license-files = ["LICENSE"]
 requires-python = ">=3.12,<3.13"

{coreml_diffusion-0.1.3 → coreml_diffusion-0.1.4}/tests/smoke/test_coreml_adapters.py

@@ -118,6 +118,30 @@ def test_coreml_text_encoder_sd15_shape(tmp_path):
 
     assert out[0].shape == (1, SEQ_LEN, HIDDEN)  # no pooled -> [0] is embeds
     assert out.hidden_states[-2].shape == (1, SEQ_LEN, HIDDEN)
+    assert out[0].dtype == torch.float16  # default: fp16 for an all-Core ML pipeline
+
+
+def test_coreml_text_encoder_output_dtype_fp32(tmp_path):
+    # output_dtype=fp32 bridges a Core ML text encoder feeding a torch fp32 UNet
+    # (an OAT config): the embeddings must come back fp32 so diffusers does not
+    # propagate fp16 into the fp32 UNet and crash on a Half/Float mismatch.
+    from transformers import CLIPTextModel
+
+    from coreml_diffusion.inference import CoreMLTextEncoder
+
+    torch.manual_seed(0)
+    encoder = CLIPTextModel(_tiny_clip_config())
+    pkg = tmp_path / "text_encoder.mlpackage"
+    _convert_text_encoder(encoder, ["hidden_states"], pkg, index=None, pooled=False)
+
+    coreml_te = CoreMLTextEncoder(
+        str(pkg), encoder, compute_unit="CPU_ONLY", output_dtype=torch.float32
+    )
+    out = coreml_te(torch.zeros(1, SEQ_LEN, dtype=torch.long))
+
+    assert coreml_te.dtype == torch.float32
+    assert out[0].dtype == torch.float32
+    assert out.last_hidden_state.dtype == torch.float32
 
 
 def test_coreml_text_encoder_sdxl_pooled(tmp_path):

coreml_diffusion-0.1.4/tests/smoke/test_lcm_conversion.py

@@ -0,0 +1,108 @@
+"""Tier 1 smoke: the generic LCM conversion path on a synthetic micro-UNet.
+
+Locks the E-LCM fixes: the ``timestep_cond`` trace input takes its dim from
+``config.time_cond_proj_dim`` (not a hardcoded 256), an LCM-LoRA merge (no
+guidance embedding) is rejected with a clear error, and diffusers-layout
+UNet-only dumps (the canonical full-distill LCM artifact shape) load through
+``load_unet_from_diffusers_state_dict``.
+
+Auto-skips on non-Apple-Silicon hosts so Tier 0 CI on Linux ignores it.
+"""
+
+import platform
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(
+    platform.system() != "Darwin" or platform.machine() != "arm64",
+    reason="Tier 1 requires macOS on Apple Silicon",
+)
+
+GUIDANCE_DIM = 16  # deliberately != 256 to prove the dim comes from the config
+
+TINY_ARCH = dict(
+    sample_size=8,
+    layers_per_block=1,
+    block_out_channels=(32, 64),
+    down_block_types=("CrossAttnDownBlock2D", "DownBlock2D"),
+    up_block_types=("UpBlock2D", "CrossAttnUpBlock2D"),
+    cross_attention_dim=32,
+    attention_head_dim=8,
+    norm_num_groups=8,
+)
+
+
+def _tiny_unet(time_cond_proj_dim=None):
+    from diffusers import UNet2DConditionModel
+
+    return UNet2DConditionModel(**TINY_ARCH, time_cond_proj_dim=time_cond_proj_dim)
+
+
+def test_lcm_unet_converts_with_config_guidance_dim(tmp_path):
+    import coremltools as ct
+
+    from coreml_diffusion.convert import convert_unet
+    from coreml_diffusion.model_version import ModelVersion
+
+    torch.manual_seed(0)
+    out_path = str(tmp_path / "lcm_unet.mlpackage")
+    convert_unet(
+        _tiny_unet(time_cond_proj_dim=GUIDANCE_DIM),
+        ModelVersion.LCM,
+        out_path,
+        sample_size=(8, 8),
+    )
+
+    spec = ct.models.MLModel(out_path, skip_model_load=True).get_spec()
+    inputs = {
+        i.name: tuple(i.type.multiArrayType.shape) for i in spec.description.input
+    }
+    assert "timestep_cond" in inputs
+    assert inputs["timestep_cond"] == (1, GUIDANCE_DIM)
+
+
+def test_lcm_merge_checkpoint_is_rejected(tmp_path):
+    from coreml_diffusion.convert import convert_unet
+    from coreml_diffusion.model_version import ModelVersion
+
+    torch.manual_seed(0)
+    with pytest.raises(ValueError, match="LCM-LoRA merge"):
+        convert_unet(
+            _tiny_unet(time_cond_proj_dim=None),
+            ModelVersion.LCM,
+            str(tmp_path / "merge.mlpackage"),
+            sample_size=(8, 8),
+        )
+
+
+def test_diffusers_layout_dump_loads_with_guidance_embedding(tmp_path):
+    from safetensors.torch import save_file
+
+    from coreml_diffusion.convert import load_unet_from_diffusers_state_dict
+
+    torch.manual_seed(0)
+    source = _tiny_unet(time_cond_proj_dim=GUIDANCE_DIM)
+    dump = tmp_path / "tiny_lcm_dump.safetensors"
+    save_file(source.state_dict(), str(dump))
+
+    # cross_attention_dim and time_cond_proj_dim must be read from the weights;
+    # only the miniature architecture is supplied as overrides.
+    loaded = load_unet_from_diffusers_state_dict(
+        str(dump), **{k: v for k, v in TINY_ARCH.items() if k != "cross_attention_dim"}
+    )
+
+    assert loaded.config.cross_attention_dim == TINY_ARCH["cross_attention_dim"]
+    assert loaded.config.time_cond_proj_dim == GUIDANCE_DIM
+    assert torch.equal(
+        loaded.time_embedding.cond_proj.weight, source.time_embedding.cond_proj.weight
+    )
+
+
+def test_lcm_inputs_legacy_call_keeps_256():
+    # The Suite's LCM converter calls lcm_inputs without a ref UNet; the legacy
+    # 256 fallback is part of the import contract until the Suite migrates.
+    from coreml_diffusion.convert import lcm_inputs
+
+    inputs = lcm_inputs({"sample": torch.rand(2, 4, 8, 8)})
+    assert inputs["timestep_cond"].shape == (2, 256)

coreml_diffusion-0.1.4/tests/unit/test_state_dict_layout.py

@@ -0,0 +1,43 @@
+"""Tier 0: state-dict layout detection routing ``load_unet``.
+
+Locks the predicate that decides whether a single-file checkpoint is an LDM
+checkpoint (``from_single_file`` path) or a diffusers-layout UNet-only dump
+(direct state-dict load — the canonical full-distill LCM artifact shape).
+"""
+
+from coreml_diffusion.conversion.state_dict import is_diffusers_unet_layout
+
+DIFFUSERS_UNET_KEYS = [
+    "conv_in.weight",
+    "time_embedding.linear_1.weight",
+    "time_embedding.cond_proj.weight",
+    "down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_k.weight",
+    "mid_block.resnets.0.conv1.weight",
+    "up_blocks.3.resnets.2.conv2.weight",
+]
+
+LDM_CHECKPOINT_KEYS = [
+    "model.diffusion_model.time_embed.0.weight",
+    "model.diffusion_model.input_blocks.0.0.weight",
+    "first_stage_model.decoder.conv_in.weight",
+    "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight",
+]
+
+
+def test_diffusers_unet_dump_is_detected():
+    assert is_diffusers_unet_layout(DIFFUSERS_UNET_KEYS)
+
+
+def test_ldm_checkpoint_is_not_diffusers_layout():
+    assert not is_diffusers_unet_layout(LDM_CHECKPOINT_KEYS)
+
+
+def test_mixed_prefixes_resolve_to_ldm():
+    # An LDM file whose extras coincidentally include diffusers-looking keys
+    # must still go through from_single_file.
+    assert not is_diffusers_unet_layout(LDM_CHECKPOINT_KEYS + DIFFUSERS_UNET_KEYS)
+
+
+def test_unrelated_keys_are_not_diffusers_layout():
+    assert not is_diffusers_unet_layout(["text_model.encoder.layers.0.mlp.fc1.weight"])
+    assert not is_diffusers_unet_layout([])

{coreml_diffusion-0.1.3 → coreml_diffusion-0.1.4}/uv.lock

@@ -106,7 +106,7 @@ wheels = [
 
 [[package]]
 name = "coreml-diffusion"
-version = "0.1.2"
+version = "0.1.3"
 source = { editable = "." }
 dependencies = [
     { name = "coremltools" },

coreml_diffusion-0.1.3/.release-please-manifest.json

@@ -1,3 +0,0 @@
-{
-  ".": "0.1.3"
-}