# Standalone Z-Image Turbo CLI Design

## Purpose

This document describes the first implementation milestone for this
repository: a standalone Python command line program that generates one
non-gibberish image from local Z-Image Turbo weights.

The goal is intentionally narrower than full ComfyUI parity. The first
working version only needs to prove that the project can load the local
model stack, run text-to-image inference without importing ComfyUI at
runtime, decode the latent, and write a PNG. LoRA support is explicitly
out of scope for this milestone.

The ComfyUI workflow is still the reference for model filenames, image
size, sampler defaults, prompt handling, and generation parameters. The
runtime program must not depend on ComfyUI as a package or import from
ComfyUI modules during generation.

## Requirements

The product requirements in `prd.md` say the tool must:

- Be a Python CLI.
- Generate images using the local Z-Image Turbo model stack referenced
  by `~/programs/ComfyUI/user/default/workflows/Z-image Turbo.json`.
- Avoid any runtime dependency on ComfyUI.
- Use the ComfyUI workflow only as a reference.
- Avoid downloading model files or other files from the internet.
- Assume model files are already local.
- Use `uv` for dependency and virtual environment management.
- Support prompt, negative prompt, seed, width, height, batch size,
  steps, CFG, output path, and device options.
- First prove one image can be generated from local weights without
  importing ComfyUI.

The current user-level clarification narrows the milestone:

- The first output only needs to be an image that is not gibberish.
- LoRA support should be skipped for now.
- ComfyUI source should be used as implementation reference material.

## Reference Workflow

The workflow file is:

`/home/mw/programs/ComfyUI/user/default/workflows/Z-image Turbo.json`

The active path in that graph uses these nodes:

- `UNETLoader`
- `CLIPLoader`
- `VAELoader`
- `ModelSamplingAuraFlow`
- `CLIPTextEncode`
- `EmptyLatentImage`
- `KSampler`
- `VAEDecode`
- `SaveImage`

The workflow also contains LoRA nodes, but this design ignores them for
the first milestone.

### Model Files

ComfyUI resolves models through:

`/home/mw/programs/ComfyUI/extra_model_paths.yaml`

That file maps the default model root to:

`/home/mw/documents/ai/diffusion/models`

For the first milestone, the CLI should use this model root by default
and allow it to be overridden with a command line option.

Required files:

| Role | Relative Path |
| --- | --- |
| Diffusion model | `diffusion_models/z-image_turbo/moodyPornMix_zitV3.safetensors` |
| Text encoder | `text_encoders/z-image_turbo/qwen_3_4b.safetensors` |
| VAE | `vae/z-image_turbo/ae.safetensors` |

The workflow also references an all-in-one checkpoint:

`checkpoints/z-image_turbo/redcraftRedzimageUpdatedDEC03_redzimage15AIO.safetensors`

That checkpoint node is not connected to the active generation path, so
the CLI should not use it in the first milestone.

### Generation Defaults

The first CLI defaults should match the active workflow:

| Setting | Value |
| --- | --- |
| Width | `832` |
| Height | `1248` |
| Batch size | `1` |
| Steps | `10` |
| CFG | `1.0` |
| Sampler | `seeds_3` |
| Scheduler | `beta` |
| Denoise | `1.0` |
| Model sampling shift | `3.0` |
| Positive prompt | User supplied |
| Negative prompt | `text, watermark, full-body` |

## Important ComfyUI Findings

This section records the ComfyUI source behavior that matters for the
standalone implementation.

### Z-Image Model Configuration

ComfyUI defines `ZImage` in:

`/home/mw/programs/ComfyUI/comfy/supported_models.py`

The relevant behavior is:

- Z-Image is a `Lumina2` model variant.
- Its `unet_config` contains `image_model = "lumina2"` and `dim = 3840`.
- Its sampling settings are `multiplier = 1.0` and `shift = 3.0`.
- Its supported inference dtypes are `bfloat16` and `float32`, with
  optional `float16` only when ComfyUI detects extended fp16 support.
- It uses the same latent format as `Lumina2`, which is `Flux`.

The standalone implementation should run on CUDA only for the first
milestone. It should start with `bfloat16` if the selected CUDA device
supports it, then fall back to `float32` if needed. This choice is
conservative because ComfyUI lists bf16/fp32 as supported.

