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

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

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

coreml_diffusion-0.1.3/CHANGELOG.md

@@ -0,0 +1,26 @@
+# Changelog
+
+## [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.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coreml-diffusion
-Version: 0.1.2
+Version: 0.1.3
 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.3}/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.3}/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.3}/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.3/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.3}/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.3/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.3/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)

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

@@ -19,10 +19,19 @@ import torch
 from diffusers import UNet2DConditionModel
 
 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.text_encoder import (
+    CoreMLTextEncoderWrapper,
+    static_causal_mask,
+)
 from coreml_diffusion.conversion.trace import prepare_unet_for_coreml_trace
 from coreml_diffusion.conversion.unet import CoreMLUNetWrapper
+from coreml_diffusion.conversion.vae import (
+    CoreMLVAEDecoderWrapper,
+    CoreMLVAEEncoderWrapper,
+)
 from coreml_diffusion.logger import logger
 from coreml_diffusion.model_version import ModelVersion
 
@@ -274,9 +283,17 @@ def convert_unet(
     del traced_unet
     gc.collect()
 
+    _palettize_and_save(coreml_unet, unet_out_path, quantize_nbits, "unet")
+
+
+def _palettize_and_save(coreml_model, out_path, quantize_nbits, label):
+    """Optionally k-means palettize the weights, then save the ``.mlpackage``.
+
+    The default path (``quantize_nbits="none"``) saves the model untouched. Shared
+    by the UNet and the VAE / text-encoder conversions so every component gets the
+    same opt-in palettization behaviour and filename-driven cache semantics.
+    """
     if quantize_nbits != "none":
-        # Opt-in k-means weight palettization. The default path
-        # (quantize_nbits="none") leaves the traced UNet untouched.
         from coremltools.optimize.coreml import (
             OpPalettizerConfig,
             OptimizationConfig,
@@ -284,16 +301,144 @@ def convert_unet(
         )
 
         nbits = int(quantize_nbits)
-        logger.info(f"Palettizing UNet weights to {nbits}-bit (kmeans)..")
+        logger.info(f"Palettizing {label} weights to {nbits}-bit (kmeans)..")
         t0 = time.time()
         cfg = OptimizationConfig(
             global_config=OpPalettizerConfig(mode="kmeans", nbits=nbits)
         )
-        coreml_unet = palettize_weights(coreml_unet, config=cfg)
+        coreml_model = palettize_weights(coreml_model, config=cfg)
         logger.info(f"Palettization took {time.time() - t0:.1f}s")
 
-    coreml_unet.save(unet_out_path)
-    logger.info(f"Saved unet into {unet_out_path}")
+    coreml_model.save(out_path)
+    logger.info(f"Saved {label} into {out_path}")
+
+
+def convert_vae_decoder(
+    ref_vae,
+    out_path: str,
+    *,
+    batch_size: int = 1,
+    sample_size: tuple[int, int] = (64, 64),
+    quantize_nbits: str = "none",
+):
+    """Convert ``AutoencoderKL``'s decoder (latent -> image) to a ``.mlpackage``.
+
+    ``sample_size`` is the *latent* spatial size (H/8, W/8); the image output is
+    8x that. The mid-block self-attention is routed through the ORIGINAL
+    (full, fp32-score) processor — diffusers' stock sdpa attention fails to
+    convert under coremltools 9, the same reason ``ORIGINAL`` exists for the UNet.
+    """
+    apply_attention_implementation(ref_vae, "ORIGINAL")
+    wrapper = CoreMLVAEDecoderWrapper(ref_vae.eval()).eval()
+
+    latent_shape = (
+        batch_size,
+        ref_vae.config.latent_channels,
+        sample_size[0],
+        sample_size[1],
+    )
+    example = torch.rand(*latent_shape)
+
+    logger.info(f"JIT tracing VAE decoder (latent {latent_shape})..")
+    traced = torch.jit.trace(wrapper, example)
+
+    inputs = [ct.TensorType(name="latent", shape=latent_shape, dtype=np.float16)]
+    coreml_model = convert_to_coreml("vae_decoder", traced, inputs, ["image"], out_path)
+    del traced
+    gc.collect()
+    _palettize_and_save(coreml_model, out_path, quantize_nbits, "vae_decoder")
+
+
+def convert_vae_encoder(
+    ref_vae,
+    out_path: str,
+    *,
+    batch_size: int = 1,
+    sample_size: tuple[int, int] = (64, 64),
+    quantize_nbits: str = "none",
+):
+    """Convert ``AutoencoderKL``'s encoder (image -> latent moments) to a ``.mlpackage``.
+
+    ``sample_size`` is the latent spatial size; the image input is 8x that. The
+    output ``latent_moments`` is the raw mean‖logvar tensor (2*latent_channels
+    channels) — the pipeline samples from it. Attention handling matches the
+    decoder.
+    """
+    apply_attention_implementation(ref_vae, "ORIGINAL")
+    wrapper = CoreMLVAEEncoderWrapper(ref_vae.eval()).eval()
+
+    image_shape = (batch_size, 3, sample_size[0] * 8, sample_size[1] * 8)
+    example = torch.rand(*image_shape)
+
+    logger.info(f"JIT tracing VAE encoder (image {image_shape})..")
+    traced = torch.jit.trace(wrapper, example)
+
+    inputs = [ct.TensorType(name="image", shape=image_shape, dtype=np.float16)]
+    coreml_model = convert_to_coreml(
+        "vae_encoder", traced, inputs, ["latent_moments"], out_path
+    )
+    del traced
+    gc.collect()
+    _palettize_and_save(coreml_model, out_path, quantize_nbits, "vae_encoder")
+
+
+def convert_text_encoder(
+    ckpt_path: str,
+    model_version: ModelVersion,
+    out_path: str,
+    *,
+    which: int = 1,
+    batch_size: int = 1,
+    quantize_nbits: str = "none",
+):
+    """Convert a CLIP text encoder (token ids -> embeddings) to a ``.mlpackage``.
+
+    ``which`` selects the encoder: 1 = primary (SD1.5/SDXL), 2 = SDXL's second
+    (``CLIPTextModelWithProjection``). SDXL uses the penultimate hidden state from
+    both encoders plus encoder 2's projected pooled output; SD1.5 uses encoder 1's
+    final ``last_hidden_state``. ``input_ids`` crosses the Core ML boundary as
+    int32 over the fixed 77-token sequence.
+    """
+    encoders = load_text_encoders(ckpt_path, model_version)
+    if which == 2:
+        if len(encoders) < 2:
+            raise ValueError(
+                f"{model_version.name} has no second text encoder "
+                "(text_encoder_2 is SDXL-only)."
+            )
+        encoder = encoders[1]
+        hidden_states_index = -2
+        output_pooled = True
+        output_names = ["hidden_states", "pooled_embeds"]
+    else:
+        encoder = encoders[0]
+        is_sdxl = model_version in {ModelVersion.SDXL, ModelVersion.SDXL_REFINER}
+        hidden_states_index = -2 if is_sdxl else None
+        output_pooled = False
+        output_names = ["hidden_states"]
+
+    wrapper = CoreMLTextEncoderWrapper(
+        encoder.eval(),
+        hidden_states_index=hidden_states_index,
+        output_pooled=output_pooled,
+    ).eval()
+
+    ids_shape = (batch_size, TEXT_TOKEN_SEQUENCE_LENGTH)
+    # Trace with int64 (torch embedding gather needs long); declare int32 at the
+    # Core ML boundary (coremltools casts the gather index op).
+    example = torch.zeros(ids_shape, dtype=torch.int64)
+
+    logger.info(f"JIT tracing text encoder {which} (input_ids {ids_shape})..")
+    with static_causal_mask(TEXT_TOKEN_SEQUENCE_LENGTH):
+        traced = torch.jit.trace(wrapper, example)
+
+    inputs = [ct.TensorType(name="input_ids", shape=ids_shape, dtype=np.int32)]
+    coreml_model = convert_to_coreml(
+        f"text_encoder_{which}", traced, inputs, output_names, out_path
+    )
+    del traced
+    gc.collect()
+    _palettize_and_save(coreml_model, out_path, quantize_nbits, f"text_encoder_{which}")
 
 
 def convert(
@@ -301,6 +446,7 @@ def convert(
     model_version: ModelVersion,
     out_path: str,
     *,
+    component: str = Component.UNET.value,
     batch_size: int = 1,
     sample_size: tuple[int, int] = (64, 64),
     controlnet_support: bool = False,
@@ -309,16 +455,53 @@ def convert(
     config_path: str = None,
     quantize_nbits: str = "none",
 ):
-    """Convert a single-file checkpoint's UNet to a Core ML ``.mlpackage``.
-
-    Keyword-only past the three required positionals so the package can add
-    capabilities (new keyword args) without breaking an older caller — the
-    versioned interface contract. Writes ``out_path``; returns None.
+    """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.
     """
     if os.path.exists(out_path):
         logger.info(f"Found existing model at {out_path}! Skipping..")
         return
 
+    comp = Component(component)
+
+    if comp is Component.VAE_DECODER:
+        convert_vae_decoder(
+            load_vae(ckpt_path),
+            out_path,
+            batch_size=batch_size,
+            sample_size=sample_size,
+            quantize_nbits=quantize_nbits,
+        )
+        return
+
+    if comp is Component.VAE_ENCODER:
+        convert_vae_encoder(
+            load_vae(ckpt_path),
+            out_path,
+            batch_size=batch_size,
+            sample_size=sample_size,
+            quantize_nbits=quantize_nbits,
+        )
+        return
+
+    if comp in {Component.TEXT_ENCODER, Component.TEXT_ENCODER_2}:
+        convert_text_encoder(
+            ckpt_path,
+            model_version,
+            out_path,
+            which=2 if comp is Component.TEXT_ENCODER_2 else 1,
+            batch_size=batch_size,
+            quantize_nbits=quantize_nbits,
+        )
+        return
+
     if attn_impl not in ATTENTION_IMPLEMENTATIONS:
         raise ValueError(
             f"Unsupported attention implementation {attn_impl!r}. "
@@ -350,3 +533,42 @@ def load_unet(ckpt_path, config_path):
         ckpt_path,
         original_config=config_path,
     )
+
+
+def load_vae(ckpt_path):
+    """Load ``AutoencoderKL`` from a single-file checkpoint."""
+    from diffusers import AutoencoderKL
+
+    return AutoencoderKL.from_single_file(ckpt_path)
+
+
+# Pipeline class whose ``from_single_file`` extracts the text encoder(s) for a
+# given model version. SDXL pipelines expose ``text_encoder_2``; SD15/LCM do not.
+_TEXT_ENCODER_PIPELINE = {
+    ModelVersion.SD15: ("diffusers", "StableDiffusionPipeline"),
+    ModelVersion.LCM: ("diffusers", "StableDiffusionPipeline"),
+    ModelVersion.SDXL: ("diffusers", "StableDiffusionXLPipeline"),
+    ModelVersion.SDXL_REFINER: ("diffusers", "StableDiffusionXLPipeline"),
+}
+
+
+def load_text_encoders(ckpt_path, model_version):
+    """Return the checkpoint's CLIP text encoder(s) as a list (1 for SD1.5, 2 for SDXL).
+
+    Loads the matching diffusers pipeline from the single-file checkpoint and
+    extracts its text encoder(s). This pulls the whole pipeline (UNet/VAE
+    included) — acceptable for an offline, one-shot conversion and the most robust
+    way to recover correctly-configured CLIP weights from a single file.
+    """
+    import importlib
+
+    if model_version not in _TEXT_ENCODER_PIPELINE:
+        raise ValueError(f"No text-encoder pipeline mapped for {model_version!r}.")
+    module_name, class_name = _TEXT_ENCODER_PIPELINE[model_version]
+    pipeline_cls = getattr(importlib.import_module(module_name), class_name)
+
+    pipe = pipeline_cls.from_single_file(ckpt_path)
+    encoders = [pipe.text_encoder]
+    if getattr(pipe, "text_encoder_2", None) is not None:
+        encoders.append(pipe.text_encoder_2)
+    return encoders