scripts/prepare_release.py

#!/usr/bin/env python3
"""Package benchmarks, full-resolution comparisons, and the practical model card."""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
import sys
from pathlib import Path
from typing import Any

from PIL import Image

SCRIPT_DIR = Path(__file__).resolve().parent
LAB_ROOT = SCRIPT_DIR.parent
sys.path.insert(0, str(LAB_ROOT))

from release_tools import create_full_resolution_matrix, validate_paired_records  # noqa: E402

RELEASE_NAME = "Krea-2-Turbo-OrbitQuant-W4A4"
SOURCE_ID = "krea/Krea-2-Turbo"
SOURCE_REVISION = "98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7"
REPO_ID = f"WaveCut/{RELEASE_NAME}"


def read_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line]


def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
    )


def tree_bytes(path: Path) -> int:
    return sum(item.stat().st_size for item in path.rglob("*") if item.is_file())


def gib(value: int | float | None) -> str:
    return "n/a" if value is None else f"{float(value) / 1024**3:.2f} GiB"


def seconds(value: int | float | None) -> str:
    return "n/a" if value is None else f"{float(value):.3f} s"


def memory(value: int | float | None) -> str:
    return "n/a" if value is None else f"{float(value) / 1024:.2f} GiB"


def delta(original: int | float | None, quantized: int | float | None) -> str:
    if original in (None, 0) or quantized is None:
        return "n/a"
    return f"{(float(quantized) / float(original) - 1) * 100:+.1f}%"


def component_rows(manifest: dict[str, Any]) -> str:
    rows = []
    for item in manifest["components"]:
        rows.append(
            "| `{component}` | `{class_name}` | {source} | {artifact} | {orbit} | {kept} | {coverage:.2%} |".format(
                component=item["component"],
                class_name=item["class_name"],
                source=gib(item["source_weight_bytes"]),
                artifact=gib(item["artifact_bytes"]),
                orbit=item["orbitquant_module_count"],
                kept=item["source_precision_linear_module_count"],
                coverage=item["linear_parameter_coverage"],
            )
        )
    return "\n".join(rows)


def benchmark_rows(original: dict[str, Any], quantized: dict[str, Any]) -> str:
    metrics = [
        ("Checkpoint load", "load_seconds", seconds),
        ("First generation", "first_generation_seconds", seconds),
        ("Hot generation median", "hot_generation_median_seconds", seconds),
        ("Hot generation mean", "hot_generation_mean_seconds", seconds),
        ("Generation peak, nvidia-smi", "gpu_peak_mb", memory),
        ("Generation peak, torch allocated", "torch_peak_mb", memory),
        ("Load peak, nvidia-smi", "load_gpu_peak_mb", memory),
        ("Load peak, torch allocated", "load_torch_peak_mb", memory),
    ]
    return "\n".join(
        f"| {label} | {formatter(original.get(key))} | {formatter(quantized.get(key))} | "
        f"{delta(original.get(key), quantized.get(key))} |"
        for label, key, formatter in metrics
    )


def prompt_rows(prompts: list[dict[str, Any]]) -> str:
    return "\n".join(
        f"| {index + 1:02d} | `{item['id']}` | {item['category']} | {61000 + index} |"
        for index, item in enumerate(prompts)
    )


