# Krea 2 Backend Design

## Purpose

This document describes how `diffusion-cli` should add Krea 2
text-to-image inference without removing or regressing the existing
Z-Image Turbo backend.

The project already owns the user-facing parts of local image generation:

- command-line parsing;
- TOML configuration;
- local safetensors path resolution;
- staged and CPU-cached model residency;
- image encoding and file output;
- a local web UI;
- a SillyTavern stable-diffusion.cpp compatibility API.

Krea 2 should therefore be another inference backend inside this project.
It should not be a separate CLI, a subprocess wrapper around the official
repository, a Git submodule, or a runtime dependency on ComfyUI. The
official Krea 2 and ComfyUI implementations are reference material for
model behavior and checkpoint compatibility.

The first milestone succeeds when the CLI uses the local scaled-FP8 Krea 2
Turbo diffusion model, scaled-FP8 Qwen3-VL-4B text encoder, Qwen Image VAE,
and tokenizer files to generate a recognizable image without making any
network request.

## Decision Summary

The implementation should make these decisions:

- Keep one application with multiple model backends.
- Introduce named model profiles selected with `--model-profile`.
- Preserve the current root-level one-shot CLI rather than requiring a new
  `generate` subcommand.
- Select one model profile when a server starts. HTTP requests must not
  switch profiles.
- Put model-family-specific loading and sampling behind an
  `InferenceBackend` interface.
- Retain Z-Image as `ZImageBackend` and add `Krea2Backend`.
- Reuse the existing generation requests, image output, server, API, and UI.
- Load every model artifact from an explicit local path.
- Never use a remote model identifier or download fallback.
- Support the observed ComfyUI scaled-FP8 format explicitly.
- Start with a correctness-first FP8 path that keeps stored weights in FP8
  and dequantizes one weight for each linear operation.
- Optimize with scaled matrix multiplication only after parity is proven.
- Construct the Diffusers Qwen Image VAE from in-code configuration and
  convert the local original-format state dict in memory.
- Keep Krea Raw and Turbo behavior separate through profile variants and
  variant-aware defaults.

## Goals

The feature must:

- Preserve existing Z-Image CLI, server, UI, configuration, and output
  behavior.
- Add named, selectable local model profiles.
- Add Krea 2 Turbo text-to-image generation.
- Leave a direct path for later Krea 2 Raw support.
- Accept separate local diffusion, text encoder, VAE, and tokenizer paths.
- Load the exact local checkpoint formats documented below.
- Avoid importing ComfyUI at runtime.
- Avoid downloading model files or tokenizer files.
- Recognize scaled-FP8 tensors, scales, and ComfyUI metadata.
- Keep scaled weights compressed while resident on CPU or CUDA.
- Produce the 12-layer Qwen3-VL conditioning stack required by Krea 2.
- Reproduce Krea 2 prompt templating and prefix stripping.
- Reproduce Krea 2 flow-matching sampling.
- Apply Qwen Image latent normalization and VAE decoding.
- Retain the existing generation lock for GPU residency changes.
- Expose the selected model through the existing UI and API.
- Add unit tests that do not require full models or CUDA.
- Add an optional local real-checkpoint integration test.
- Give actionable errors for incompatible files and hardware.

## Non-Goals

The first milestone should not:

- Add image-to-image or reference-image generation.
- Add LoRA loading, fine-tuning, or training.
- Add multi-GPU tensor parallelism.
- Add concurrent GPU generation.
- Import ComfyUI as a package.
- Support arbitrary quantization formats.
- Promise optimized FP8 throughput.
- Switch profiles per HTTP request.
- Load Raw and Turbo simultaneously.
- Reproduce Krea-hosted safety systems.
- Replace the existing API profile or UI.
- Add the native stable-diffusion.cpp asynchronous job API.

These exclusions keep the milestone focused on one local Krea 2 Turbo
image. Later features should build on a verified base model path.

## External References

Implementation should be checked against primary sources:

