Changes
diff --git a/designs/design-6-krea.md b/designs/design-6-krea.md
new file mode 100644
index 0000000..baf664e
--- /dev/null
+++ b/designs/design-6-krea.md
@@ -0,0 +1,1144 @@
+# 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.