### Text Encoder

ComfyUI defines Z-Image text encoding in:

`/home/mw/programs/ComfyUI/comfy/text_encoders/z_image.py`

The relevant behavior is:

- The tokenizer is based on `Qwen2Tokenizer` from `transformers`.
- The text encoder model is `Qwen3_4B`.
- The prompt is wrapped with a chat-style template before tokenization:

```text
<|im_start|>user
{prompt}<|im_end|>
<|im_start|>assistant
```

- Token weighting is disabled for this tokenizer path.
- The model returns attention masks.
- The default selected hidden layer is `-2`.

The standalone implementation must reproduce the prompt template. If it
does not, the model can receive text embeddings in a distribution it was
not trained to use, which is a common cause of visually incoherent
images.

### Latent Format

ComfyUI's `Lumina2` model uses `latent_formats.Flux`.

That means:

- The diffusion latent has 16 channels.
- The spatial downscale ratio is 8.
- A latent for `832x1248` should have shape:

```text
[batch_size, 16, 156, 104]
```

The order is `[N, C, H, W]`. The latent height is image height divided
by 8. The latent width is image width divided by 8.

The `Flux` latent format also applies scaling when entering and leaving
the diffusion model. The implementation must port the exact scale,
shift, or mean/std logic from ComfyUI's `latent_formats.Flux` class
rather than assuming vanilla Stable Diffusion VAE scaling.

### Model Sampling

The workflow node `ModelSamplingAuraFlow` calls the same patch method as
ComfyUI's SD3-style model sampling node, but with:

```text
shift = 3.0
multiplier = 1.0
```

The implementation maps to ComfyUI's `ModelSamplingDiscreteFlow` in:

`/home/mw/programs/ComfyUI/comfy/model_sampling.py`

The important methods are:

```python
def time_snr_shift(alpha, t):
    if alpha == 1.0:
        return t
    return alpha * t / (1 + (alpha - 1) * t)

def sigma(timestep):
    return time_snr_shift(shift, timestep / multiplier)

def timestep(sigma):
    return sigma * multiplier
```

For this workflow, `multiplier` is `1.0`, so the sigma schedule is a
flow schedule in the `[0, 1]` range shifted by `shift = 3.0`.

### Scheduler

The workflow scheduler is `beta`.

ComfyUI implements this in:

`/home/mw/programs/ComfyUI/comfy/samplers.py`

Conceptually:

1. Build the model sampling object.
2. Read its `sigmas` table.
3. Generate `steps` values using a beta distribution percent point
   function with alpha `0.6` and beta `0.6`.
4. Convert those values into timestep indices.
5. Deduplicate adjacent repeated timesteps.
6. Append a final zero sigma.

This requires `scipy.stats.beta.ppf`, so `scipy` should be a dependency
unless the beta PPF is reimplemented. Using SciPy is simpler and less
risky for the first milestone.

### Sampler

The workflow sampler is `seeds_3`.

ComfyUI implements `sample_seeds_3` in:

`/home/mw/programs/ComfyUI/comfy/k_diffusion/sampling.py`

SEEDS-3 is a three-stage stochastic solver. At each sigma interval, it:

1. Calls the denoising model at the current latent.
2. Converts sigma to half-log-SNR.
3. Builds two intermediate points at one-third and two-thirds through
   the interval.
4. Calls the denoising model at each intermediate point.
5. Combines the denoised predictions into the next latent.
6. Optionally injects stochastic noise when `eta` and `s_noise` are
   positive.

The sampler depends on model-specific conversions between sigma and
half-log-SNR. Those conversions need to be ported together with the
sampler math. Porting only the visible loop is not enough.

For the first milestone, the implementation may expose `--sampler
seeds_3` as the default but keep the sampler code private and minimal.
It does not need every sampler ComfyUI supports.

### VAE Decode

ComfyUI's `VAE` wrapper is in:

`/home/mw/programs/ComfyUI/comfy/sd.py`

The active VAE file is:

`vae/z-image_turbo/ae.safetensors`

The VAE loader detects the VAE architecture from state dict keys. For
Z-Image Turbo, the first implementation should inspect the `ae` state
dict and choose the exact matching VAE class or a compatible local
implementation. It should not assume SD 1.x VAE layout unless the keys
prove that layout.