- [Official Krea 2 inference repository](https://github.com/krea-ai/krea-2)
  for architecture constants, conditioning, timestep shifting, and
  recommended Raw and Turbo settings.
- [ComfyUI Krea 2 model](https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/ldm/krea2/model.py)
  for checkpoint-compatible module structure and tensor names.
- [ComfyUI Krea 2 text encoder](https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/text_encoders/krea2.py)
  for prompt templating, layer taps, prefix stripping, and layout.
- [ComfyUI quantized operations](https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/ops.py)
  and [quantization layouts](https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/quant_ops.py)
  for `comfy_quant`, `weight_scale`, and scaled FP8.
- [PyTorch `scaled_mm`](https://docs.pytorch.org/docs/main/generated/torch.nn.functional.scaled_mm.html)
  for a later accelerated FP8 path.
- [PyTorch float8 design note](https://dev-discuss.pytorch.org/t/float8-in-pytorch-1-x/1815)
  for raw float8 values and separately tracked scales.
- [Transformers Qwen3-VL model](https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py)
  for text-model and hidden-state behavior.
- [Diffusers Qwen Image autoencoder](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py)
  for the causal 3D VAE and latent normalization.
- [Transformers offline mode](https://huggingface.co/docs/transformers/v4.49.0/en/installation#offline-mode)
  for `local_files_only` and `HF_HUB_OFFLINE`.
- [Safetensors documentation](https://huggingface.co/docs/safetensors/main/index)
  for selective local tensor and metadata access.

The local ComfyUI checkout is a version-pinned practical reference at:

```text
/home/mw/programs/ComfyUI
```

The program must not import that checkout at runtime.

The official Krea 2 implementation is also available as a local source
checkout at:

```text
~/programs/krea-2
```

This checkout is approved source material for the implementation. Code may
be copied and adapted from it directly instead of being reimplemented from
the public GitHub view. Copied files and fragments must retain the applicable
Apache 2.0 notices, identify modifications, and be reshaped around
`diffusion-cli`'s local checkpoint loaders and backend interface. The Krea 2
checkout is a development reference, not a runtime package dependency.

## Current Repository State

### Existing Layers

| Module | Current responsibility |
| --- | --- |
| `cli.py` | Parse one-shot and server arguments. |
| `config.py` | Load model paths and generation defaults. |
| `paths.py` | Resolve separate or all-in-one sources. |
| `checkpoint.py` | Load selected safetensors tensors. |
| `model_inspect.py` | Inspect tensor keys and dtypes. |
| `generation_service.py` | Serialize requests and manage residency. |
| `zimage_model.py` | Implement Z-Image diffusion. |
| `text_encoder.py` | Implement Z-Image Qwen conditioning. |
| `vae.py` | Implement the Z-Image VAE. |
| `sampling.py` | Implement Z-Image sampling. |
| `server.py` | Create the Flask application. |
| `api_profiles.py` | Translate compatibility API payloads. |
| `ui.py` | Serve and handle the local UI. |

The HTTP and image-output layers are already model-independent in practice.
The generation service and configuration are Z-Image-specific.

### Duplicate Generation Paths

One-shot generation in `cli.py` directly loads the text encoder, diffusion
model, and VAE. Server generation repeats that sequence in
`GenerationService`.

Before adding Krea 2, one-shot generation must call
`GenerationService.generateToFiles()`. This creates one shared generation
path and prevents backend logic from being duplicated again.

### Test Baseline

At design time, all 99 existing unit tests pass. The backend refactor must
preserve this baseline before Krea-specific behavior is added.

## Observed Local Model Artifacts

The intended model root is:

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

The first profile should use:

| Role | Local path |
| --- | --- |
| Diffusion | `diffusion_models/krea2/krea2_turbo_fp8_scaled.safetensors` |
| Text encoder | `text_encoders/krea2/qwen3vl_4b_fp8_scaled.safetensors` |
| VAE | `vae/krea2/qwen_image_vae.safetensors` |
| Tokenizer | `/home/mw/programs/ComfyUI/comfy/text_encoders/qwen25_tokenizer` |

ComfyUI's Qwen3-VL Krea path deliberately uses the Qwen 2.5 tokenizer
assets, so no new tokenizer download is required. The CLI should still
accept another explicit compatible local tokenizer directory.

### Diffusion File

The observed Turbo file is approximately 12.24 GiB:

| Dtype | Tensor count |
| --- | ---: |
| `F8_E4M3` | 256 |
| `BF16` | 174 |
| `F32` | 256 |

Quantized layers contain pairs such as:

```text
blocks.0.attn.wq.weight
blocks.0.attn.wq.weight_scale
```

File-level `_quantization_metadata` JSON identifies quantized layers and
marks some with `full_precision_matrix_mult`. The loader must parse it rather
than inferring behavior from dtype alone.

### Text Encoder File

The Qwen3-VL-4B file is approximately 4.88 GiB. Quantized layers contain:

```text
model.layers.0.self_attn.q_proj.weight
model.layers.0.self_attn.q_proj.weight_scale
model.layers.0.self_attn.q_proj.comfy_quant
```

`comfy_quant` is byte-encoded JSON. The loader must consume it as metadata,
not pass it to Transformers as an unexpected model parameter. Embeddings and
normalizations remain BF16.

### VAE File

The Qwen Image VAE is approximately 0.24 GiB with 194 BF16 tensors. It uses
original Wan/Qwen names including:

```text
decoder.conv1.weight
encoder.conv1.weight
conv1.weight
conv2.weight
```

This is a causal 3D autoencoder, not the existing Z-Image 2D VAE. A still
image is represented with a time dimension of one.

### Other Krea Files

The model root also has a Raw scaled-FP8 checkpoint and a mixed
INT8/INT4/FP8 checkpoint. Raw support comes after Turbo. The mixed checkpoint
is unsupported and must not be accepted merely because its keys resemble
Krea 2.

## Terminology

### Model Profile

A model profile is a named collection of paths, architecture identity,
variant, and model-specific defaults. Examples are `z-image-turbo` and
`krea2-turbo`.

### Architecture and Variant

Architecture selects the inference implementation. Initial values are
`z-image` and `krea2`. Variant selects training and sampling expectations
inside an architecture. Krea variants are `raw` and `turbo`.

Configured architecture is authoritative. Inspection verifies compatibility
but must not silently switch backends.

### Runtime and Storage Dtypes

Runtime dtype is used for activations and ordinary computation, normally
BF16. Storage dtype is retained for parameters. A scaled-FP8 linear weight
must remain `torch.float8_e4m3fn` even when activations use BF16.

Calling `.to(dtype=torch.bfloat16)` on the whole quantized model would expand
every FP8 weight and defeat the checkpoint format.

### Correctness-First FP8

The first path reconstructs one BF16 linear weight temporarily:

```text
weight = fp8_weight.to(bfloat16) * weight_scale
output = linear(input, weight, bias)
```

Only one expanded layer should exist at a time. This is slower than an
optimized FP8 GEMM but establishes a clear correctness reference.

## User Interface

### Backward Compatibility

This must continue to work with legacy configuration:

```bash
uv run diffusion-cli --prompt "a ceramic mug"
```

Legacy `[models]` configuration becomes one implicit Z-Image profile named
`legacy` in memory. The file is not rewritten.

### Profile Selection

Add global option:

```text
--model-profile NAME
```

Selection precedence is:

1. command-line `--model-profile`;
2. TOML `default_model`;
3. implicit `legacy` when legacy paths exist;
4. an error explaining that no model is configured.

Examples:

```bash
uv run diffusion-cli \
    --model-profile krea2-turbo \
    --prompt "a fox walking through fresh snow"
```

```bash
uv run diffusion-cli \
    --model-profile krea2-turbo \
    serve \
    --api-profile sillytavern-sdcpp
```

Because it is global, the option appears before `serve` with the current
`argparse` layout.

### Path Overrides

Existing options continue to override one selected-profile component:

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

Overrides do not change architecture or variant. A Krea checkpoint selected
under a Z-Image profile should fail compatibility validation.

### Krea Sampling Overrides

Add optional values that remain unset during argument parsing:

```text
--mu FLOAT
--shift-y1 FLOAT
--shift-y2 FLOAT
```

Existing `--steps` and `--cfg` remain shared. A backend should reject an
override it does not understand; Z-Image should not silently ignore Krea
shift options.

### Inspection

`--inspect-models` should add profile, architecture, and variant. For each
component it should report:

- path and selected checkpoint prefix;
- tensor and dtype counts;
- detected architecture or component format;
- quantized layer count;
- missing scale count;
- unsupported formats;
- compatibility with the selected profile.

Inspection must not initialize CUDA or allocate full checkpoints.

## Configuration

### Named Profiles

Recommended configuration:

```toml
default_model = "krea2-turbo"

[model_profiles.krea2-turbo]
architecture = "krea2"
variant = "turbo"
diffusion_model = "/home/mw/documents/ai/diffusion/models/diffusion_models/krea2/krea2_turbo_fp8_scaled.safetensors"
text_encoder = "/home/mw/documents/ai/diffusion/models/text_encoders/krea2/qwen3vl_4b_fp8_scaled.safetensors"
vae = "/home/mw/documents/ai/diffusion/models/vae/krea2/qwen_image_vae.safetensors"
tokenizer = "/home/mw/programs/ComfyUI/comfy/text_encoders/qwen25_tokenizer"

[model_profiles.z-image-turbo]
architecture = "z-image"
variant = "turbo"
diffusion_model = "/models/z-image/diffusion.safetensors"
text_encoder = "/models/z-image/qwen_3_4b.safetensors"
vae = "/models/z-image/ae.safetensors"
tokenizer = "/models/qwen25_tokenizer"

[generation]
width = 1024
height = 1024
output = "output.png"
device = "cuda"
dtype = "auto"
```

### Data Structures

Replace the single `ModelPathConfig` with:

```python
@dataclass(frozen=True)
class ModelProfile:
    """Local components and behavior for one named backend."""

    name: str
    architecture: str
    variant: str
    checkpoint: Path | None = None
    diffusion_model: Path | None = None
    text_encoder: Path | None = None
    vae: Path | None = None
    tokenizer: Path | None = None


@dataclass(frozen=True)
class UserConfig:
    """Validated profiles and generation defaults."""

    model_profiles: dict[str, ModelProfile]
    default_model: str | None
    generation: GenerationDefaults
```

Every public item retains intention comments per project style.

### Validation

Reject:

- empty profile names;
- unknown architectures;
- unsupported variants;
- unknown keys or non-string paths;
- a nonexistent `default_model`;
- profiles with neither component paths nor a supported checkpoint.

Krea all-in-one checkpoints are not required initially. A Krea profile with
only `checkpoint` should explain that separate files are required.

## Backend Architecture

### Registry

Add one registry in `backends.py`:

```python
BACKEND_FACTORIES = {
    "z-image": ZImageBackend,
    "krea2": Krea2Backend,
}
```

CLI validation, inspection, and service construction use this registry as
the source of architecture names.

### Interface

```python
class InferenceBackend(Protocol):
    """Generate image tensors for one local model profile."""

    @property
    def modelId(self) -> str:
        """Return the stable identifier exposed to clients."""

    def generate(self, config: GenerationConfig):
        """Return NCHW RGB image tensors in the range zero through one."""
```

The constructor receives a `ModelProfile`, resolved `ModelSources`, and
residency policy. The backend owns component instances and movement.

### Residency Ownership

Krea and Z-Image movement differs. A generic `.to(device, dtype)` call would
expand Krea FP8 weights. Each backend must own:

- staged component creation and destruction;
- CPU cache creation;
- device movement;
- CUDA cache release;
- model-specific dtype preservation.

Small lifecycle helpers may be shared, but model correctness is more
important than forcing every operation into one abstraction.

### Generation Service

The service should only validate, lock, delegate, and encode:

```python
class GenerationService:
    """Serialize requests to one selected inference backend."""

    def generateImages(self, request):
        with self._lock:
            config = buildGenerationConfigFromRequest(
                request,
                self._backend.profile,
                self.user_config,
            )
            images = self._backend.generate(config)
            return encodeGeneratedImages(images, config)
```

The real implementation continues returning existing `GeneratedImage`
objects. One-shot CLI creates the same service and calls `generateToFiles()`.
It uses staged residency because the process exits after one request.

## Source Layout

Add:

```text
diffusion_cli/
    backends.py
    krea2_backend.py
    krea2_model.py
    krea2_sampling.py
    krea2_text_encoder.py
    qwen_image_vae.py
    quantization.py
```

Keep current Z-Image filenames for the first change. Renaming
`text_encoder.py`, `vae.py`, and `sampling.py` while adding the boundary would
increase review noise without helping the milestone.

## Local-Only Contract

Runtime loading must obey all these rules:

- Model tensors come from resolved local regular files.
- Tokenizer assets come from a resolved local directory.
- No loader receives a Hugging Face repository ID.
- No loader calls a Hub download function.
- Any `from_pretrained()` receives a local path and
  `local_files_only=True`.
- Missing configured files fail immediately.
- A cached remote artifact must not satisfy a missing configured path.

Set before importing Transformers or Diffusers:

```text
HF_HUB_OFFLINE=1
TRANSFORMERS_OFFLINE=1
DIFFUSERS_OFFLINE=1
```

These variables are defense in depth; explicit local construction is the
primary guarantee.

The tokenizer requires `vocab.json`, `merges.txt`, and
`tokenizer_config.json`. Load it with `Qwen2Tokenizer`, a local path,
right-padding, and pad token ID `151643`.

Tests should patch Hub download helpers to raise. Successful tiny-component
construction and missing-path failures must never call them.

## Quantization

### Inspection

Recognize safetensors `F8_E4M3` as `torch.float8_e4m3fn`, but do not treat it
as a component runtime dtype. Krea chooses BF16 activations when supported
while preserving individual parameter storage dtypes.

### Metadata Model

```python
@dataclass(frozen=True)
class QuantizedLayerSpec:
    """Quantization settings for one parameterized layer."""

    format: str
    full_precision_matrix_mult: bool = False


@dataclass(frozen=True)
class QuantizationManifest:
    """Quantization settings indexed by module path."""

    layers: dict[str, QuantizedLayerSpec]
```

Diffusion stores this in file metadata under `_quantization_metadata`.
The encoder stores equivalent JSON in each `*.comfy_quant` byte tensor.
Normalize both into one manifest before model construction.

### Supported Format

Initially support only:

```text
float8_e4m3fn with one scalar F32 weight_scale per linear weight
```

Reject missing or nonscalar scales, unknown formats, FP8 convolutions,
quantized biases, INT8/FP4 storage, and inconsistent metadata. Rejecting an
unknown format is safer than producing plausible but incorrect images.

### Scaled FP8 Linear

`ScaledFp8Linear` owns an FP8 weight, F32 scalar scale, and optional ordinary
bias. Its first forward path:

1. reads the activation dtype;
2. converts only this weight to that dtype;
3. multiplies by the scale on the activation device;
4. calls `F.linear`;
5. releases the expanded temporary on return.

The scale reconstructs the original weight:

```text
dequantized = stored_weight * weight_scale
```

Tests must verify this against explicit reconstruction.

### Device Movement

Never call `.to(device=device, dtype=runtime_dtype)` on the full quantized
model. Movement must distinguish:

- FP8 parameters: move device, preserve FP8;
- BF16 parameters: move and use runtime dtype;
- F32 scales or sensitive parameters: move and preserve F32;
- integer metadata: consume at load time.

Use an explicit movement helper with tests that assert stored weight dtype.

### Optimized Path

Later, `ScaledFp8Linear` may use `torch.nn.functional.scaled_mm`. Dispatch
must verify CUDA, PyTorch API availability, device capability, matrix
alignment, and supported activation/output dtypes. Unsupported systems use
the correctness path. Optimization cannot change scale interpretation.

## Krea Diffusion Model

### Supported Constants

| Setting | Value |
| --- | ---: |
| Feature width | `6144` |
| Timestep width | `256` |
| Text width | `2560` |
| Main heads | `48` |
| KV heads | `12` |
| Main blocks | `28` |
| MLP multiplier | `4` |
| Patch size | `2` |
| Latent channels | `16` |
| Text layers | `12` |
| Text-fusion heads | `20` |
| Text-fusion KV heads | `20` |
| RoPE theta | `1000` |

Validate these against checkpoint shapes; do not expose them as user
settings in the first milestone.

### Derived Code

`krea2_model.py` should begin with the Apache-2.0 official implementation in
`~/programs/krea-2`. Code may be copied from that checkout and adapted to use
the backend interface, local scaled-FP8 layers, and strict state-dict
accounting. Derived files must retain applicable notices and state that they
were modified. Use the matching ComfyUI implementation to verify checkpoint
compatibility, but do not port ComfyUI framework hooks, patch extensions,
reference images, or model-management code.

### Components

The model contains latent patch projection, timestep embedding, a text-fusion
transformer over 12 Qwen layers, text projection, 28 single-stream blocks,
grouped-query attention, QK RMS normalization, sigmoid attention gating,
SwiGLU, AdaLN-style modulation, three-axis RoPE, and final patch projection.

### Loading

Construct on the meta device. For each linear module:

1. inspect the manifest;
2. construct `ScaledFp8Linear` for supported quantized layers;
3. otherwise construct ordinary `nn.Linear`;
4. consume weight, optional bias, and optional scale;
5. record every consumed state key;
6. reject missing and remaining keys outside a documented allowlist.

Strict accounting is essential because ignored scales preserve tensor shapes
while corrupting values.

### Forward Contract

The wrapper accepts NCHW latent, timestep, and `Krea2Conditioning`. Validate
four dimensions, 16 channels, batch sizes, 12 text layers, 2560 text
features, mask length, output shape, and finite values. It returns flow
velocity with the same shape as the input latent.

## Krea Text Encoder

### Conditioning

Tap Qwen hidden states:

```text
2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35
```

Return:

```text
[batch, sequence, 12, 2560]
```

Use a Krea-specific type rather than silently reusing Z-Image's
three-dimensional conditioning type:

```python
@dataclass(frozen=True)
class Krea2Conditioning:
    """Twelve Qwen hidden-state taps and their token mask."""

    hidden_states: torch.Tensor
    attention_mask: torch.Tensor
```

### Prompt Template

Apply exactly:

```text
<|im_start|>system
Describe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>
<|im_start|>user
{prompt}<|im_end|>
<|im_start|>assistant
```

After encoding, remove the system and user-opening prefix. Prefix stripping
must be token-aware and tested against the known special-token boundary, not
implemented as a character offset.

### Construction

Reuse compatible Qwen configuration work, but create a Krea encoder because
it needs 12 taps, the Krea template, prefix stripping, Krea layout, and FP8
linear replacement. Instantiate only the text language backbone from in-code
configuration. Do not allocate a vision tower or language-model head when
the checkpoint does not contain them.

Normalize `model.` prefixes and consume `comfy_quant` metadata before strict
state-dict validation.

When CFG is zero, encode only the positive prompt. Turbo's default therefore
avoids an unnecessary second text-encoder pass.

## Krea Sampler

### Noise and Dimensions

The latent downscale is eight and MMDiT patch size is two, so image dimensions
align to 16. For official compatibility, round nonmultiples up with a visible
warning. Image `i` uses `base_seed + i` with an independent CUDA generator.

### Timesteps

Create a uniform grid from one to zero with `steps + 1` entries and apply:

```text
shifted(t) = exp(mu) / (exp(mu) + (1 / t - 1) ** sigma)
```

When `mu` is unset, interpolate it from image-token sequence length between
minimum resolution 256 and maximum resolution 1280, with `y1 = 0.5` and
`y2 = 1.15`. Turbo pins `mu = 1.15`; Raw derives it from resolution.

### Euler Integration

For each adjacent timestep pair:

1. evaluate conditional velocity;
2. evaluate unconditional velocity only when CFG is enabled;
3. combine them using the official Krea guidance equation;
4. update latent with explicit Euler integration;
5. report progress through an optional callback;
6. reject NaN or Inf.

Do not reuse Z-Image's denoised-prediction CFG helper; the contracts differ.

### Defaults

| Setting | Turbo | Raw |
| --- | ---: | ---: |
| Steps | `8` | `52` |
| CFG | `0.0` | `3.5` |
| Mu | `1.15` | unset |
| Width | `1024` | `1024` |
| Height | `1024` | `1024` |

Precedence is request/CLI, `[generation]`, profile variant, then architecture
fallback. This replaces the current universal Z-Image fallback constants.

## Qwen Image VAE

### Dependency and Construction

Add `diffusers` and use `AutoencoderKLQwenImage`. Do not call
`from_pretrained()` or `from_single_file()`, because a high-level loader may
seek remote configuration when the single file has none.

Instead:

1. instantiate from explicit constructor values;
2. load local safetensors;
3. convert original Wan/Qwen keys to Diffusers keys in memory;
4. strictly load the converted dict;
5. use decode behavior only.

Use checkpoint-compatible defaults: base dimension 96, latent dimension 16,
channel multipliers `[1, 2, 4, 4]`, two residual blocks, temporal downsample
flags `[False, True, True]`, and three input channels. Use the exact
16-element Qwen Image latent means and standard deviations from the public
Diffusers configuration.

### Key Conversion

Own an attributed conversion based on Diffusers'
`convert_wan_vae_to_diffusers`; do not import that private helper. Convert:

- encoder and decoder `conv1` inputs;
- `conv1` and `conv2` quant projections;
- middle residual and attention blocks;
- residual norms and convolutions;
- downsample and upsample blocks;
- encoder and decoder output heads.

Tests use representative keys from every mapping family and ensure no source
key disappears silently.

### Decode

Starting from NCHW latent:

1. insert a time dimension of one;
2. apply per-channel standard deviation and mean;
3. decode the five-dimensional tensor;
4. remove the time dimension;
5. convert model range `[-1, 1]` to `[0, 1]`;
6. validate batch, RGB channels, size, and finite values.

The complete public VAE may be instantiated initially because it is small.
A decoder-only optimization should wait for memory measurements.

## Residency

### Staged

Staged mode loads and releases the text encoder, then diffusion model, then
VAE. It minimizes simultaneous VRAM and is the one-shot default.

### CPU Cache

CPU-cache mode constructs each once on CPU. For each request it moves one
stage to CUDA, executes it, and moves it back. FP8 weights remain FP8 during
movement. The existing service lock covers the complete sequence.

The two quantized files total about 17.12 GiB before allocator overhead, but
they are not on CUDA simultaneously. Diffusion peak also includes one
expanded linear weight, conditioning, image tokens, attention intermediates,
and latents.

CUDA OOM errors should name the active stage and profile. The program must
not silently retry with a different numeric policy.

## API and UI

The server selects one profile at startup. `/v1/models` returns that profile
name instead of hard-coded `z-image-local`. Health responses add profile,
architecture, and variant but never local filesystem paths.

The existing txt2img request fields remain unchanged. At CFG zero, Krea
operationally skips `negative_prompt`, but response metadata retains it.
Krea shift values need not be exposed through HTTP initially.

The UI receives effective defaults for the selected profile. A Krea Turbo
server displays eight steps and zero CFG unless user configuration overrides
them. Model identity is read-only; the UI does not switch large backends.

## Error Handling

Expected errors use `DiffusionCliError`. Examples include:

```text
Unknown model profile: krea-two
Unsupported Krea 2 variant: distilled
Missing model_profiles.krea2-turbo.vae
Krea layer blocks.0.attn.wq is missing weight_scale
Unsupported quantization format in blocks.0.attn.wq: nvfp4
Expected scalar weight scale, got shape (6144,)
Text encoder checkpoint is not Qwen3-VL-4B
Qwen Image VAE conversion left unexpected keys
Krea diffusion output contains NaN or Inf
```

Server logs retain unexpected tracebacks while HTTP clients receive the
existing generic generation-failure response.

## Testing Strategy

### Backend and Configuration

Test the backend registry, unknown architectures, named-profile parsing,
default and CLI selection, direct overrides, legacy conversion, coexistence
of legacy and named profiles, variant validation, and defaults. Verify
one-shot CLI uses `GenerationService` and all old tests continue to pass.

### Inspection and Quantization

Create tiny safetensors fixtures for Krea diffusion, Qwen encoder, and Qwen
VAE. Test FP8 counts, file metadata, byte JSON, missing scales, unsupported
formats, and mixed checkpoint rejection.

For `ScaledFp8Linear`, quantize a tiny known weight, compare output with
explicit reconstruction, and ensure stored dtype survives CPU/CUDA movement.
CUDA tests skip only when CUDA is unavailable.

### MMDiT

Use a reduced Krea config to test patch round-trips, KV-head repetition, RMS
normalization convention, sigmoid gating, text fusion, conditioning shape,
finite output, and strict state-key accounting.

### Text Encoder

Use fake tokenizer and model objects to test the exact template, token-aware
prefix boundary, layer indices, output layout, mask trimming, CFG-zero skip,
and `local_files_only=True`.

### Sampler

Test dimension alignment, per-image seeds, Raw and Turbo `mu`, descending
timesteps, Euler updates with constant fake velocity, exact guidance, no
unconditional call at CFG zero, and finite-value errors.

### VAE

Test every conversion mapping family, key accounting, latent normalization,
time-axis handling, RGB output range and shape, and construction without
high-level pretrained loaders.

### Service, API, and UI

Test dynamic model IDs, health information, profile defaults, Krea requests
through a fake backend, serialization, and absence of filesystem paths in
responses.

### Real Integration

Add an opt-in script outside ordinary unit discovery. It requires explicit
local paths, generates one 1024-square Turbo image with eight steps, CFG zero,
mu 1.15, and a fixed seed, then reports time and peak CUDA memory. Verify
dimensions and non-degenerate pixel statistics. Optionally compare a
perceptual hash with a known ComfyUI image; exact pixels may vary by kernel
and GPU.

## Implementation Sequence

### Phase 1: Backend Boundary

1. Add profile types and legacy conversion.
2. Add CLI profile selection.
3. Add backend registry and interface.
4. Move existing generation into `ZImageBackend`.
5. Route one-shot generation through the service.
6. Make API model identity dynamic.
7. Restore all 99 tests before proceeding.

### Phase 2: FP8 Feasibility Gate

1. Extend inspection for FP8.
2. Parse both metadata forms.
3. Validate supported format strictly.
4. Implement correctness-first `ScaledFp8Linear`.
5. Add dtype-preserving movement.
6. Validate one real diffusion and text linear weight.

Do not build the full model until a real scaled layer produces finite,
correct output.

### Phase 3: MMDiT

1. Adapt the architecture with attribution.
2. Replace linear modules from the manifest.
3. Add strict loading and reduced forward tests.
4. Load the real model on CPU.
5. Run a real block or reduced CUDA forward.

### Phase 4: Text Conditioning

1. Build explicit Qwen configuration.
2. Reuse local tokenizer assets.
3. Add scaled-FP8 Qwen loading.
4. Implement template, trimming, and 12 taps.
5. Compare conditioning shape and mask with ComfyUI.

### Phase 5: VAE

1. Add Diffusers.
2. Construct the VAE explicitly.
3. Convert local keys.
4. Strictly load the real file.
5. Decode a synthetic latent.

### Phase 6: End-to-End

1. Implement noise, timesteps, Euler, and CFG.
2. Add Raw and Turbo defaults.
3. Connect stages in `Krea2Backend`.
4. Generate one fixed-seed Turbo image.
5. Compare with ComfyUI and record memory and time.

### Phase 7: Server Completion

1. Expose dynamic identity and defaults.
2. Exercise fake-backend API and UI tests.
3. Exercise one real API request.
4. Update README examples and offline guarantee.

### Phase 8: Optional Optimization

1. Benchmark per-call dequantization.
2. Prototype `scaled_mm`.
3. Add hardware and alignment detection.
4. Compare numerics with the fallback.
5. Enable only after parity passes.

## Acceptance Criteria

The feature is complete when:

- all old tests pass;
- named Z-Image and Krea profiles coexist;
- legacy configuration still works;
- inspection recognizes all local Krea components;
- loaders make no network request;
- scaled weights remain FP8 while resident;
- conditioning contains the 12 selected layers;
- Turbo defaults are eight steps, CFG zero, and mu 1.15;
- the Qwen Image VAE decodes a synthetic latent;
- one-shot Turbo generation produces a recognizable 1024-square image;
- the same backend generates through the server;
- API model identity matches the selected profile;
- errors name the component and stage;
- README documents configuration and offline behavior;
- adapted Apache code retains attribution.

## Risks and Mitigations

### FP8 Correctness

A scale can be inverted or ignored while shapes remain valid. Compare tiny
outputs with explicit dequantization, account for every scale key, and compare
real behavior with ComfyUI.

### FP8 Performance

Per-call reconstruction may be slow. Keep it as the correctness reference,
then add `scaled_mm` behind detection and parity tests.

### VRAM

The diffusion model plus activations may exceed device memory. Preserve FP8,
stage components, release temporaries, and measure before designing block
offload.

### Dependency Drift

Transformers structure and Diffusers private conversion helpers may change.
Pin dependencies, use explicit configuration, own the VAE mapping, and keep
strict missing/unexpected-key tests.

### Z-Image Regression

Profiles and backend extraction can break working behavior. Land the boundary
separately, preserve legacy config, and require the existing suite to pass
before Krea code lands.

### Reference Divergence

Official Krea and ComfyUI can differ in trimming, timesteps, or names. Use
official Krea for intent and recommended sampling, ComfyUI for actual local
format, document each choice, and verify a fixed local workflow.

## Licensing

Official Krea 2 source is Apache 2.0. The local source checkout at
`~/programs/krea-2` may be copied from directly. Adapted files must retain
applicable notices, state that they were modified, ship the license, and
preserve attribution. If ComfyUI code is copied rather than independently
adapted, its license obligations must also be reviewed and preserved.

Model weights use a separate Krea community license. Documentation must not
imply that the source-code license governs weights.

## Future Work

Later designs may cover Raw validation, optimized FP8 kernels, LoRAs,
reference images, explicit profile unloading and switching, VAE tiling,
diffusion block offload, other quantization formats, and native
stable-diffusion.cpp API profiles.

Each should build on the backend boundary instead of adding model conditions
to shared HTTP or image-output code.