def build_readme(
    manifest: dict[str, Any],
    original: dict[str, Any],
    quantized: dict[str, Any],
    prompts: list[dict[str, Any]],
    matrix: dict[str, Any],
) -> str:
    component_source = manifest["totals"]["source_weight_bytes"]
    component_quantized = manifest["totals"]["artifact_bytes"]
    release_size = tree_bytes(Path(manifest["release_path"]))
    hardware = quantized["hardware"]
    measured_runtime = ", ".join(quantized.get("effective_runtime_modes", [])) or "n/a"
    return f"""---
language:
- en
license: other
license_name: krea-2-community-license
license_link: https://cdn.jsdelivr.net/gh/krea-ai/krea-2@db3984fbc6e13b34c0064990fc2d95ac64d00058/assets/hf_samples/LICENSE.pdf
base_model:
- {SOURCE_ID}
base_model_relation: quantized
library_name: diffusers
pipeline_tag: text-to-image
tags:
- diffusers
- image-generation
- krea2
- orbitquant
- w4a4
- 4-bit
- quantized
---

# Krea 2 Turbo OrbitQuant W4A4

OrbitQuant W4A4 deployment checkpoint for [{SOURCE_ID}](https://huggingface.co/{SOURCE_ID}). Both transformer-class components are quantized: the Qwen3-VL text encoder and the Krea 2 diffusion transformer. The VAE, scheduler, tokenizer, embeddings, normalization parameters, convolutions, and policy-protected projections remain in source precision.

<a href="https://huggingface.co/{REPO_ID}/resolve/main/assets/original_vs_orbitquant_w4a4.webp"><img src="https://huggingface.co/{REPO_ID}/resolve/main/assets/original_vs_orbitquant_w4a4_preview.webp" alt="BF16 versus OrbitQuant W4A4, ten paired prompts" width="100%"></a>

The embedded image is a reduced preview linked to the **{matrix['matrix_size'][0]}×{matrix['matrix_size'][1]} lossless matrix**. Every source tile in the linked original remains at the model's native **2048×2048** benchmark output size; the builder adds labels and concatenates tiles without resizing. Individual PNGs are in [`artifacts/generations/`](https://huggingface.co/{REPO_ID}/tree/main/artifacts/generations).

## Quick facts

| Item | Value |
| --- | --- |
| Source revision | `{SOURCE_REVISION}` |
| Quantized components | `text_encoder` (`Qwen3VLModel`), `transformer` (`Krea2Transformer2DModel`) |
| Recipe | OrbitQuant W4A4, universal policy, RP-BH rotation, no calibration dataset |
| Runtime mode in artifact | `auto_fused`, automatic CUDA kernel selection |
| Runtime actually measured | `{measured_runtime}` with strict packed mode; no full-weight dequantization cache |
| Official Turbo settings tested | 2048×2048, 8 steps, guidance 0, distilled schedule (`mu=1.15`) |
| Hardware | {hardware} |
| Quantized learned-component storage | {gib(component_quantized)} vs {gib(component_source)} ({delta(component_source, component_quantized)}) |
| Complete runtime release tree | {gib(release_size)} |

## Install and run

The repository is publicly downloadable. Use and redistribution remain subject to the Krea 2 Community License Agreement copied below.

```bash
python -m venv .venv
source .venv/bin/activate
pip install -r https://huggingface.co/{REPO_ID}/resolve/main/runtime-requirements.txt
orbitquant kernels-install --build
```

`kernels-install` downloads a matching prebuilt kernel when available. `--build` permits an exact local build when the Torch/CUDA ABI is not among the published variants; that path requires the CUDA toolkit and `ninja`. The included runner sets `ORBITQUANT_STRICT_PACKED=1` so a missing packed kernel is an error instead of a silent BF16 fallback.

For the measured 2048×2048 path, use the included chunked-attention runner:

```bash
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
```

Direct Diffusers loading also works. Importing `orbitquant` registers both Hugging Face quantizers before the pipeline is loaded:

```python
import os
import torch
os.environ.setdefault("ORBITQUANT_STRICT_PACKED", "1")
import orbitquant  # registers the OrbitQuant HF integrations
from diffusers import Krea2Pipeline

pipe = Krea2Pipeline.from_pretrained(
    "{REPO_ID}",
    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]
image.save("krea2-orbitquant.png")
```

At 2048×2048, attention memory rather than packed linear weights is the dominant transient. `scripts/run_inference.py` installs the same 1024-query chunked native-attention path used for the measurements below.

## What is quantized

| Component | Class | Source weights | Packed artifact | OrbitQuant linears | Source-precision linears | Linear parameter coverage |
| --- | --- | ---: | ---: | ---: | ---: | ---: |
{component_rows(manifest)}

The three source-precision linear layers reported for the DiT are the two time-embedding projections and the final output projection selected by the universal policy. The text encoder's linear projections—including its visual branch—are packed. This table does not imply that embeddings, normalizations, or convolutions are 4-bit.

Machine-readable module names, parameter counts, source/artifact bytes, and per-component load/quantize peaks are in [`quantization_manifest.json`](https://huggingface.co/{REPO_ID}/blob/main/quantization_manifest.json).

## Measured latency and memory

Both variants were run in separate processes with the same pipeline class, BF16 non-linear tensors, prompts, seeds, 2048×2048 resolution, 8 steps, guidance 0, and chunked native attention. Checkpoint load measurements used already-downloaded local artifacts and exclude network transfer. The OrbitQuant process used strict packed mode and was accepted only when every executed quantized linear reported `native_packed_matmul` and no full dequantized weight cache remained. The first generation is reported separately from the median and mean of the remaining nine hot-path generations.

| Metric | Original BF16 | OrbitQuant W4A4 | Change |
| --- | ---: | ---: | ---: |
{benchmark_rows(original, quantized)}

Raw per-prompt records are available as both CSV and JSONL in [`benchmark/`](https://huggingface.co/{REPO_ID}/tree/main/benchmark). `nvidia-smi` peaks include the CUDA context and non-Torch allocations; Torch peaks are `torch.cuda.max_memory_allocated()`.

## Comparison protocol

| # | Prompt ID | Stress category | Seed |
| ---: | --- | --- | ---: |
{prompt_rows(prompts)}

The set covers product detail, portrait fidelity, public-domain style prompts, poster typography, a technical cutaway, long Latin text, long Cyrillic text, a mixed-script diagram, and a dense city scene. It is a practical paired deployment check, not an FID, CLIP, or human-preference benchmark.

## Repository contents

- `text_encoder/` and `transformer/`: clean-loadable OrbitQuant packed components.
- `vae/`, `scheduler/`, and `tokenizer/`: pinned source components.
- `artifacts/generations/original/`: ten original BF16 PNG outputs.
- `artifacts/generations/orbitquant/`: ten paired W4A4 PNG outputs.
- `assets/original_vs_orbitquant_w4a4_preview.webp`: reduced card preview linked to the original.
- `assets/original_vs_orbitquant_w4a4.webp`: lossless, full-resolution comparison matrix.
- `benchmark/`: prompts, raw metrics, summaries, matrix metadata, environment, and dummy preflight report.
- `scripts/`: inference, quantization, benchmark, and packaging scripts.
- `quantization_manifest.json`, `SHA256SUMS`, `NOTICE`, and `MODIFICATIONS.md`: provenance and integrity metadata.

## Limitations

- Quantization can change fine texture, typography, object counts, and composition. Inspect the paired originals for your target workload.
- Long text and small labels remain difficult for the source model and may change after quantization.
- Latency and memory depend on the GPU, driver, Torch, Triton, attention backend, resolution, and cache state.
- The comparison uses ten fixed prompts and seeds and should not be treated as a broad quality score.
- 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`](https://huggingface.co/{REPO_ID}/blob/main/LICENSE.pdf). The required upstream notice and modification notice are in [`NOTICE`](https://huggingface.co/{REPO_ID}/blob/main/NOTICE), with a technical summary in [`MODIFICATIONS.md`](https://huggingface.co/{REPO_ID}/blob/main/MODIFICATIONS.md). The agreement includes recipient-binding and attribution requirements for redistribution, a commercial-use revenue threshold, and a requirement to implement reasonable content filters for deployments. Review the agreement itself before use or redistribution; no endorsement by Krea is implied.
"""


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def write_checksums(release: Path) -> None:
    paths = sorted(
        item
        for item in release.rglob("*")
        if item.is_file() and item.name != "SHA256SUMS"
    )
    lines = [f"{sha256(path)}  {path.relative_to(release).as_posix()}" for path in paths]
    (release / "SHA256SUMS").write_text("\n".join(lines) + "\n", encoding="utf-8")


