coreml-diffusion 0.1.2__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.4/CHANGELOG.md

@@ -0,0 +1,34 @@
+# 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)
+
+
+### ✨ Features
+
+* **convert:** add VAE and CLIP text-encoder conversion ([dc1f85b](https://github.com/aszc-dev/coreml-diffusion/commit/dc1f85bafe50d36655ff7ece0c052a30fd77bb81))
+* **inference:** end-to-end Core ML pipeline (VAE + text-encoder swap) ([ca08b16](https://github.com/aszc-dev/coreml-diffusion/commit/ca08b16729529afbdf610d0e8ec2d09b849080c6))
+
+
+### 🐛 Bug Fixes
+
+* **inference:** expose .device on the Core ML adapters ([30a673e](https://github.com/aszc-dev/coreml-diffusion/commit/30a673eebe3927722214d0ab6a44fbc344d18f3a))
+
+
+### 📚 Documentation
+
+* **readme:** link the log.aszc.dev energy benchmark writeup ([77927b5](https://github.com/aszc-dev/coreml-diffusion/commit/77927b5dd5f1311a3b3c317692f3a347e3976a54))
+
+## [0.1.2](https://github.com/aszc-dev/coreml-diffusion/compare/v0.1.1...v0.1.2) (2026-05-27)
+
+
+### 🐛 Bug Fixes
+
+* **attention:** convertible fp32 ORIGINAL attention for the Core ML GPU path ([#2](https://github.com/aszc-dev/coreml-diffusion/issues/2)) ([28e56fc](https://github.com/aszc-dev/coreml-diffusion/commit/28e56fcf8c2242ebbe4c05abd05f7e796069d7d1))

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coreml-diffusion
-Version: 0.1.2
+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
@@ -44,6 +44,12 @@ GPU-free, embeddable in a Swift/iOS app. ANE is the differentiator — this is a
 feasibility and power efficiency for SD1.5/SDXL on ANE, not a raw-throughput claim
 against desktop GPUs.
 
+The power-efficiency claim is measured: in a cross-backend benchmark the ct9
+converter here runs the SD1.5 UNet on the ANE at **6-7x lower energy** than
+GPU/MPS, at the same speed — see the writeup,
+[The ANE runs the SD1.5 UNet at 6-7x lower energy than GPU/MPS](https://log.aszc.dev/ane-vs-gpu-mps-sd15-unet-energy/),
+for the methodology and the numerical-divergence tradeoff.
+
 The scope is diffusion architectures generally, not Stable Diffusion specifically.
 The project aims to gather, in one place: the conversion path, a reproducible
 benchmarking suite for objective comparison, a per-model catalogue documenting the

{coreml_diffusion-0.1.2 → coreml_diffusion-0.1.4}/README.md

@@ -15,6 +15,12 @@ GPU-free, embeddable in a Swift/iOS app. ANE is the differentiator — this is a
 feasibility and power efficiency for SD1.5/SDXL on ANE, not a raw-throughput claim
 against desktop GPUs.
 
+The power-efficiency claim is measured: in a cross-backend benchmark the ct9
+converter here runs the SD1.5 UNet on the ANE at **6-7x lower energy** than
+GPU/MPS, at the same speed — see the writeup,
+[The ANE runs the SD1.5 UNet at 6-7x lower energy than GPU/MPS](https://log.aszc.dev/ane-vs-gpu-mps-sd15-unet-energy/),
+for the methodology and the numerical-divergence tradeoff.
+
 The scope is diffusion architectures generally, not Stable Diffusion specifically.
 The project aims to gather, in one place: the conversion path, a reproducible
 benchmarking suite for objective comparison, a per-model catalogue documenting the

{coreml_diffusion-0.1.2 → coreml_diffusion-0.1.4}/coreml_diffusion/__init__.py

@@ -23,9 +23,11 @@ because a saved workflow JSON references these strings verbatim.
 from enum import Enum
 
 from coreml_diffusion.attention import ATTENTION_IMPLEMENTATIONS
+from coreml_diffusion.component import CONVERTIBLE_COMPONENTS
 from coreml_diffusion.model_version import ModelVersion
 from coreml_diffusion.naming import (
     QUANT_NBITS_VALUES,
+    compose_component_name,
     compose_out_name,
     lora_names_from_params,
 )
@@ -36,12 +38,16 @@ __all__ = [
     "list_model_versions",
     "list_attention_impls",
     "list_quant_modes",
+    "list_convertible_components",
     "CONTRACT_VERSION",
     "compose_out_name",
+    "compose_component_name",
     "lora_names_from_params",
     "convert",
     "build_pipeline",
     "CoreMLUNet",
+    "CoreMLVAE",
+    "CoreMLTextEncoder",
 ]
 
 
@@ -91,9 +97,20 @@ def list_quant_modes() -> list[str]:
     return list(QUANT_NBITS_VALUES)
 
 
+def list_convertible_components() -> list[str]:
+    """Convertible components, e.g. ``["unet", "vae_decoder", ...]``.
+
+    ``"unet"`` is the historical default; the rest are the VAE / text-encoder
+    extension. ``"text_encoder_2"`` is only meaningful for SDXL — validity per
+    model version is enforced at convert time, not advertised here.
+    """
+    return list(CONVERTIBLE_COMPONENTS)
+
+
 # Discovery-contract version. Bump per the additive-only rules in this module's
 # docstring and CONVERTER_EXTRACTION_SPEC.md "Interface contract".
-CONTRACT_VERSION = "1.0"
+# 1.1: added list_convertible_components (VAE + text-encoder conversion).
+CONTRACT_VERSION = "1.1"
 
 
 def __getattr__(name):
@@ -107,7 +124,7 @@ def __getattr__(name):
         from coreml_diffusion.convert import convert as _convert
 
         return _convert
-    if name in ("build_pipeline", "CoreMLUNet"):
+    if name in ("build_pipeline", "CoreMLUNet", "CoreMLVAE", "CoreMLTextEncoder"):
         from coreml_diffusion import inference
 
         return getattr(inference, name)

{coreml_diffusion-0.1.2 → coreml_diffusion-0.1.4}/coreml_diffusion/cli.py

@@ -37,6 +37,7 @@ def _convert_cmd(args):
         ckpt,
         coreml_diffusion.ModelVersion[args.model_version],
         args.out,
+        component=args.component,
         batch_size=args.batch_size,
         sample_size=sample_size,
         controlnet_support=args.controlnet,
@@ -99,6 +100,13 @@ def build_parser():
         choices=coreml_diffusion.list_model_versions(include_experimental=True),
         help="Model architecture (verified: SD15, SDXL; experimental otherwise)",
     )
+    conv.add_argument(
+        "--component",
+        choices=coreml_diffusion.list_convertible_components(),
+        default="unet",
+        help="Checkpoint component to convert (default unet). VAE/text-encoder "
+        "components ignore --attn-impl/--controlnet/--lora. text_encoder_2 is SDXL-only",
+    )
     conv.add_argument("--out", required=True, help="Output .mlpackage path to write")
     conv.add_argument(
         "--height", type=int, default=512, help="Target image height (default 512)"

coreml_diffusion-0.1.4/coreml_diffusion/component.py

@@ -0,0 +1,32 @@
+"""Convertible model components — framework-free leaf.
+
+A single-file checkpoint bundles several sub-models; ``coreml_diffusion`` can
+convert each into its own ``.mlpackage``. This enum is the canonical identifier
+set, mirrored into the discovery contract (``list_convertible_components``) and
+the naming contract (``compose_component_name``).
+
+``UNET`` keeps its historical conversion path and filename (``compose_out_name``);
+the other components are the additive VAE / text-encoder extension. ``.value`` is
+the wire string used by the CLI ``--component`` flag and the discovery list — kept
+lowercase and stable, since a saved workflow / benchmark manifest references it
+verbatim (same additive-only rule as ``ModelVersion``).
+
+``TEXT_ENCODER_2`` is SDXL-only (its second CLIP, ``CLIPTextModelWithProjection``);
+on SD1.5 only ``TEXT_ENCODER`` exists. Validity per model version is enforced at
+convert time, not encoded here.
+"""
+
+from enum import Enum
+
+
+class Component(Enum):
+    UNET = "unet"
+    VAE_DECODER = "vae_decoder"
+    VAE_ENCODER = "vae_encoder"
+    TEXT_ENCODER = "text_encoder"
+    TEXT_ENCODER_2 = "text_encoder_2"
+
+
+# Declaration order is the discovery-list order; UNET leads to match the
+# historical, primary conversion path.
+CONVERTIBLE_COMPONENTS = tuple(c.value for c in Component)

{coreml_diffusion-0.1.2 → coreml_diffusion-0.1.4}/coreml_diffusion/conversion/attention.py

@@ -120,9 +120,13 @@ def _attention_forward(
     input_ndim = hidden_states.ndim
     if input_ndim == 4:
         batch_size, channel, height, width = hidden_states.shape
-        hidden_states = hidden_states.view(
-            batch_size, channel, height * width
-        ).transpose(1, 2)
+        # flatten(2) instead of view(B, C, height * width): the explicit
+        # height * width multiplies two traced size ints, emitting an aten::mul ->
+        # aten::Int that coremltools 9 cannot fold to a const (it fails the
+        # conversion). flatten collapses the spatial dims with a single reshape and
+        # no symbolic product. Only the 4D path (VAE self-attention) hits this; the
+        # UNet routes attention through a 3D tensor, so ORIGINAL there is untouched.
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
     else:
         batch_size, _, channel = hidden_states.shape
         height = None

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.4/coreml_diffusion/conversion/text_encoder.py

@@ -0,0 +1,85 @@
+"""CLIP text-encoder wrapper adapting it to a flat Core ML tensor contract.
+
+Wraps a transformers ``CLIPTextModel`` (SD1.5, SDXL encoder 1) or
+``CLIPTextModelWithProjection`` (SDXL encoder 2) so the traced graph takes a
+single ``input_ids`` tensor ``(B, 77)`` and returns the embeddings the diffusion
+pipeline consumes — nothing else.
+
+Which hidden state and whether a pooled vector is emitted depends on the model:
+
+  SD1.5            : final ``last_hidden_state`` (768), no pooled
+  SDXL encoder 1   : penultimate ``hidden_states[-2]`` (768), no pooled
+  SDXL encoder 2   : penultimate ``hidden_states[-2]`` (1280) + projected pooled (1280)
+
+The penultimate selection is SDXL's documented behaviour (it concatenates the
+two encoders' penultimate states and uses encoder 2's projected pooled output as
+the ``add_embeds`` conditioning). ``hidden_states_index=None`` selects the final
+``last_hidden_state``; an int indexes ``hidden_states`` directly.
+
+The pooled vector prefers ``text_embeds`` (the projection head, encoder 2) and
+falls back to ``pooler_output`` — so the same wrapper serves both CLIP variants.
+``input_ids`` is fed as int32 at the Core ML boundary (see ``convert``).
+"""
+
+import contextlib
+
+import torch
+
+
+@contextlib.contextmanager
+def static_causal_mask(seq_len):
+    """Patch CLIP's causal-mask builder to a constant for the trace duration.
+
+    transformers builds the causal mask from ``query_length + past_key_values_length``,
+    both traced 0-dim size tensors; the resulting ``aten::Int`` cannot be folded to a
+    const under coremltools 9 (conversion fails). Since the converted sequence length
+    is fixed at trace time, swap ``_create_4d_causal_attention_mask`` for a closure
+    that materialises the upper-triangular ``-inf`` mask from a Python-int ``seq_len``
+    — a pure constant the frontend folds away. Mirrors ``prepare_unet_for_coreml_trace``:
+    a trace-only enablement patch, restored on exit. Shape ``(1, 1, seq, seq)``
+    broadcasts over batch/heads, so no symbolic batch leaks back in.
+    """
+    from transformers.models.clip import modeling_clip
+
+    original = modeling_clip._create_4d_causal_attention_mask
+
+    def _const_mask(input_shape, dtype, device=None, *args, **kwargs):
+        mask = torch.full(
+            (seq_len, seq_len), torch.finfo(dtype).min, dtype=dtype, device=device
+        )
+        return torch.triu(mask, diagonal=1)[None, None]
+
+    modeling_clip._create_4d_causal_attention_mask = _const_mask
+    try:
+        yield
+    finally:
+        modeling_clip._create_4d_causal_attention_mask = original
+
+
+class CoreMLTextEncoderWrapper(torch.nn.Module):
+    """token ids ``(B, 77)`` -> embeddings (+ optional pooled)."""
+
+    def __init__(self, text_encoder, *, hidden_states_index=None, output_pooled=False):
+        super().__init__()
+        self.text_encoder = text_encoder
+        self.hidden_states_index = hidden_states_index
+        self.output_pooled = output_pooled
+
+    def forward(self, input_ids):
+        out = self.text_encoder(
+            input_ids,
+            output_hidden_states=self.hidden_states_index is not None,
+            return_dict=True,
+        )
+        if self.hidden_states_index is None:
+            embeds = out.last_hidden_state
+        else:
+            embeds = out.hidden_states[self.hidden_states_index]
+
+        if not self.output_pooled:
+            return embeds
+
+        pooled = getattr(out, "text_embeds", None)
+        if pooled is None:
+            pooled = out.pooler_output
+        return embeds, pooled

coreml_diffusion-0.1.4/coreml_diffusion/conversion/vae.py

@@ -0,0 +1,49 @@
+"""VAE wrappers adapting ``AutoencoderKL`` to a flat Core ML tensor contract.
+
+Two thin modules, one per direction, mirroring ``CoreMLUNetWrapper``: they expose
+a single positional tensor in / single tensor out so the traced graph has a clean,
+named Core ML I/O signature. They call the VAE *submodules* directly
+(``post_quant_conv``/``decoder``, ``encoder``/``quant_conv``) rather than
+``decode``/``encode`` — the same op graph, but free of the ``return_dict``
+plumbing and the ``DiagonalGaussianDistribution`` wrapper that complicate tracing.
+
+Scaling is intentionally NOT baked in. The pipeline owns ``scaling_factor``
+(``latent = latent / scaling_factor`` before decode, ``moments`` -> distribution
+-> sample -> ``* scaling_factor`` after encode), keeping these artifacts 1:1 with
+the reference VAE.
+
+The mid-block self-attention is converted via the ORIGINAL (full, fp32-score)
+processor; see ``convert.convert_vae_*``.
+"""
+
+import torch
+
+
+class CoreMLVAEDecoderWrapper(torch.nn.Module):
+    """latent ``(B, latent_channels, h, w)`` -> image ``(B, 3, h*8, w*8)``."""
+
+    def __init__(self, vae):
+        super().__init__()
+        self.vae = vae
+
+    def forward(self, latent):
+        z = self.vae.post_quant_conv(latent)
+        return self.vae.decoder(z)
+
+
+class CoreMLVAEEncoderWrapper(torch.nn.Module):
+    """image ``(B, 3, h*8, w*8)`` -> latent moments ``(B, 2*latent_channels, h, w)``.
+
+    Outputs the raw moments (mean ‖ logvar) exactly as ``AutoencoderKL.encode``
+    produces them before wrapping in a ``DiagonalGaussianDistribution``. Sampling
+    (mean + std·noise) is deferred to the pipeline so the converted encoder stays
+    deterministic and noise-source agnostic.
+    """
+
+    def __init__(self, vae):
+        super().__init__()
+        self.vae = vae
+
+    def forward(self, image):
+        h = self.vae.encoder(image)
+        return self.vae.quant_conv(h)