Changes
diff --git a/designs/design-1-config.md b/designs/design-1-config.md
new file mode 100644
index 0000000..7f04a72
--- /dev/null
+++ b/designs/design-1-config.md
@@ -0,0 +1,546 @@
+# Config File Design
+
+## Purpose
+
+This document describes how `diffusion-cli` should support a user
+configuration file.
+
+The configuration file gives the command stable defaults for local model
+files and generation parameters. This matters because the project runs
+against locally stored model files and should not assume that every
+machine uses the same directory layout. A single model root plus fixed
+relative paths is too rigid. Each required model file should be named
+independently.
+
+The command line interface remains the highest-priority source of user
+intent. The config file provides defaults, but explicit command line
+options override those defaults for the current run.
+
+## Requirements
+
+The config file path is fixed:
+
+```text
+~/.config/diffusion.toml
+```
+
+The file uses TOML. The project requires Python 3.11 or newer, so the
+implementation can use Python's standard `tomllib` module to read TOML
+without adding a dependency. `tomllib` only reads TOML; it does not
+write TOML. That is enough for this feature because the CLI only needs
+to load user-provided defaults. See the
+[tomllib documentation](https://docs.python.org/3/library/tomllib.html).
+
+The config file should contain:
+
+- Absolute or user-expanded paths for each model file.
+- A tokenizer directory path.
+- Default generation parameters such as width, height, steps, CFG,
+ negative prompt, output path, device, and dtype.
+
+The config file should not contain:
+
+- A model root used to build relative model file paths.
+- Any rule that scans directories to discover model files.
+- Any requirement that model files live under the same parent directory.
+
+The command should still expose command line options for configurable
+generation values. When a value appears in both the config file and the
+command line, the command line value wins.
+
+## Current State
+
+The current path resolver is centered on a model root:
+
+```text
+diffusion_cli/paths.py
+```
+
+It defines a default model root and joins fixed relative paths under
+that root. The CLI exposes `--model-root`, and the resolver also checks
+`DIFFUSION_CLI_MODEL_ROOT`.
+
+That behavior should be replaced for model files. The new behavior
+should treat the diffusion model, text encoder, VAE, and tokenizer as
+four independent paths.
+
+The current generation defaults are defined as constants in:
+
+```text
+diffusion_cli/config.py
+```
+
+Those constants should remain as built-in fallback values. The config
+file should override the built-in fallbacks, and command line options
+should override both.
+
+## Config Format
+
+The file should use two top-level tables:
+
+- `[models]` for local model and tokenizer paths.
+- `[generation]` for generation defaults.
+
+Example:
+
+```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"
+
+[generation]
+negative_prompt = "text, watermark, full-body"
+width = 832
+height = 1248
+batch_size = 1
+steps = 10
+cfg = 1.0
+output = "output.png"
+device = "cuda"
+dtype = "auto"
+```
+
+### `[models]`
+
+`diffusion_model`
+
+: Path to the Z-Image diffusion model safetensors file. This value
+ replaces the old root-plus-relative path behavior for the diffusion
+ model.
+
+`text_encoder`
+
+: Path to the Qwen text encoder safetensors file.
+
+`vae`
+
+: Path to the VAE safetensors file.
+
+`tokenizer`
+
+: Path to the local tokenizer directory used by `transformers`.
+
+Every model path should be parsed as a `pathlib.Path`, expanded with
+`expanduser()`, resolved, and validated before generation starts. Model
+files must exist and be regular files. The tokenizer path must exist,
+must be a directory, and must contain the required tokenizer files.
+
+### `[generation]`
+
+`negative_prompt`
+
+: Default negative prompt. This is used when `--negative-prompt` is not
+ passed.
+
+`width`
+
+: Default output width. It must be a positive multiple of the latent
+ downscale, currently 8.
+
+`height`
+
+: Default output height. It must be a positive multiple of the latent
+ downscale, currently 8.
+
+`batch_size`
+
+: Default number of images to generate. It must be at least 1.
+
+`steps`
+
+: Default number of denoising steps. It must be at least 1.
+
+`cfg`
+
+: Default classifier-free guidance scale. It must be non-negative.
+
+`output`
+
+: Default output image path. If the parent directory does not exist, the
+ CLI should create it when possible, matching the current behavior.
+
+`device`
+
+: Default inference device string. The current milestone is CUDA-only,
+ so non-CUDA values should fail with a clear unsupported-device error.
+
+`dtype`
+
+: Default inference dtype string. Allowed values should remain `auto`,
+ `bf16`, and `fp32`.
+
+## Values That Should Stay CLI-First
+
+`prompt` should stay CLI-first and should not be required in the config
+file.
+
+The prompt is usually different for every generation. Putting it in the
+config file would make repeated invocations easier to run accidentally
+with stale prompt text. The existing behavior, where generation fails
+without a prompt, is clearer.
+
+`seed` may remain CLI-only with a random default when omitted. A seed is
+part of a specific generation request more often than it is a global
+user preference. Keeping it CLI-only avoids accidentally making every
+generation deterministic because a seed was left in the config file.
+
+If deterministic default behavior becomes useful later, `seed` can be
+added to `[generation]` without changing the rest of the merge model.
+
+## Precedence
+
+Effective values should be selected in this order:
+
+1. Command line option.
+2. Config file value.
+3. Built-in default.
+
+This rule is intentionally simple. The command line is closest to the
+current request, so it should win. The config file represents persistent
+user preferences, so it should beat built-in project defaults. Built-in
+defaults are the last fallback so the command remains usable when the
+config file is absent or incomplete.
+
+For example, if the config contains:
+
+```toml
+[generation]
+steps = 20
+```
+
+and the user runs:
+
+```bash
+uv run diffusion-cli --prompt "a mug" --steps 12
+```
+
+the effective step count is `12`.
+
+If the user runs:
+
+```bash
+uv run diffusion-cli --prompt "a mug"
+```
+
+the effective step count is `20`.
+
+If there is no config file and no `--steps` option, the effective step
+count is the built-in default from `diffusion_cli/config.py`.
+
+## Argparse Handling
+
+The implementation must be careful with `argparse` defaults.
+
+Today several options are registered with defaults directly in the
+parser. For config merging, those parser defaults should become `None`
+for values that can come from the config file. If argparse injects the
+built-in value before the merge step, the program cannot tell whether
+the user explicitly supplied the option or whether argparse supplied the
+default.
+
+For configurable values, parser setup should look conceptually like:
+
+```python
+parser.add_argument("--width", type=int)
+parser.add_argument("--height", type=int)
+parser.add_argument("--steps", type=int)
+parser.add_argument("--cfg", type=float)
+```
+
+The merge layer then fills missing values:
+
+```python
+value = cli_value
+if value is None:
+ value = config_value
+if value is None:
+ value = built_in_default
+```
+
+This keeps command line precedence precise.
+
+Options that are not generation defaults can keep normal argparse
+behavior. For example, `--inspect-models` is a command mode flag, so it
+can remain a normal boolean flag with `False` as the default.
+
+## Loading Flow
+
+The CLI should load and merge configuration before validation that needs
+effective values.
+
+Recommended flow:
+
+1. Build the parser.
+2. Parse command line arguments.
+3. Load `~/.config/diffusion.toml` if it exists.
+4. Merge command line values, config values, and built-in defaults.
+5. Validate the merged model paths.
+6. Validate the merged generation parameters.
+7. Build immutable `ModelFiles` and `GenerationConfig` objects.
+8. Continue into inspection or generation.
+
+The config file should be optional. If it does not exist, the command
+should continue with built-in defaults and command line values.
+
+If the file exists but cannot be parsed, the command should fail
+clearly. A malformed config file represents user intent that the program
+cannot safely interpret.
+
+## Proposed Code Structure
+
+The feature can be implemented with small changes to existing modules.
+
+### `config.py`
+
+`config.py` should continue to own immutable value objects and
+validation helpers.
+
+Add a config-file representation such as:
+
+```python
+@dataclass(frozen=True)
+class FileModelConfig:
+ diffusion_model: Path | None = None
+ text_encoder: Path | None = None
+ vae: Path | None = None
+ tokenizer: Path | None = None
+
+
+@dataclass(frozen=True)
+class FileGenerationConfig:
+ negative_prompt: str | None = None
+ width: int | None = None
+ height: int | None = None
+ batch_size: int | None = None
+ steps: int | None = None
+ cfg: float | None = None
+ output: Path | None = None
+ device: str | None = None
+ dtype: str | None = None
+
+
+@dataclass(frozen=True)
+class FileConfig:
+ models: FileModelConfig
+ generation: FileGenerationConfig
+```
+
+The file config objects are intentionally different from
+`GenerationConfig`. File config objects represent optional raw defaults.
+`GenerationConfig` represents validated effective generation intent.
+
+Add a constant:
+
+```python
+DEFAULT_CONFIG_PATH = Path("~/.config/diffusion.toml")
+```
+
+The loader should expand this path before reading:
+
+```python
+DEFAULT_CONFIG_PATH.expanduser()
+```
+
+Add a loader:
+
+```python
+def loadFileConfig(path: Path = DEFAULT_CONFIG_PATH) -> FileConfig:
+ ...
+```
+
+The loader should return an empty `FileConfig` if the path does not
+exist. It should raise `DiffusionCliError` if the file exists but is not
+valid TOML or has invalid top-level structure.
+
+### `paths.py`
+
+`paths.py` should stop exposing root-based model resolution as the main
+path. It can keep `requireFile()` because that helper is still useful.
+
+Replace:
+
+```python
+resolveModelFiles(model_root: Path | None = None)
+```
+
+with a function that accepts already merged path values:
+
+```python
+def validateModelFiles(
+ diffusion_model: Path,
+ text_encoder: Path,
+ vae: Path,
+) -> ModelFiles:
+ ...
+```
+
+This function should validate exactly the paths it receives. It should
+not join them to a root.
+
+### `cli.py`
+
+`cli.py` should:
+
+- Remove `--model-root`.
+- Keep command line options for generation defaults.
+- Keep `--tokenizer-path` as a command line override for the tokenizer
+ directory.
+- Parse configurable options with `None` defaults.
+- Load the fixed config file path.
+- Merge values before building `GenerationConfig`.
+- Use merged model paths for model inspection and generation.
+
+The `--tokenizer-path` name can remain even though the config key is
+`models.tokenizer`. The CLI option name is already clear to users, and
+the config key belongs with the other model-loading paths.
+
+## Error Handling
+
+Missing config file:
+
+: Not an error. Continue with command line values and built-in defaults.
+
+Malformed TOML:
+
+: Error. Include the config path and TOML parse error message.
+
+Unknown top-level table:
+
+: Error. This catches misspellings such as `[model]` instead of
+ `[models]`.
+
+Unknown key in a known table:
+
+: Error. This catches misspellings such as `step = 10` instead of
+ `steps = 10`.
+
+Wrong value type:
+
+: Error. For example, `steps = "10"` should fail because `steps` must be
+ an integer.
+
+Missing model path after merge:
+
+: Error. Name the missing key, such as `models.diffusion_model`.
+
+Missing tokenizer path after merge:
+
+: Error. Name the missing key or selected path. Tokenizer validation
+ should continue to report missing tokenizer files precisely.
+
+Invalid generation value:
+
+: Error. Reuse the existing validation messages where possible.
+
+Unknown keys should be treated as errors instead of ignored. Ignoring
+unknown keys makes config mistakes hard to find because the command
+silently falls back to a different value than the user expected.
+
+## Backward Compatibility
+
+This feature intentionally removes the root-based model path model from
+the preferred interface.
+
+Recommended compatibility behavior:
+
+- Remove `--model-root` from the parser.
+- Remove `DIFFUSION_CLI_MODEL_ROOT` from model resolution.
+- Keep the built-in concrete model paths if the project should remain
+ runnable on the original development machine without a config file.
+
+The last point is a product choice. There are two reasonable options:
+
+1. Keep built-in concrete model defaults matching the current machine.
+2. Require model paths in `~/.config/diffusion.toml`.
+
+Option 1 is more convenient for existing local usage. Option 2 is more
+portable and makes local machine assumptions explicit. If the goal is to
+avoid hidden path assumptions entirely, choose option 2 and make missing
+model config fail clearly.
+
+This design recommends option 2. Built-in generation defaults are stable
+workflow defaults, but built-in model file paths are local machine
+defaults. Local machine paths should live in the user config file.
+
+## Tests
+
+Add or update tests for these behaviors.
+
+### Config Loading
+
+- Missing `~/.config/diffusion.toml` produces an empty file config.
+- Valid TOML loads `[models]` paths.
+- Valid TOML loads `[generation]` values.
+- Malformed TOML raises `DiffusionCliError`.
+- Unknown top-level table raises `DiffusionCliError`.
+- Unknown key in `[models]` raises `DiffusionCliError`.
+- Unknown key in `[generation]` raises `DiffusionCliError`.
+- Wrong scalar types raise `DiffusionCliError`.
+
+### Merge Precedence
+
+- CLI value overrides config value.
+- Config value overrides built-in default.
+- Built-in default is used when both CLI and config omit a generation
+ value.
+- Missing prompt still fails for generation.
+- Missing model path after merge fails with the specific config key.
+
+### Path Validation
+
+- Independent model paths are accepted even when they are in unrelated
+ directories.
+- Missing diffusion model reports `Missing diffusion model`.
+- Missing text encoder reports `Missing text encoder`.
+- Missing VAE reports `Missing VAE`.
+- Tokenizer directory validation still checks required tokenizer files.
+
+### CLI Parser
+
+- Configurable arguments parse to `None` when omitted.
+- Explicit CLI arguments parse to their supplied values.
+- `--inspect-models` remains a boolean command flag.
+- `--model-root` is no longer accepted.
+
+## Documentation Updates
+
+Update `README.md` to describe:
+
+- The fixed config path: `~/.config/diffusion.toml`.
+- The `[models]` and `[generation]` tables.
+- The command line override rule.
+- The removal of `--model-root`.
+
+The README should include a small complete TOML example and at least one
+override example:
+
+```bash
+uv run diffusion-cli --prompt "a ceramic mug" --steps 12
+```
+
+The explanation should state that `--steps 12` overrides any
+`generation.steps` value in the config file for that run only.
+
+## Open Decisions
+
+`seed`
+
+: This design keeps `seed` CLI-only. Add it to `[generation]` later only
+ if deterministic default runs become a real workflow requirement.
+
+Built-in model paths
+
+: This design recommends requiring model paths from
+ `~/.config/diffusion.toml`. If preserving out-of-the-box behavior on
+ the current development machine is more important, the implementation
+ can keep built-in concrete model file defaults. It should still avoid
+ root-plus-relative path joining.
+
+Environment variables
+
+: This design does not include environment variable overrides. They can
+ be useful later, but adding them now creates a fourth precedence layer.
+ The requested behavior only needs built-ins, config, and command line
+ options.
diff --git a/designs/design-2-checkpoint.md b/designs/design-2-checkpoint.md
new file mode 100644
index 0000000..db8db8b
--- /dev/null
+++ b/designs/design-2-checkpoint.md
@@ -0,0 +1,897 @@
+# 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