Krea-2-Turbo-OrbitQuant-W4A4

Krea 2 Turbo OrbitQuant W4A4

Deployment checkpoint for krea/Krea-2-Turbo. Both transformer-class components use OrbitQuant: the Qwen3-VL text encoder and the Krea 2 diffusion transformer. The Qwen linears are fully W4A4. The DiT uses W4A4 for its attention and gate/up projections; its 28 quality-sensitive main-block ff.down projections, which the quantization policy keeps out of 4-bit, run as per-row INT8 weights with INT8 activations (W8A8), and the latent input projection img_in stays in BF16.

<a href="https://huggingface.co/WaveCut/Krea-2-Turbo-OrbitQuant-W4A4/resolve/main/assets/original_vs_orbitquant_w4a4.webp"><img src="https://huggingface.co/WaveCut/Krea-2-Turbo-OrbitQuant-W4A4/resolve/main/assets/original_vs_orbitquant_w4a4_preview.webp" alt="BF16 versus quality-protected OrbitQuant W4A4 on ten paired prompts" width="100%"></a>

The embedded preview links to the lossless 8192×10720 matrix. Every source tile in the original matrix remains at the model's native benchmark output size of 2048×2048; labels are added and tiles are concatenated without resizing. The twenty individual PNGs are in artifacts/generations/.

At a glance

Install and run

python -m venv .venv
source .venv/bin/activate
pip install -r https://huggingface.co/WaveCut/Krea-2-Turbo-OrbitQuant-W4A4/resolve/main/runtime-requirements.txt
hf download WaveCut/Krea-2-Turbo-OrbitQuant-W4A4 scripts/run_inference.py --local-dir .
orbitquant kernels-install --build

kernels-install uses a compatible prebuilt kernel when available. --build allows an exact local build for an unlisted Torch/CUDA ABI and requires the CUDA toolkit plus ninja.

The recommended runner stages Qwen and the DiT sequentially, enables tiled VAE decoding, and requires Flash SDPA plus packed OrbitQuant kernels:

python scripts/run_inference.py \
  --prompt "A rain-soaked Warsaw street seen through a tram window" \
  --width 2048 --height 2048 --steps 8 --seed 0 \
  --output krea2-orbitquant.png

This staging policy matters at 2048×2048. Qwen is used to produce prompt embeddings and released before the DiT and VAE move to CUDA. Keeping every component resident at once is simpler, but wastes VRAM without improving the result.

Direct Diffusers loading is also supported when the GPU can hold the whole pipeline:

import os
import torch

os.environ.setdefault("ORBITQUANT_STRICT_PACKED", "1")
import orbitquant  # registers the OrbitQuant Hugging Face integrations
from diffusers import Krea2Pipeline

pipe = Krea2Pipeline.from_pretrained(
    "WaveCut/Krea-2-Turbo-OrbitQuant-W4A4",
    torch_dtype=torch.bfloat16,
    is_distilled=True,
).to("cuda")

image = pipe(
    prompt="A clean technical poster with readable labels",
    width=1024,
    height=1024,
    num_inference_steps=8,
    guidance_scale=0.0,
    generator=torch.Generator(device="cuda").manual_seed(0),
).images[0]

Fused transformer (OrbitQuant 0.11)

transformer/ stores the DiT as fused blocks, so the ordinary Diffusers loader returns it fused and the pipeline call above runs it; no extra file or runner is involved. Per block: one grouped Q|K|V|gate W4A4 GEMM with the sigmoid gate in its epilogue, a Q/K RMSNorm + RoPE kernel, the attention output projection, one SwiGLU gate|up GEMM, and the 28 protected ff.down projections as per-row INT8 weights for a W8A8 GEMM; the RMSNorm and modulation run in the activation quantization prologues and the gated residual updates in the GEMM epilogues. Attention computes Q·Kᵀ in INT8 and P·V in FP16 (the SageAttention v1 scheme); blocks 0 and 26, whose Q/K RMSNorm scales one channel far above the others, keep BF16 Flash attention. Batch-1 pipeline calls drop the padded prompt rows before the blocks (OrbitQuant 0.11.1), and the text-fusion stack runs once per prompt instead of once per step.

RTX 4060 Ti 16 GB, torch 2.10.0+cu128, Diffusers 80c7ed26, 1024×1024, 8 steps, guidance 0, components staged with enable_model_cpu_offload(), the first six prompts of benchmark/prompts.json:

Step times exclude the first step of an image, which also moves the DiT to the GPU. OrbitQuant 0.11.0 kept the padded prompt rows and ran the fused checkpoint about 7% slower than the 0.10 runner. The fused kernels round differently from the packed path, so the eight-step trajectory moves: the same seed keeps its subject, detail and lettering quality, while individual compositions can change, as they do between the packed path and the 0.10 runner. The 0.10 runner's transformer_fused/fused_blocks.safetensors held a second copy of the DiT weights and is no longer shipped; it remains at revision 56f91291 for orbitquant.runtime.krea2.

What is quantized

