# All-In-One Checkpoint Support Design

## Purpose

This document describes how `diffusion-cli` should support Z-Image
all-in-one checkpoint files.

The current program expects three independent safetensors files:

- A Z-Image diffusion model.
- A Qwen text encoder.
- A VAE.

Some local Z-Image distributions package those same components into one
checkpoint file. One example is:

```text
~/documents/ai/diffusion/models/checkpoints/z-image_turbo/redcraftRedzimageUpdatedDEC03_redzimage15AIO.safetensors
```

The goal is to let users point the CLI at that one checkpoint instead of
requiring them to provide three separate component files. The generation
algorithm, model architecture, tokenizer handling, sampler, scheduler,
and VAE decode math should stay unchanged. Only the way weights are
located and loaded should change.

## Requirements

The feature must:

- Add command line support for a local all-in-one checkpoint path.
- Add config file support for the same checkpoint path.
- Continue supporting the existing separate model files.
- Let explicit component paths override the all-in-one checkpoint for
  that component.
- Keep tokenizer configuration separate from the checkpoint.
- Avoid importing or depending on ComfyUI at runtime.
- Avoid downloading any files.
- Avoid loading unrelated checkpoint tensors when only one component is
  needed.
- Produce clear errors when the checkpoint is missing a required
  component.
- Extend inspection output so users can see what an all-in-one
  checkpoint contains.

The feature should not:

- Add support for LoRA loading.
- Add support for checkpoint formats other than safetensors.
- Add a new Z-Image architecture.
- Change prompts, sampling, scheduler, latent scaling, or output image
  writing.
- Extract component files to disk as an implementation detail.

## Current State

The current command line interface exposes independent model paths in
`diffusion_cli/cli.py`:

```text
--diffusion-model
--text-encoder
--vae
--tokenizer-path
```

The current config model mirrors those paths in `[models]`:

```toml
[models]
diffusion_model = "/models/z-image/diffusion.safetensors"
text_encoder = "/models/z-image/qwen_3_4b.safetensors"
vae = "/models/z-image/ae.safetensors"
tokenizer = "/models/z-image/tokenizer"
```

`diffusion_cli/paths.py` validates those paths and returns a
`ModelFiles` object. That object assumes every component is backed by a
different regular file.

Generation then loads the components in this order:

1. `ZImageTextEncoder` loads the Qwen encoder and encodes the prompt.
2. `ZImageModel` loads the diffusion model and samples latents.
3. `ZImageVae` loads the VAE decoder and decodes the final latent.

The order is important. It lets the program release Python and CUDA
memory between large stages. The all-in-one implementation should keep
that staged loading behavior.

## Observed Checkpoint Layout

The example all-in-one safetensors checkpoint has this top-level tensor
layout:

```text
model: 454 tensors
text_encoders: 399 tensors
vae: 244 tensors
```

Representative diffusion keys:

```text
model.diffusion_model.cap_embedder.0.weight
model.diffusion_model.layers.0.attention.qkv.weight
model.diffusion_model.final_layer.linear.weight
```

Representative text encoder keys:

```text
text_encoders.qwen3_4b.transformer.model.embed_tokens.weight
text_encoders.qwen3_4b.transformer.model.layers.0.self_attn.q_proj.weight
text_encoders.qwen3_4b.transformer.model.layers.0.mlp.down_proj.weight
```

Representative VAE keys:

```text
vae.decoder.conv_in.weight
vae.decoder.mid.attn_1.q.weight
vae.decoder.conv_out.weight
```

This layout means the checkpoint already contains the same three
components that the program loads today. Each component is namespaced
with an extra prefix. Supporting this file is therefore mostly a
state-dict extraction and prefix-normalization problem.

## External References

The implementation should use the existing project dependencies:

- [`safetensors.safe_open`](https://huggingface.co/docs/safetensors/en/index)
  for reading safetensors headers and selected tensors.
- [`safetensors.torch.load_file`][safetensors-torch]
  only when loading a file that is already known to contain one
  component.
- Python [`pathlib.Path`](https://docs.python.org/3/library/pathlib.html)
  for path expansion and validation.
- Python [`dataclasses`](https://docs.python.org/3/library/dataclasses.html)
  for immutable source descriptions.
- Python [`tomllib`](https://docs.python.org/3/library/tomllib.html)
  for reading the existing TOML config file.

The reason to prefer `safe_open` for all-in-one checkpoints is that it
can inspect keys and load tensors by name. That matters because the
diffusion model, text encoder, and VAE are large. Loading the entire
all-in-one file just to use one component would waste memory and make
the staged loading design less effective.

## User Interface

### Command Line

Add a new option:

```text
--checkpoint PATH
```

The option means:

```text
Use PATH as an all-in-one safetensors checkpoint that may contain the
diffusion model, text encoder, and VAE.
```

Example:

```bash
AIO=~/documents/ai/diffusion/models/checkpoints/z-image_turbo
uv run diffusion-cli \
    --prompt "a ceramic mug on a desk" \
    --checkpoint "$AIO/redcraftRedzimageUpdatedDEC03_redzimage15AIO.safetensors"
```

The tokenizer path remains required through either CLI or config. The
checkpoint contains model tensors, not the tokenizer JSON and vocabulary
files used by `transformers`.

### Config File

Allow a new `[models]` key:

```toml
[models]
checkpoint = "/models/z-image/redcraftRedzimage15AIO.safetensors"
tokenizer = "/home/mw/programs/ComfyUI/comfy/text_encoders/qwen25_tokenizer"
```

The existing component keys remain valid:

```toml
[models]
diffusion_model = "/models/z-image/diffusion.safetensors"
text_encoder = "/models/z-image/qwen_3_4b.safetensors"
vae = "/models/z-image/ae.safetensors"
tokenizer = "/home/mw/programs/ComfyUI/comfy/text_encoders/qwen25_tokenizer"
```

Mixed configuration is also valid:

```toml
[models]
checkpoint = "/models/z-image/redcraftRedzimage15AIO.safetensors"
vae = "/models/z-image/alternate_ae.safetensors"
tokenizer = "/home/mw/programs/ComfyUI/comfy/text_encoders/qwen25_tokenizer"
```

In that example, diffusion and text encoder weights come from the
checkpoint, while the VAE comes from the separate `vae` file.

## Precedence

Source selection should follow this order for each model component:

1. Explicit component command line option.
2. Component path from config.
3. Explicit checkpoint command line option.
4. Checkpoint path from config.

For example:

```bash
uv run diffusion-cli \
    --prompt "a red apple" \
    --checkpoint aio.safetensors \
    --vae ae.safetensors
```

The effective sources are:

```text
diffusion model: aio.safetensors under model.diffusion_model.
text encoder: aio.safetensors under text_encoders.qwen3_4b.transformer.
VAE: ae.safetensors as a standalone VAE file
```

The reason component paths override checkpoint paths is that a component
path is more specific. If a user provides both, the specific path should
win without requiring a separate disable flag.

## Data Model

The current `ModelFiles` dataclass is not expressive enough because it
can only say "this component is in this file." It cannot say "this
component is inside this file under this prefix."

Replace or augment it with a source model.

### `ComponentRole`

Use constants or an enum-like string set for the three supported roles:

```python
DIFFUSION_ROLE = "diffusion_model"
TEXT_ENCODER_ROLE = "text_encoder"
VAE_ROLE = "vae"
```

The role is used for:

- Error messages.
- Prefix selection.
- Inspection summaries.
- Tests.

If an enum is used, its cases should follow the repository style for
global constants. That means enum cases should be `UPPER_CASE`.

### `ModelSource`

Represent one component source as:

```python
@dataclass(frozen=True)
class ModelSource:
    """A local safetensors source for one model component."""

    path: Path
    role: str
    checkpoint_prefix: str | None = None
```

Fields:

`path`

: Absolute path to a safetensors file. The file must exist and must be a
  regular file.

`role`

: The component role. This is not derived from the file name because
  local model files often have informal names.

`checkpoint_prefix`

: Prefix to select and strip when reading tensors from an all-in-one
  checkpoint. `None` means the file is a standalone component file.

Examples:

```python
ModelSource(
    path=Path("/models/aio.safetensors"),
    role="diffusion_model",
    checkpoint_prefix="model.diffusion_model.",
)
```

```python
ModelSource(
    path=Path("/models/qwen_3_4b.safetensors"),
    role="text_encoder",
    checkpoint_prefix=None,
)
```

### `ModelSources`

Represent the resolved generation sources as:

```python
@dataclass(frozen=True)
class ModelSources:
    """Resolved local sources for all model components."""

    diffusion_model: ModelSource
    text_encoder: ModelSource
    vae: ModelSource
```

The CLI should call one resolver and receive this object before
generation starts.

## Config Model Changes

Update `ModelPathConfig` to include an optional checkpoint:

```python
@dataclass(frozen=True)
class ModelPathConfig:
    """Optional model path defaults loaded from user configuration."""

    checkpoint: Path | None = None
    diffusion_model: Path | None = None
    text_encoder: Path | None = None
    vae: Path | None = None
    tokenizer: Path | None = None
```

Update `MODEL_CONFIG_KEYS` to include:

```python
"checkpoint"
```

Update `loadUserConfig()` so it reads:

```python
checkpoint=_optionalPath(models, "models", "checkpoint")
```

The config parser should continue rejecting unknown keys. That behavior
is useful because a misspelled key such as `checkpont` would otherwise
silently fall back to missing model paths.

## Path Resolution

Rename or replace `resolveModelFiles()` with a resolver that returns
`ModelSources`.

The resolver should:

1. Read CLI component paths.
2. Read config component paths.
3. Read CLI checkpoint path.
4. Read config checkpoint path.
5. Validate every chosen file path.
6. Build a source for each required component.
7. Raise a clear error if a component cannot be resolved.

Pseudo-code:

```python
def resolveModelSources(args, user_config: UserConfig) -> ModelSources:
    models = user_config.models
    checkpoint = args.checkpoint or models.checkpoint

    diffusion_model = resolveComponentSource(
        args.diffusion_model,
        models.diffusion_model,
        checkpoint,
        "diffusion_model",
        "model.diffusion_model.",
        "models.diffusion_model",
    )
    text_encoder = resolveComponentSource(
        args.text_encoder,
        models.text_encoder,
        checkpoint,
        "text_encoder",
        "text_encoders.qwen3_4b.transformer.",
        "models.text_encoder",
    )
    vae = resolveComponentSource(
        args.vae,
        models.vae,
        checkpoint,
        "vae",
        "vae.",
        "models.vae",
    )
    return ModelSources(diffusion_model, text_encoder, vae)
```

`resolveComponentSource()` should choose paths in this order:

1. CLI component path.
2. Config component path.
3. Checkpoint path.

If a component path is chosen, `checkpoint_prefix` is `None`.

If a checkpoint path is chosen, `checkpoint_prefix` is the component
prefix for that role.

If no source is available, the error should name both accepted config
styles:

```text
Missing models.diffusion_model or models.checkpoint
```

For command line users, argparse help already documents the matching
options. The config key in the error is still useful because it points
to the persistent fix.

## Selective Safetensors Loading

Add a small helper module or helper functions near model loading. A
separate module is preferable because all three component loaders need
the same behavior.

Suggested file:

```text
diffusion_cli/checkpoint.py
```

Suggested API:

```python
def loadStateDict(source: ModelSource) -> dict[str, torch.Tensor]:
    """Load tensors for one model component from a safetensors source."""
```

Behavior:

1. If `source.checkpoint_prefix is None`, call
   `safetensors.torch.load_file(source.path, device="cpu")`.
2. If `source.checkpoint_prefix` is not `None`, open the file with
   `safe_open(source.path, framework="pt", device="cpu")`.
3. Iterate over all keys in the file.
4. Select only keys that start with `source.checkpoint_prefix`.
5. Strip the prefix from each selected key.
6. Load each selected tensor with `tensors.get_tensor(key)`.
7. Return the stripped-key state dict.
8. If no keys matched, raise `DiffusionCliError`.

Pseudo-code:

```python
def loadStateDict(source: ModelSource) -> dict[str, torch.Tensor]:
    if source.checkpoint_prefix is None:
        return load_file(source.path, device="cpu")

    state_dict = {}
    with safe_open(source.path, framework="pt", device="cpu") as tensors:
        for key in tensors.keys():
            if key.startswith(source.checkpoint_prefix):
                stripped_key = key.removeprefix(source.checkpoint_prefix)
                state_dict[stripped_key] = tensors.get_tensor(key)

    if not state_dict:
        raise DiffusionCliError(
            f"Checkpoint does not contain {source.role}: {source.path}"
        )
    return state_dict
```

This helper deliberately returns a normal PyTorch state dict. That keeps
the existing model classes simple. They can continue doing architecture
detection, local key normalization, and `load_state_dict()` checks.

## Component Prefixes

The all-in-one prefixes should be centralized so tests and inspection
use the same values as generation.

Suggested constants:

```python
CHECKPOINT_DIFFUSION_PREFIX = "model.diffusion_model."
CHECKPOINT_TEXT_ENCODER_PREFIX = "text_encoders.qwen3_4b.transformer."
CHECKPOINT_VAE_PREFIX = "vae."
```

After these prefixes are stripped, the resulting keys match the
standalone component conventions closely:

```text
cap_embedder.0.weight
model.embed_tokens.weight
decoder.conv_in.weight
```

The text encoder still needs the existing Qwen normalization because it
expects Transformers model keys without the leading `model.` prefix.
The current `normalizeQwenStateDict()` already handles that conversion.

The VAE still needs its existing decoder extraction because it expects
to load a `VaeDecoder`, not a full autoencoder object. After stripping
`vae.`, the keys still begin with `decoder.`, so the existing decoder
filter can remain in place.

## Loader Changes

### Text Encoder

Change `ZImageTextEncoder` to accept a `ModelSource` instead of a
plain `Path`, or add an overload-style internal path that accepts either
source or path.

Preferred constructor:

```python
class ZImageTextEncoder:
    """Load the local Qwen3-4B encoder and produce text conditioning."""

    def __init__(
        self,
        model_source: ModelSource,
        tokenizer_path: Path,
        device,
        dtype,
        *,
        model=None,
        tokenizer=None,
    ) -> None:
```

Then `_loadModel()` should call:

```python
state_dict = normalizeQwenStateDict(loadStateDict(model_source))
```

The rest of `_loadModel()` can stay the same.

### Diffusion Model

Change `ZImageModel` in the same way. `_loadModel()` should call:

```python
state_dict = normalizeZImageStateDict(loadStateDict(model_source))
```

For the example all-in-one checkpoint, `loadStateDict()` will already
strip `model.diffusion_model.`, so `normalizeZImageStateDict()` may not
need to strip anything. Keeping the normalization call is still correct
because it protects the standalone component path and preserves current
behavior.

### VAE

Change `ZImageVae` in the same way. The VAE loader should call:

```python
state_dict = loadStateDict(model_source)
```

Then the existing code can:

1. Detect the VAE config from `decoder.conv_in.weight`.
2. Strip `decoder.` for the decoder module.
3. Load the decoder state dict.

## CLI Generation Flow

`generate()` should keep its current sequence:

1. Build generation config.
2. Resolve model sources.
3. Load text encoder.
4. Encode prompts.
5. Release memory.
6. Load diffusion model.
7. Sample latents.
8. Release memory.
9. Load VAE.
10. Decode and save images.

Only step 2 and constructor arguments change.

Conceptually:

```python
config = buildGenerationConfig(args, user_config)
model_sources = resolveModelSources(args, user_config)

text_encoder = ZImageTextEncoder(
    model_sources.text_encoder,
    config.tokenizer_path,
    config.device,
    config.dtype,
)
```

This preserves the main memory-management property of the current CLI.
Even when all components live in one file, the program only materializes
the active stage's tensors.

## Inspection

`--inspect-models` currently prints one summary per component path. With
checkpoint support, it should support both separate files and all-in-one
sources.

For standalone component paths, the current output is still fine.

For checkpoint-backed sources, the output should make the prefix clear:

```text
diffusion model: /models/aio.safetensors
  source prefix: model.diffusion_model.
  tensors: 454
  dtypes: BF16: 454
  architecture guess: z_image_diffusion
```

The inspection helper should avoid loading tensor data. It should use
`safe_open()` and `get_slice()` just like the current
`inspectSafetensors()` implementation.

A practical approach is to add:

```python
def inspectModelSource(source: ModelSource) -> TensorSummary:
    """Read metadata for one resolved model source."""
```

When `source.checkpoint_prefix is None`, this can behave like the
current `inspectSafetensors()`.

When `source.checkpoint_prefix` is set, it should:

1. Open `source.path`.
2. Select only matching keys.
3. Strip the prefix before architecture guessing.
4. Count only matching tensors.
5. Raise `DiffusionCliError` if no matching keys exist.

The existing `guessArchitecture()` should be updated for stripped AIO
keys if needed. For example, after stripping the text encoder prefix,
`model.embed_tokens.weight` is already recognized as
`qwen_text_encoder`.

## Error Handling

Use `DiffusionCliError` for user-fixable problems.

### Missing Checkpoint File

If a selected checkpoint path does not exist:

```text
Missing checkpoint: /path/to/model.safetensors
```

### Missing Component in Checkpoint

If a checkpoint exists but has no matching tensors for a role:

```text
Checkpoint does not contain text encoder: /path/to/model.safetensors
```

This should happen before constructing the heavy model object. A missing
component is a source-selection problem, not an architecture problem.

### Unsupported Component Keys

If matching tensors exist but do not match the expected architecture,
keep the existing component errors:

```text
Unsupported Z-Image checkpoint: missing Lumina2 keys
Unsupported VAE state dict: no known decoder keys found
Could not load Qwen text encoder: missing keys: [...]
```

These errors mean the selected source had tensors for that component
role, but they were not the supported Z-Image/Qwen/VAE layout.

### Conflicting Inputs

Providing both `--checkpoint` and `--diffusion-model` should not be an
error. The component path should win for that component. This makes the
CLI easy to use for experiments where only one component is swapped.

## Testing Strategy

The tests should avoid real model weights. Small synthetic safetensors
files are enough for path resolution, prefix selection, and metadata
inspection. For loader-level tests, existing unit tests can cover
normalization functions without allocating real model sizes.

### Config Tests

Add tests that verify:

- `[models].checkpoint` is accepted.
- `ModelPathConfig.checkpoint` is populated.
- Unknown model keys are still rejected.
- Non-string `checkpoint` values fail with a clear type error.

Example assertion:

```python
self.assertEqual(config.models.checkpoint.name, "aio.safetensors")
```

### CLI Parser Tests

Add parser assertions:

- `args.checkpoint is None` when not provided.
- `--checkpoint aio.safetensors` parses to a `Path`.

The existing parser test that configurable defaults stay unset should
include `checkpoint`.

### Path Resolution Tests

Add tests for these cases:

- Separate component paths still resolve.
- Checkpoint-only config resolves all three roles to the same path with
  different prefixes.
- CLI checkpoint overrides config checkpoint.
- CLI component path overrides checkpoint for one role.
- Missing component and missing checkpoint produce clear errors.

A mixed-source assertion should check both path and prefix:

```python
self.assertIsNone(sources.vae.checkpoint_prefix)
self.assertEqual(
    sources.diffusion_model.checkpoint_prefix,
    CHECKPOINT_DIFFUSION_PREFIX,
)
```

### Checkpoint Helper Tests

Create a tiny safetensors file with keys like:

```text
model.diffusion_model.cap_embedder.1.weight
text_encoders.qwen3_4b.transformer.model.embed_tokens.weight
vae.decoder.conv_in.weight
unrelated.weight
```

Verify:

- Loading the diffusion source returns only diffusion keys.
- Returned keys have the checkpoint prefix stripped.
- Unrelated keys are ignored.
- Missing prefixes raise `DiffusionCliError`.

Use tiny tensors such as `torch.zeros((1, 1))` so tests are fast.

### Normalization Tests

Extend existing normalization tests:

- `normalizeQwenStateDict()` should still turn
  `model.embed_tokens.weight` into `embed_tokens.weight`.
- VAE prefix stripping should leave `decoder.conv_in.weight` for config
  detection.
- Diffusion prefix stripping should leave `cap_embedder.1.weight` and
  `x_embedder.weight` for config detection.

### Inspection Tests

Add tests that inspect a synthetic all-in-one file and verify:

- The selected tensor count excludes unrelated keys.
- Top-level counts are based on stripped keys.
- Architecture guessing uses stripped keys.

## Implementation Plan

1. Add checkpoint fields to config dataclasses and config parsing.
2. Add `--checkpoint` to the CLI parser.
3. Introduce `ModelSource` and `ModelSources`.
4. Replace `resolveModelFiles()` with `resolveModelSources()`, or keep
   a compatibility wrapper only if tests or callers still need it.
5. Add checkpoint prefix constants.
6. Add a selective safetensors loading helper.
7. Update `ZImageTextEncoder`, `ZImageModel`, and `ZImageVae` to load
   from `ModelSource`.
8. Update `generate()` to pass resolved sources.
9. Update inspection to inspect resolved sources.
10. Add unit tests for config, parser, path resolution, prefix loading,
    and inspection.
11. Run the unit test suite.
12. If local GPU memory allows, run a small real generation smoke test
    with the example all-in-one checkpoint.

## Backward Compatibility

Existing commands that specify separate component files should continue
to work unchanged:

```bash
uv run diffusion-cli \
    --prompt "a mug" \
    --diffusion-model diffusion.safetensors \
    --text-encoder qwen_3_4b.safetensors \
    --vae ae.safetensors \
    --tokenizer-path qwen25_tokenizer
```

Existing config files should continue to work unchanged because the new
`checkpoint` key is optional.

The only intentional config behavior change is that `checkpoint` is no
longer rejected as an unknown `[models]` key.

## Risks

### Memory Use

The largest risk is accidentally loading the entire all-in-one
checkpoint for every component. That would undermine the current staged
memory release behavior. The selective loader addresses this by reading
only matching tensors.

### Prefix Variants

Other all-in-one checkpoints may use different prefixes. The first
implementation should support the observed Z-Image prefix layout only.
If another local checkpoint uses a different layout, the error should be
clear rather than trying to guess too broadly.

A later feature could add auto-detection or configurable prefixes. That
should be deferred until there are at least two real layouts to support.

### Duplicate Component Names

A future checkpoint could contain multiple text encoders under
`text_encoders.*`. The current feature should target
`text_encoders.qwen3_4b.transformer.` specifically because the current
Z-Image implementation only supports Qwen3-4B.

### Silent Wrong Component

Broad prefix matching could accidentally select the wrong tensors. The
prefixes should therefore be exact and role-specific. Architecture
detection and `load_state_dict()` validation provide a second layer of
protection.

## Open Questions

The first implementation can proceed without resolving these questions:

- Should `--inspect-models` show one combined checkpoint overview before
  per-component summaries?
- Should the config key be named `checkpoint` or `aio_checkpoint`?
- Should future releases support user-configurable prefixes?

Recommended answers for the first implementation:

- Show per-component summaries only, because that maps directly to the
  resolved sources used for generation.
- Use `checkpoint`, because this matches common diffusion terminology
  and keeps the CLI concise.
- Do not support configurable prefixes yet, because the project has one
  known Z-Image layout and exact prefixes are safer.

[safetensors-torch]: https://huggingface.co/docs/safetensors/en/api/torch