def audit(release: Path) -> dict[str, Any]:
    required = [
        "README.md",
        "LICENSE.pdf",
        "NOTICE",
        "MODIFICATIONS.md",
        "model_index.json",
        "runtime-requirements.txt",
        "quantization_manifest.json",
        "SHA256SUMS",
        "text_encoder/config.json",
        "transformer/config.json",
        "vae/config.json",
        "scheduler/scheduler_config.json",
        "tokenizer/tokenizer_config.json",
        "assets/original_vs_orbitquant_w4a4.webp",
        "assets/original_vs_orbitquant_w4a4_preview.webp",
        "benchmark/summary.json",
        "benchmark/original.metrics.csv",
        "benchmark/orbitquant.metrics.csv",
        "scripts/run_inference.py",
    ]
    missing = [name for name in required if not (release / name).exists()]
    forbidden = []
    for path in release.rglob("*"):
        relative = path.relative_to(release)
        if any(part in {".cache", "__pycache__", ".git", "logs", "tmp"} for part in relative.parts):
            forbidden.append(relative.as_posix())
        if path.is_file() and path.suffix in {".pyc", ".log"}:
            forbidden.append(relative.as_posix())
    original_pngs = list((release / "artifacts" / "generations" / "original").glob("*.png"))
    quantized_pngs = list((release / "artifacts" / "generations" / "orbitquant").glob("*.png"))
    configs = {
        name: read_json(release / name / "config.json").get("quantization_config", {}).get("quant_method")
        for name in ("text_encoder", "transformer")
    }
    report = {
        "release": str(release),
        "missing": missing,
        "forbidden": sorted(set(forbidden)),
        "original_png_count": len(original_pngs),
        "orbitquant_png_count": len(quantized_pngs),
        "quantization_methods": configs,
        "ok": (
            not missing
            and not forbidden
            and len(original_pngs) == 10
            and len(quantized_pngs) == 10
            and configs == {"text_encoder": "orbitquant", "transformer": "orbitquant"}
        ),
    }
    if not report["ok"]:
        raise RuntimeError(f"release audit failed: {report!r}")
    return report


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--scripts-source", type=Path, required=True)
    parser.add_argument("--dummy-report", type=Path)
    args = parser.parse_args()

    root = args.root.resolve()
    release = root / "release" / RELEASE_NAME
    scripts_source = args.scripts_source.resolve()
    if not release.is_dir():
        raise RuntimeError(f"release directory is missing: {release}")
    manifest_path = release / "quantization_manifest.json"
    if not manifest_path.is_file():
        raise RuntimeError("quantization is not complete")
    manifest = read_json(manifest_path)
    manifest["release_path"] = str(release)

    benchmark_root = root / "results" / "benchmark"
    original_dir = benchmark_root / "original"
    quantized_dir = benchmark_root / "orbitquant"
    original_summary = read_json(original_dir / "summary.json")
    quantized_summary = read_json(quantized_dir / "summary.json")
    original_rows = read_jsonl(original_dir / "metrics.jsonl")
    quantized_rows = read_jsonl(quantized_dir / "metrics.jsonl")
    validate_paired_records(original_rows, quantized_rows)
    prompts = read_json(root / "prompts.json")

    benchmark_release = release / "benchmark"
    benchmark_release.mkdir(parents=True, exist_ok=True)
    for label, source in (("original", original_dir), ("orbitquant", quantized_dir)):
        shutil.copy2(source / "summary.json", benchmark_release / f"{label}.summary.json")
        shutil.copy2(source / "metrics.csv", benchmark_release / f"{label}.metrics.csv")
        shutil.copy2(source / "metrics.jsonl", benchmark_release / f"{label}.metrics.jsonl")
    shutil.copy2(root / "prompts.json", benchmark_release / "prompts.json")
    environment_path = root / "state" / "environment.json"
    if environment_path.is_file():
        shutil.copy2(environment_path, benchmark_release / "environment.json")
    if args.dummy_report and args.dummy_report.is_file():
        shutil.copy2(args.dummy_report, benchmark_release / "dummy-preflight.json")

    artifact_root = release / "artifacts" / "generations"
    pairs = []
    for original_row, quantized_row in zip(original_rows, quantized_rows):
        filename = f"{int(original_row['prompt_idx']):02d}-{original_row['prompt_id']}.png"
        original_target = artifact_root / "original" / filename
        quantized_target = artifact_root / "orbitquant" / filename
        original_target.parent.mkdir(parents=True, exist_ok=True)
        quantized_target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(original_row["image_path"], original_target)
        shutil.copy2(quantized_row["image_path"], quantized_target)
        pairs.append(
            {
                "title": original_row["title"],
                "seed": original_row["seed"],
                "original": original_target,
                "orbitquant": quantized_target,
            }
        )

    matrix = create_full_resolution_matrix(
        pairs,
        release / "assets" / "original_vs_orbitquant_w4a4.webp",
        tile_size=(2048, 2048),
        prompt_pairs_per_row=2,
        label_height=96,
    )
    matrix_path = release / "assets" / "original_vs_orbitquant_w4a4.webp"
    preview_path = release / "assets" / "original_vs_orbitquant_w4a4_preview.webp"
    with Image.open(matrix_path) as full_matrix:
        preview = full_matrix.convert("RGB")
        preview.thumbnail((2400, 10000), Image.Resampling.LANCZOS)
        preview.save(preview_path, format="WEBP", quality=88, method=6)
        matrix["preview_path"] = str(preview_path)
        matrix["preview_size"] = list(preview.size)
    write_json(release / "assets" / "original_vs_orbitquant_w4a4.json", matrix)
    combined_summary = {
        "source_model_id": SOURCE_ID,
        "source_revision": SOURCE_REVISION,
        "repo_id": REPO_ID,
        "settings": original_summary["settings"],
        "original": original_summary,
        "orbitquant": quantized_summary,
    }
    write_json(benchmark_release / "summary.json", combined_summary)

    release_scripts = release / "scripts"
    release_scripts.mkdir(parents=True, exist_ok=True)
    for name in (
        "run_inference.py",
        "quantize_full_components.py",
        "benchmark_full_pipeline.py",
        "prepare_release.py",
    ):
        shutil.copy2(scripts_source / name, release_scripts / name)
    shutil.copy2(LAB_ROOT / "release_tools.py", release_scripts / "release_tools.py")

    (release / "README.md").write_text(
        build_readme(manifest, original_summary, quantized_summary, prompts, matrix),
        encoding="utf-8",
    )
    manifest.pop("release_path", None)
    manifest["release_bytes_before_checksums"] = tree_bytes(release)
    manifest["comparison_matrix"] = "assets/original_vs_orbitquant_w4a4.webp"
    manifest["raw_generation_artifacts"] = "artifacts/generations"
    manifest["benchmark_summary"] = "benchmark/summary.json"
    write_json(manifest_path, manifest)
    write_checksums(release)
    report = audit(release)
    write_json(root / "state" / "release_audit.json", report)
    print(json.dumps(report, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())