The DiT remains an OrbitQuant component. The sensitive output projection of each of its 28 main feed-forward blocks and the small img_in projection are protected from 4-bit quantization: the fused checkpoint stores the 28 ff.down weights as per-row INT8 (the 0.10 fused runtime's down="int8" mode), and img_in adds only 0.76 MiB of BF16 weights. Linear weight coverage counts W4A4 projections only. Attention projections, FFN gate/up projections, and the text-fusion stack remain packed W4A4. The universal policy also keeps two time-embedding projections and the final output projection in source precision. Embeddings, normalization parameters, convolutions, biases, the VAE, scheduler, and tokenizer are not 4-bit tensors.

Exact module names and parameter counts are in quantization_manifest.json.

Latency and VRAM

These measurements were made with the per-projection layout of revision 56f91291 (packed W4A4 layers, BF16 ff.down). They use the same A40, BF16 non-linear tensors, 2048×2048 output, 8 denoising steps, guidance 0, strict BF16 Flash SDPA, tiled VAE, lossless valid-token prompt compaction, and sequential Qwen→DiT component staging. Network downloads and PNG writes are excluded. No callback, latent copy, SageAttention, quantized attention, or full-model torch.compile is present in the timed path.

An intentionally empty Triton cache took 68.875 s for the first OrbitQuant image: 7.057 s for Qwen encode and 61.818 s for DiT plus VAE. Triton specializes some kernels for new prompt lengths, so the first encounter with another uncached shape can also pay a one-time compilation cost. The hot table is a second ten-prompt sweep after all ten tested shapes were compiled; it is not presented as cold-start latency.

The runtime prepares RP-BH activations directly as INT8 centroid surrogates and reuses that exact prepared tensor across adjacent Q/K/V/gate and SwiGLU gate/up projections. On the controlled eight-step denoiser benchmark this reduced the previous packed runtime from 49.470 s to 47.878 s without changing its output hash; the BF16 denoiser took 52.626 s. Final denoiser peak allocation was 12.59 GiB, versus 26.40 GiB for BF16. The small BF16 img_in protection did not erase the runtime gain.

The previously surprising ~44 GiB behavior is avoided by the included runner: it releases Qwen before loading the denoiser, compacts padded prompt lanes before unmasked Flash attention, does not retain decoded INT8 weights, and does not keep every pipeline component resident simultaneously. The text encoder itself is slower than BF16, but contributes only about 0.19 s to a hot 2048² request; the DiT dominates the net speedup.

Raw records are in benchmark/. nvidia-smi includes the CUDA context and non-Torch allocations; Torch figures use torch.cuda.max_memory_allocated().

Quality-preservation policy

The protected DiT policy was selected with component isolation, per-step latent trajectories, and native-size image inspection:

For the 34-token prompt, the quantized Qwen embeddings had cosine similarity 0.9674 to BF16 over valid tokens. This is a meaningful numerical difference, but it was not the source of the observed texture failure. The model therefore keeps Qwen fully W4A4.

The residual smooth-region problem was localized separately to DiT img_in. Mid-frequency residual ratios are normalized to BF16 (1.0 is closest): sky improved from 1.8497 to 1.0667, studio gradient from 1.3489 to 0.8760, and flat vector fills remained neutral (0.9853 → 1.0102). These are targeted diagnostic ratios, not a general image-quality score.

The diagnostic reports and causal matrices are retained under artifacts/diagnostics/.

Comparison protocol

All pairs use the same prompt, seed, 2048×2048 output, 8-step distilled schedule, and guidance 0. The set probes product texture, portraits, line art, fine diagrams, Latin and Cyrillic text, mixed scripts, and dense reflections. It is a practical deployment check, not an FID, CLIP, or human-preference benchmark.

Repository contents

• text_encoder/: fully packed OrbitQuant W4A4 Qwen3-VL component.
• transformer/: OrbitQuant W4A4 Krea DiT as fused blocks, with the 28 protected FFN outputs as per-row INT8 and BF16 img_in.
• vae/, scheduler/, and tokenizer/: pinned source components.
• artifacts/generations/original/ and artifacts/generations/orbitquant/: ten paired native-size PNG sets.
• assets/original_vs_orbitquant_w4a4.webp: full-resolution lossless comparison matrix.
• assets/original_vs_orbitquant_w4a4_preview.webp: reduced card preview linked to the original matrix.
• artifacts/diagnostics/: Qwen/DiT isolation and protection-ladder evidence.
• benchmark/: raw timings, memory measurements, prompts, environment, and kernel microbenchmarks.
• scripts/run_inference.py: recommended low-VRAM sequential runner.
• quantization_manifest.json, NOTICE, and MODIFICATIONS.md: provenance and modification details.

Limitations

• The first call pays a packed-kernel compilation cost; measure hot latency after warm-up.
• W4A4 may still change composition, contrast, fine typography, object counts, and small labels. Inspect the paired PNGs for your workload.
• The fused DiT needs CUDA and Triton. The runner's strict Flash SDPA (used by the text-fusion stack) requires a compatible CUDA/PyTorch build; the runner fails instead of silently selecting a slower high-memory attention path.
• Latency and memory depend on GPU architecture, drivers, Torch, Triton, resolution, and allocator state.
• This derivative inherits the source model's intended-use, safety, and license restrictions.

License and attribution

This is a modified derivative of Krea 2 Turbo. The upstream Krea 2 Community License Agreement is copied as LICENSE.pdf. The required upstream notice and modification notice are in NOTICE, with a technical summary in MODIFICATIONS.md. Review the agreement itself before use or redistribution; no endorsement by Krea is implied.