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

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

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

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

@@ -1,5 +1,20 @@
 # Changelog
 
+## [0.1.5](https://github.com/aszc-dev/coreml-diffusion/compare/v0.1.4...v0.1.5) (2026-06-13)
+
+
+### ✨ Features
+
+* **convert:** auto-detect model version from the checkpoint ([2a24d4e](https://github.com/aszc-dev/coreml-diffusion/commit/2a24d4efd196100dbdd0bf9d5dd61c6cce31d2ac))
+
+## [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.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coreml-diffusion
-Version: 0.1.3
+Version: 0.1.5
 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.3 → coreml_diffusion-0.1.5}/coreml_diffusion/__init__.py

@@ -44,6 +44,7 @@ __all__ = [
     "compose_component_name",
     "lora_names_from_params",
     "convert",
+    "detect_model_version",
     "build_pipeline",
     "CoreMLUNet",
     "CoreMLVAE",
@@ -123,7 +124,23 @@ def __getattr__(name):
     if name == "convert":
         from coreml_diffusion.convert import convert as _convert
 
+        # Importing the submodule binds ``coreml_diffusion.convert`` to the
+        # MODULE as a side effect, which shadows this function on every later
+        # access (a module object isn't callable). Rebind the package attribute
+        # to the function so ``coreml_diffusion.convert(...)`` stays callable in
+        # long-lived processes (e.g. a ComfyUI server doing >1 conversion).
+        globals()["convert"] = _convert
         return _convert
+    if name == "detect_model_version":
+        # Lives in the framework-free state_dict module (reads only the
+        # safetensors header), so exposing it never drags coremltools/diffusers
+        # into the import path.
+        from coreml_diffusion.conversion.state_dict import (
+            detect_model_version as _detect,
+        )
+
+        globals()["detect_model_version"] = _detect
+        return _detect
     if name in ("build_pipeline", "CoreMLUNet", "CoreMLVAE", "CoreMLTextEncoder"):
         from coreml_diffusion import inference
 

{coreml_diffusion-0.1.3 → coreml_diffusion-0.1.5}/coreml_diffusion/cli.py

@@ -33,9 +33,14 @@ def _convert_cmd(args):
     sample_size = (args.height // 8, args.width // 8)
     lora_weights = [_parse_lora(spec) for spec in (args.lora or [])]
     ckpt = sources.resolve_checkpoint(args.ckpt, args.source)
+    model_version = (
+        coreml_diffusion.ModelVersion[args.model_version]
+        if args.model_version
+        else None
+    )
     coreml_diffusion.convert(
         ckpt,
-        coreml_diffusion.ModelVersion[args.model_version],
+        model_version,
         args.out,
         component=args.component,
         batch_size=args.batch_size,
@@ -94,11 +99,13 @@ def build_parser():
     )
     conv.add_argument(
         "--model-version",
-        required=True,
-        # include experimental: the CLI is the power-user path. Experimental
-        # versions (LCM, SDXL_REFINER) convert but are not golden-verified.
+        default=None,
+        # Auto-detected from the checkpoint when omitted. Choices stay available
+        # as an explicit override (the CLI is the power-user path; experimental
+        # versions convert but are not golden-verified).
         choices=coreml_diffusion.list_model_versions(include_experimental=True),
-        help="Model architecture (verified: SD15, SDXL; experimental otherwise)",
+        help="Model architecture; auto-detected from the checkpoint when omitted "
+        "(verified: SD15, SDXL; experimental otherwise)",
     )
     conv.add_argument(
         "--component",

coreml_diffusion-0.1.5/coreml_diffusion/conversion/state_dict.py

@@ -0,0 +1,99 @@
+"""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.
+
+Also home to ``detect_model_version``: it reads only the safetensors header (no
+coremltools/diffusers), so the conversion entrypoint can auto-pick the model
+version without dragging the heavy stack into the discovery path.
+"""
+
+from coreml_diffusion.model_version import ModelVersion
+
+DIFFUSERS_UNET_KEY_PREFIXES = ("down_blocks.", "up_blocks.", "mid_block.")
+LDM_UNET_KEY_PREFIX = "model.diffusion_model."
+
+# cross_attention_dim -> model version. The context (key/value) dim of the UNet's
+# cross-attention is the architecture fingerprint: SD1.5 conditions on a single
+# 768-dim CLIP, SDXL on the 2048-dim concat of both encoders, the SDXL refiner on
+# the 1280-dim OpenCLIP-bigG alone. A guidance embedding (time_embedding.cond_proj)
+# on top of the 768-dim SD1.5 stack marks a full-distill LCM.
+_CROSS_ATTENTION_DIM_TO_VERSION = {
+    768: ModelVersion.SD15,
+    2048: ModelVersion.SDXL,
+    1280: ModelVersion.SDXL_REFINER,
+}
+
+
+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
+
+
+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 detect_model_version(ckpt_path):
+    """Infer the ``ModelVersion`` from a checkpoint's UNet weights.
+
+    Reads two architecture fingerprints straight from the safetensors header
+    (no full model load): the cross-attention context dim (``attn2.to_k``) and
+    whether a guidance embedding (``cond_proj``) is present. Works for both LDM
+    and diffusers key layouts. Raises ``ValueError`` carrying the observed
+    evidence when the architecture is unrecognised, so a bad guess is debuggable
+    rather than silent.
+    """
+    keys = safetensors_keys(ckpt_path)
+    if keys is None:
+        raise ValueError(
+            f"Cannot auto-detect model version from {ckpt_path!r}: not a readable "
+            "safetensors file. Pass model_version explicitly."
+        )
+    cross_attn_key = next((k for k in keys if k.endswith("attn2.to_k.weight")), None)
+    if cross_attn_key is None:
+        raise ValueError(
+            f"Cannot auto-detect model version from {ckpt_path!r}: no cross-attention "
+            "(attn2.to_k) weights found. Pass model_version explicitly."
+        )
+    from safetensors import safe_open
+
+    with safe_open(ckpt_path, framework="pt") as f:
+        cross_attention_dim = f.get_slice(cross_attn_key).get_shape()[1]
+    has_guidance_embedding = any(k.endswith("cond_proj.weight") for k in keys)
+
+    if has_guidance_embedding:
+        if cross_attention_dim == 768:
+            return ModelVersion.LCM
+        raise ValueError(
+            f"Cannot auto-detect model version from {ckpt_path!r}: found a guidance "
+            f"embedding (LCM) but cross_attention_dim={cross_attention_dim}; only "
+            "SD1.5-class LCM (cross_attention_dim=768) is supported. Pass "
+            "model_version explicitly."
+        )
+    version = _CROSS_ATTENTION_DIM_TO_VERSION.get(cross_attention_dim)
+    if version is None:
+        raise ValueError(
+            f"Cannot auto-detect model version from {ckpt_path!r}: unrecognised "
+            f"cross_attention_dim={cross_attention_dim}. Supported: 768 (SD15/LCM), "
+            "2048 (SDXL), 1280 (SDXL_REFINER). Pass model_version explicitly."
+        )
+    return version

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

@@ -22,6 +22,11 @@ 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 (
+    detect_model_version,
+    is_diffusers_unet_layout,
+    safetensors_keys,
+)
 from coreml_diffusion.conversion.text_encoder import (
     CoreMLTextEncoderWrapper,
     static_causal_mask,
@@ -139,9 +144,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 +270,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)
@@ -443,8 +464,8 @@ def convert_text_encoder(
 
 def convert(
     ckpt_path: str,
-    model_version: ModelVersion,
-    out_path: str,
+    model_version: ModelVersion = None,
+    out_path: str = None,
     *,
     component: str = Component.UNET.value,
     batch_size: int = 1,
@@ -457,20 +478,28 @@ def convert(
 ):
     """Convert a single-file checkpoint component to a Core ML ``.mlpackage``.
 
-    ``component`` selects what to convert (default ``"unet"`` — historical
-    behaviour, all UNet-only kwargs apply). ``vae_decoder`` / ``vae_encoder`` /
-    ``text_encoder`` / ``text_encoder_2`` convert the corresponding sub-model and
-    ignore the UNet-only kwargs (``controlnet_support``, ``lora_weights``,
-    ``attn_impl``). Keyword-only past the three required positionals so the package
-    can add capabilities without breaking an older caller. Writes ``out_path``;
-    returns None.
+    ``model_version`` is auto-detected from the checkpoint when left ``None``
+    (the architecture fully determines the conversion); pass it explicitly only
+    to override a misdetection. ``component`` selects what to convert (default
+    ``"unet"`` — historical behaviour, all UNet-only kwargs apply).
+    ``vae_decoder`` / ``vae_encoder`` / ``text_encoder`` / ``text_encoder_2``
+    convert the corresponding sub-model and ignore the UNet-only kwargs
+    (``controlnet_support``, ``lora_weights``, ``attn_impl``). Keyword-only past
+    the leading positionals so the package can add capabilities without breaking
+    an older caller. Writes ``out_path``; returns None.
     """
+    if out_path is None:
+        raise TypeError("convert() requires out_path")
     if os.path.exists(out_path):
         logger.info(f"Found existing model at {out_path}! Skipping..")
         return
 
     comp = Component(component)
 
+    if model_version is None:
+        model_version = detect_model_version(ckpt_path)
+        logger.info(f"Auto-detected model version: {model_version.name}")
+
     if comp is Component.VAE_DECODER:
         convert_vae_decoder(
             load_vae(ckpt_path),
@@ -529,12 +558,48 @@ 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 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.5}/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.5}/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.5"
 license = "MIT"
 license-files = ["LICENSE"]
 requires-python = ">=3.12,<3.13"

{coreml_diffusion-0.1.3 → coreml_diffusion-0.1.5}/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.5/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.5/tests/unit/test_detect_model_version.py

@@ -0,0 +1,77 @@
+"""Tier 0: model-version auto-detection from checkpoint weights.
+
+Locks the architecture fingerprinting that lets ``convert(model_version=None)``
+pick the right conversion path: cross-attention context dim (attn2.to_k) plus
+the presence of a guidance embedding (cond_proj). Synthetic safetensors files
+carry only the two keys the detector reads, so the test stays framework-free.
+"""
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from coreml_diffusion.conversion.state_dict import detect_model_version
+from coreml_diffusion.model_version import ModelVersion
+
+
+def _write_ckpt(path, cross_attention_dim, *, guidance=False, with_cross_attn=True):
+    tensors = {}
+    if with_cross_attn:
+        # attn2.to_k maps the context dim -> inner dim; shape[1] is what we read.
+        key = "down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_k.weight"
+        tensors[key] = torch.zeros(320, cross_attention_dim)
+    if guidance:
+        tensors["time_embedding.cond_proj.weight"] = torch.zeros(320, 256)
+    if not tensors:
+        tensors["dummy"] = torch.zeros(1)
+    save_file(tensors, str(path))
+    return str(path)
+
+
+@pytest.mark.parametrize(
+    "cross_attention_dim, guidance, expected",
+    [
+        (768, False, ModelVersion.SD15),
+        (768, True, ModelVersion.LCM),
+        (2048, False, ModelVersion.SDXL),
+        (1280, False, ModelVersion.SDXL_REFINER),
+    ],
+)
+def test_detects_known_architectures(tmp_path, cross_attention_dim, guidance, expected):
+    ckpt = _write_ckpt(
+        tmp_path / "m.safetensors", cross_attention_dim, guidance=guidance
+    )
+    assert detect_model_version(ckpt) is expected
+
+
+def test_lcm_lora_merge_detects_as_sd15(tmp_path):
+    # An LCM-LoRA merge is plain SD1.5 architecture (no guidance embedding); it
+    # must NOT be mistaken for a full-distill LCM.
+    ckpt = _write_ckpt(tmp_path / "merge.safetensors", 768, guidance=False)
+    assert detect_model_version(ckpt) is ModelVersion.SD15
+
+
+def test_guidance_with_non_sd15_dim_is_rejected(tmp_path):
+    ckpt = _write_ckpt(tmp_path / "sdxl_lcm.safetensors", 2048, guidance=True)
+    with pytest.raises(ValueError, match="only SD1.5-class LCM"):
+        detect_model_version(ckpt)
+
+
+def test_unknown_cross_attention_dim_is_rejected(tmp_path):
+    # 1024 is SD2.x — unsupported; the error must name the observed dim.
+    ckpt = _write_ckpt(tmp_path / "sd2.safetensors", 1024)
+    with pytest.raises(ValueError, match="cross_attention_dim=1024"):
+        detect_model_version(ckpt)
+
+
+def test_no_cross_attention_weights_is_rejected(tmp_path):
+    ckpt = _write_ckpt(tmp_path / "weird.safetensors", 768, with_cross_attn=False)
+    with pytest.raises(ValueError, match="no cross-attention"):
+        detect_model_version(ckpt)
+
+
+def test_non_safetensors_is_rejected(tmp_path):
+    bogus = tmp_path / "model.ckpt"
+    bogus.write_bytes(b"not safetensors")
+    with pytest.raises(ValueError, match="not a readable safetensors"):
+        detect_model_version(str(bogus))

coreml_diffusion-0.1.5/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.5}/uv.lock

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

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

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