The decode output should be converted from model output range to image
range, then written as a PNG:

1. Model output tensor is expected in `[-1, 1]`.
2. Convert to `[0, 1]` with `(image + 1) / 2`.
3. Clamp to `[0, 1]`.
4. Convert to `uint8`.
5. Save with Pillow.

This is the default ComfyUI postprocess path for image VAEs.

## Architecture

The CLI should be organized as a small Python package. The package
should isolate command parsing, configuration, model loading, sampling,
and image saving so that each part can be tested independently.

Proposed file layout:

```text
diffusion_cli/
    __init__.py
    cli.py
    config.py
    errors.py
    image_io.py
    paths.py
    sampling.py
    text_encoder.py
    zimage_model.py
    vae.py
tests/
    test_paths.py
    test_sampling.py
    test_cli.py
pyproject.toml
README.md
```

### `cli.py`

This module owns the command line interface.

It should use Python's standard `argparse` library. `argparse` is part
of the Python standard library and is designed for command line option
parsing, so it avoids adding a CLI framework dependency. See the
[argparse documentation](https://docs.python.org/3/library/argparse.html).

The command should be exposed through a `pyproject.toml` console script:

```toml
[project.scripts]
diffusion-cli = "diffusion_cli.cli:main"
```

The CLI should accept:

| Option | Type | Default | Meaning |
| --- | --- | --- | --- |
| `--prompt` | string | required | Positive prompt. |
| `--negative-prompt` | string | workflow default | Negative prompt. |
| `--seed` | int | random if absent | Noise seed. |
| `--width` | int | `832` | Output width. |
| `--height` | int | `1248` | Output height. |
| `--batch-size` | int | `1` | Number of images. |
| `--steps` | int | `10` | Denoising steps. |
| `--cfg` | float | `1.0` | Classifier-free guidance scale. |
| `--output` | path | `output.png` | Output file path. |
| `--device` | string | `cuda` | CUDA device id or alias. |
| `--model-root` | path | local model root | Model root directory. |
| `--tokenizer-path` | path | required | Qwen tokenizer directory. |
| `--dtype` | string | `auto` | `auto`, `bf16`, `fp32`. |

Validation rules:

- `width` and `height` must be positive multiples of 8.
- `batch-size` must be at least 1.
- `steps` must be at least 1.
- `cfg` must be non-negative.
- `output` parent directory must exist or be creatable.
- The first milestone is CUDA-only. CPU execution should fail with a
  clear unsupported-device error rather than attempting a very slow run.
- `device=cuda` must fail clearly if CUDA is unavailable.
- `dtype=bf16` must fail clearly if unsupported on the selected device.
- `tokenizer-path` must exist and contain the tokenizer files required
  by `transformers`, including `vocab.json`, `merges.txt`, and
  `tokenizer_config.json`.

### `config.py`

This module should hold immutable configuration objects.

Use dataclasses:

```python
@dataclass(frozen=True)
class ModelFiles:
    diffusion_model: Path
    text_encoder: Path
    vae: Path

@dataclass(frozen=True)
class GenerationConfig:
    prompt: str
    negative_prompt: str
    seed: int
    width: int
    height: int
    batch_size: int
    steps: int
    cfg: float
    device: torch.device
    dtype: torch.dtype
    output: Path
```

The value object should not load models. It only records the user's
intent after validation. Keeping configuration separate from loading
makes errors easier to test.

### `paths.py`

This module should resolve model paths.

Default model root:

```text
/home/mw/documents/ai/diffusion/models
```

Resolution flow:

1. Read `--model-root` if provided.
2. Otherwise read an environment variable such as
   `DIFFUSION_CLI_MODEL_ROOT`.
3. Otherwise use the default root above.
4. Join the required relative model paths.
5. Check each path exists and is a regular file.
6. Return a `ModelFiles` instance.

The implementation should not parse ComfyUI's YAML config at runtime.
Parsing it would make the new tool's behavior depend on a ComfyUI
installation. The default root can be documented and overridden.

### `zimage_model.py`

This module should own the diffusion model architecture and state dict
loading.

The first milestone has two viable implementation strategies:

#### Strategy A: Minimal Local Port

Port the specific Z-Image/Lumina2 model pieces needed by this workflow:

- Lumina2/NextDiT architecture.
- Z-Image state dict key mapping if needed.
- Model forward pass for text-to-image.
- Flux latent format scaling.
- DiscreteFlow model sampling.

This gives maximum control and avoids requiring upstream `diffusers`
support that may not exist or may not match ComfyUI's checkpoint format.

The downside is implementation time. The architecture is not tiny.

#### Strategy B: External Pipeline If Available

Use a library implementation only if it can load the exact local
safetensors files without downloading and without ComfyUI imports.

This should be investigated during implementation, but the design should
not rely on it. Current local evidence shows ComfyUI has custom Z-Image
support, and the public Diffusers documentation does not obviously
document a Z-Image pipeline.

The first implementation should therefore plan for Strategy A, while
keeping the loader boundary narrow enough that Strategy B can replace it
later.

### `text_encoder.py`

This module should load the Qwen3-4B text encoder and produce text
conditioning tensors.

Required behavior:

1. Load `qwen_3_4b.safetensors` with `safetensors`.
2. Instantiate the Qwen3-4B architecture.
3. Use the tokenizer vocabulary expected by ComfyUI's
   `qwen25_tokenizer`.
4. Wrap prompts with the Z-Image chat template.
5. Tokenize without weighted prompt syntax.
6. Run the text encoder under `torch.inference_mode`.
7. Select the `-2` hidden layer.
8. Return hidden states and attention masks in the format expected by
   the diffusion model.

Using `torch.inference_mode` is appropriate because generation does not
need gradients. PyTorch documents it as a context manager for inference
workloads that disables autograd-related overhead:
[torch.inference_mode](https://docs.pytorch.org/docs/stable/generated/torch.inference_mode.html).

### `sampling.py`

This module should contain only the sampler and scheduler needed for the
first image.

Components:

- `ModelSamplingDiscreteFlow`
- `betaScheduler`
- sigma/log-SNR conversion helpers
- `sampleSeeds3`

The implementation should use names that follow the project style:

- Public classes in `CapCase`.
- Functions in `camelCase` when multi-word.
- Variables in `snake_case`.

The sampling flow should be:

1. Build initial Gaussian noise with shape
   `[batch_size, 16, height // 8, width // 8]`.
2. Use a seeded generator to create deterministic initial noise on the
   selected CUDA device.
3. Build an empty latent image of zeros with the same shape.
4. Build the shifted flow sigma schedule with shift `3.0`.
5. Build beta scheduler sigmas for `steps`.
6. Run SEEDS-3 over those sigmas.
7. Return the final latent.

CFG behavior:

- At `cfg = 1.0`, conditional and unconditional guidance should reduce
  to the positive prediction.
- Even though CFG is 1.0 by default, the CLI should still compute the
  negative prompt path unless a later profiling pass proves it can be
  skipped safely.
- Keeping the negative path initially makes the code match the workflow
  structure and avoids special cases.

### `vae.py`

This module should load and decode the VAE.

Required behavior:

1. Load `ae.safetensors`.
2. Detect the architecture from state dict keys.
3. Instantiate the matching VAE decoder.
4. Load weights.
5. Move to selected device and dtype.
6. Decode final latent.
7. Return an image tensor in `[0, 1]`.

The first milestone only needs decoding. Encoding can be omitted.

### `image_io.py`

This module should save image tensors as PNG.

Input tensor:

```text
[batch_size, 3, height, width]
```

Save behavior:

- If `batch_size == 1`, write exactly the requested output path.
- If `batch_size > 1`, append an index before the suffix:
  `output_0000.png`, `output_0001.png`, and so on.
- Convert tensor layout from `NCHW` to `NHWC`.
- Clamp before converting to `uint8`.

## Dependency Plan

Use `uv` as the project manager. The official uv documentation describes
uv as a Python project and package manager with lockfile support:
[uv documentation](https://docs.astral.sh/uv/).

Initial dependencies:

- `torch`
- `safetensors`
- `transformers`
- `numpy`
- `scipy`
- `pillow`
- `tqdm`

Potential dependencies:

- `einops`, if the local model port needs shape rearrangement helpers.
- `accelerate`, only if memory management becomes difficult.

Do not add `diffusers` unless implementation research proves it can load
the exact local Z-Image Turbo files without downloads and without format
conversion. Adding a large dependency without using it would obscure the
actual first milestone.

Use `safetensors` for all three model files. Hugging Face documents
Safetensors as a safe tensor storage format and shows `safe_open` for
loading tensors:
[Safetensors documentation](https://huggingface.co/docs/safetensors/index).

## Licensing

ComfyUI source files inspected for this design include GPL license text.
For this first milestone, licensing is not a blocker. The project owner
has explicitly allowed copying ComfyUI code directly if that is the
fastest way to generate the first image.

The runtime constraint still matters: copied or adapted code may live in
this repository, but the CLI should not import ComfyUI as an installed
application or require a working ComfyUI runtime. This keeps the
milestone aligned with the product requirement that the tool stand
alone.

## Implementation Plan

### Phase 1: Project Scaffolding

Create:

- `pyproject.toml`
- `diffusion_cli/`
- `tests/`
- minimal `README.md`

Add a console script named `diffusion-cli`.

Success criteria:

- `uv run diffusion-cli --help` works.
- CLI arguments are visible and have correct defaults.
- No model loading happens during `--help`.

### Phase 2: Path and Config Validation

Implement model path resolution and config validation.

Success criteria:

- Missing model files produce a direct error listing the missing path.
- Invalid dimensions produce a direct error.
- Output path validation catches unwritable or invalid paths early.

### Phase 3: Tensor Loading Probe

Implement safetensors loading utilities and a diagnostic command or
internal debug function that prints:

- number of tensors,
- top-level key patterns,
- detected dtype summary,
- selected architecture guess.

Success criteria:

- The diffusion model, text encoder, and VAE files can be inspected
  without running full inference.
- The program can identify the expected Z-Image/Qwen/VAE key patterns.
- No download attempt occurs.

### Phase 4: Text Encoder

Implement prompt templating and Qwen3-4B conditioning.

Success criteria:

- A short prompt produces a conditioning tensor.
- The conditioning tensor has plausible shape.
- Positive and negative prompt encodings both complete.
- Encoding runs under `torch.inference_mode`.

### Phase 5: Diffusion Model

Implement or integrate the Z-Image/Lumina2 diffusion model.

Success criteria:

- The model loads the local diffusion state dict.
- A single forward pass accepts:
  - latent noise,
  - sigma/timestep,
  - text conditioning,
  - attention mask,
  - token count.
- The forward output shape matches the latent shape.
- No NaN or Inf values appear in the output.

### Phase 6: Scheduler and Sampler

Implement beta scheduler and SEEDS-3.

Success criteria:

- Sigma schedule starts at the expected high-noise value and ends at
  zero.
- The sampler runs for 10 steps.
- The latent remains finite.
- The implementation can run with `cfg=1.0`.

### Phase 7: VAE Decode and PNG Save

Implement VAE decode and image saving.

Success criteria:

- Final decoded tensor has shape `[1, 3, 1248, 832]`.
- Pixel range after postprocess is `[0, 1]`.
- Output PNG exists.
- The image is not obvious random noise or severe gibberish.

## Runtime Flow

The final command should look like:

```bash
uv run diffusion-cli \
    --prompt "a cinematic photo of a red apple on a wooden table" \
    --output output.png
```

Detailed execution:

1. Parse CLI arguments.
2. Validate generation config.
3. Resolve model file paths.
4. Select device and dtype.
5. Load text encoder.
6. Encode positive prompt.
7. Encode negative prompt.
8. Load diffusion model.
9. Build initial latent noise.
10. Run sampler.
11. Load VAE.
12. Decode latent to image.
13. Save PNG.
14. Print the output path.

The implementation may load VAE before sampling if that is simpler, but
sampling first can reduce peak memory if the text encoder or diffusion
model can be unloaded before VAE decode. The first milestone should
favor clarity over aggressive memory management.

## Error Handling

Errors should be direct and actionable.

Examples:

- `Missing diffusion model: /path/to/file.safetensors`
- `Width must be a positive multiple of 8: got 830`
- `CUDA requested but torch.cuda.is_available() is false`
- `bf16 requested but selected device does not support bf16`
- `Unsupported VAE state dict: no known decoder keys found`

Do not print Python tracebacks for expected user errors. Raise custom
exceptions internally and catch them in `cli.main`.

Unexpected programming errors may still show tracebacks. That is useful
during this early milestone.

## Testing Strategy

Unit tests should cover the parts that do not require loading the full
models:

- CLI parsing.
- Path resolution.
- Dimension validation.
- Batch output naming.
- Beta scheduler shape and final zero sigma.
- Deterministic seed noise shape.

Integration tests with the full model should be manual at first because
the model files are large and local to this machine.

Manual test command:

```bash
uv run diffusion-cli \
    --prompt "a simple studio photo of a ceramic mug" \
    --negative-prompt "text, watermark, full-body" \
    --seed 1234 \
    --steps 10 \
    --cfg 1 \
    --width 832 \
    --height 1248 \
    --output test_output.png
```

Manual acceptance checklist:

- Program exits with code 0.
- Output file exists.
- Output dimensions are `832x1248`.
- Image has recognizable structure.
- Image is not pure noise, blank, all black, all white, or tiled garbage.
- Re-running with the same seed produces the same output on the same
  device and dtype.

## Known Risks

### Model Architecture Complexity

The Z-Image/Lumina2 model architecture is the biggest risk. The CLI
surface is easy; model execution is the hard part.

Mitigation:

- Start by loading state dict keys and matching them to ComfyUI's
  detection path.
- Port only the architecture pieces needed for this exact model.
- Validate one forward pass before implementing the full sampler.

### Tokenizer Assets

ComfyUI's tokenizer path is local to its source tree:

`comfy/text_encoders/qwen25_tokenizer`

On this machine that directory contains the tokenizer assets:

- `vocab.json`
- `merges.txt`
- `tokenizer_config.json`

The `.safetensors` text encoder file contains model weights, not the
tokenizer vocabulary. The standalone CLI must therefore receive a
tokenizer directory separately.

Because the PRD says not to download files, the implementation must not
fetch tokenizer files from Hugging Face at runtime.

Recommended first approach:

- Add `--tokenizer-path`.
- Require the user to specify the tokenizer directory for now.
- Validate that the directory contains `vocab.json`, `merges.txt`, and
  `tokenizer_config.json`.
- Later copy the Qwen tokenizer assets out of ComfyUI into this
  repository so the CLI is more self-contained.

### Memory

The workflow uses a large diffusion model and a Qwen3-4B text encoder.
Running all components on GPU at once may exceed VRAM.

Mitigation:

- Load and run the text encoder first.
- Move text encoder back to CPU or delete it before loading diffusion.
- Decode VAE after sampling.
- Keep batch size default at 1.
- Prefer bf16 on CUDA.
- Do not support CPU execution in the first milestone. If the selected
  device is not CUDA, fail early with a clear message.

### Exact ComfyUI Parity

The first milestone does not require exact bitwise or visual parity with
ComfyUI. It only requires a coherent image.

Mitigation:

- Match the highest-impact pieces first:
  - prompt template,
  - latent shape,
  - model sampling shift,
  - scheduler,
  - sampler,
  - VAE decode.
- Defer LoRA, UI workflow parity, and multi-sampler support.

## Out Of Scope

The first milestone should not include:

- LoRA loading.
- ComfyUI workflow execution.
- Image-to-image.
- ControlNet.
- Multiple samplers.
- Multiple schedulers.
- Model downloads.
- Web UI.
- Exact ComfyUI output parity.
- Prompt weighting syntax.

## Open Questions

No blocking open questions remain for the first milestone.

Resolved decisions:

- Tokenizer assets: require `--tokenizer-path` for now. Later, copy the
  Qwen tokenizer assets out of ComfyUI into this repository.
- Licensing: do not block on licensing for the first milestone. Copy
  ComfyUI code directly if needed.
- Device support: CUDA-only for the first milestone.
- Diffusion model: hard-code
  `diffusion_models/z-image_turbo/moodyPornMix_zitV3.safetensors` for
  the first image proof.
- All-in-one checkpoints: out of scope for the first milestone, but the
  design should leave room to support them later.

## Recommended Next Step

Start with scaffolding and model inspection, not the full sampler.

The first concrete implementation task should create:

- `pyproject.toml`
- CLI argument parsing
- path resolution
- safetensors inspection for the three model files

That gives immediate feedback about local model compatibility and avoids
spending time on sampler code before confirming the exact state dict
formats.