scripts/quantize_full_components.py

#!/usr/bin/env python3
"""Quantize both learned transformer components of Krea 2 Turbo with OrbitQuant."""

from __future__ import annotations

import argparse
import gc
import json
import os
import platform
import shutil
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import psutil
import torch
from huggingface_hub import HfApi, hf_hub_download, snapshot_download

import orbitquant
from orbitquant import recipe
from orbitquant.adaln import RTNInt4Linear
from orbitquant.layers import OrbitQuantLinear

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


@dataclass(frozen=True)
class Component:
    name: str
    framework: str
    class_name: str


COMPONENTS = (
    Component("transformer", "diffusers", "Krea2Transformer2DModel"),
    Component("text_encoder", "transformers", "Qwen3VLModel"),
)


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 read_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


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


def clean_cuda() -> None:
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()


def gpu_snapshot() -> dict[str, Any]:
    free, total = torch.cuda.mem_get_info()
    return {
        "device": torch.cuda.get_device_name(0),
        "capability": list(torch.cuda.get_device_capability(0)),
        "free_bytes": free,
        "total_bytes": total,
        "allocated_bytes": torch.cuda.memory_allocated(),
        "reserved_bytes": torch.cuda.memory_reserved(),
        "peak_allocated_bytes": torch.cuda.max_memory_allocated(),
        "peak_reserved_bytes": torch.cuda.max_memory_reserved(),
    }


def component_class(component: Component) -> type[torch.nn.Module]:
    if component.framework == "diffusers":
        import diffusers

        return getattr(diffusers, component.class_name)
    import transformers

    return getattr(transformers, component.class_name)


def source_weight_bytes(component: Component, cache_dir: Path) -> int:
    index_name = (
        f"{component.name}/diffusion_pytorch_model.safetensors.index.json"
        if component.framework == "diffusers"
        else f"{component.name}/model.safetensors.index.json"
    )
    try:
        path = Path(
            hf_hub_download(
                SOURCE_ID,
                index_name,
                revision=SOURCE_REVISION,
                cache_dir=cache_dir,
            )
        )
        total_size = read_json(path).get("metadata", {}).get("total_size")
        if total_size is not None:
            return int(total_size)
    except Exception:
        pass

    file_name = (
        f"{component.name}/diffusion_pytorch_model.safetensors"
        if component.framework == "diffusers"
        else f"{component.name}/model.safetensors"
    )
    paths = HfApi().get_paths_info(
        SOURCE_ID, file_name, revision=SOURCE_REVISION, repo_type="model"
    )
    if len(paths) != 1 or getattr(paths[0], "size", None) is None:
        raise RuntimeError(f"could not determine source size for {component.name}")
    return int(paths[0].size)


def module_inventory(model: torch.nn.Module) -> dict[str, Any]:
    orbit_modules: list[str] = []
    adaln_modules: list[str] = []
    source_precision_modules: list[str] = []
    quantized_weight_parameters = 0
    skipped_weight_parameters = 0
    packed_state_bytes = 0

    for name, module in model.named_modules():
        if isinstance(module, OrbitQuantLinear):
            orbit_modules.append(name)
            quantized_weight_parameters += module.in_features * module.out_features
            packed_state_bytes += sum(
                value.numel() * value.element_size()
                for value in module.state_dict().values()
            )
        elif isinstance(module, RTNInt4Linear):
            adaln_modules.append(name)
            quantized_weight_parameters += module.in_features * module.out_features
            packed_state_bytes += sum(
                value.numel() * value.element_size()
                for value in module.state_dict().values()
            )
        elif isinstance(module, torch.nn.Linear):
            source_precision_modules.append(name)
            skipped_weight_parameters += module.weight.numel()

    total = quantized_weight_parameters + skipped_weight_parameters
    cache_count = sum(
        isinstance(module, OrbitQuantLinear)
        and getattr(module, "_dequantized_weight_cache", None) is not None
        for module in model.modules()
    )
    return {
        "orbitquant_module_count": len(orbit_modules),
        "adaln_int4_module_count": len(adaln_modules),
        "source_precision_linear_module_count": len(source_precision_modules),
        "orbitquant_modules": orbit_modules,
        "adaln_int4_modules": adaln_modules,
        "source_precision_linear_modules": source_precision_modules,
        "quantized_linear_weight_parameters": quantized_weight_parameters,
        "source_precision_linear_weight_parameters": skipped_weight_parameters,
        "linear_weight_parameters": total,
        "linear_parameter_coverage": quantized_weight_parameters / total if total else 0.0,
        "packed_module_state_bytes": packed_state_bytes,
        "full_dequantized_cache_count": cache_count,
    }


def load_quantized(component: Component, cache_dir: Path) -> torch.nn.Module:
    cls = component_class(component)
    config = recipe(
        "w4a4",
        target_policy="universal",
        runtime_mode="auto_fused",
        activation_kernel_backend="auto",
    )
    kwargs: dict[str, Any] = {
        "revision": SOURCE_REVISION,
        "subfolder": component.name,
        "cache_dir": cache_dir,
        "quantization_config": config,
        "low_cpu_mem_usage": True,
    }
    if component.framework == "diffusers":
        kwargs["quantization_device"] = "cuda"
        kwargs["torch_dtype"] = torch.bfloat16
    else:
        kwargs["dtype"] = torch.bfloat16
    model = cls.from_pretrained(SOURCE_ID, **kwargs)
    model.eval().requires_grad_(False)
    return model


def copy_source_metadata(release: Path, cache_dir: Path) -> None:
    snapshot = Path(
        snapshot_download(
            SOURCE_ID,
            revision=SOURCE_REVISION,
            cache_dir=cache_dir,
            allow_patterns=(
                "model_index.json",
                "scheduler/*",
                "tokenizer/*",
                "vae/*",
                "LICENSE.pdf",
            ),
        )
    )
    for relative in ("model_index.json", "scheduler", "tokenizer", "vae", "LICENSE.pdf"):
        source = snapshot / relative
        target = release / relative
        if source.is_dir():
            shutil.copytree(source, target, dirs_exist_ok=True)
        else:
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(source, target)


def write_legal_and_runtime_files(release: Path) -> None:
    (release / "NOTICE").write_text(
        "Krea 2 is licensed under the Krea 2 Community License Agreement. "
        "For more information, visit https://krea.ai/krea-2-licensing.\n\n"
        "Modified distribution: the Qwen3-VL text encoder and Krea 2 diffusion "
        "transformer linear layers were converted to OrbitQuant W4A4. This "
        "distribution is not endorsed by Krea.\n",
        encoding="utf-8",
    )
    (release / "MODIFICATIONS.md").write_text(
        "# Modifications\n\n"
        "The learned linear projections in `text_encoder` (`Qwen3VLModel`) and "
        "`transformer` (`Krea2Transformer2DModel`) were converted from the pinned "
        "Krea 2 Turbo checkpoint to OrbitQuant W4A4 packed weights. The universal "
        "policy keeps explicitly protected time-embedding and final-output "
        "projections in source precision. Embeddings, normalization parameters, "
        "convolutions, biases, VAE, scheduler, and tokenizer are not quantized.\n",
        encoding="utf-8",
    )
    (release / "runtime-requirements.txt").write_text(
        "orbitquant[hf,kernels] @ git+https://github.com/iamwavecut/OrbitQuant.git@"
        f"{ORBITQUANT_REVISION}\n"
        f"diffusers=={DIFFUSERS_VERSION}\n"
        "transformers>=5.13,<6\n"
        "huggingface_hub>=1.22,<2\n"
        "accelerate\n"
        "safetensors\n",
        encoding="utf-8",
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--component", action="append", choices=[item.name for item in COMPONENTS])
    parser.add_argument("--keep-source-cache", action="store_true")
    args = parser.parse_args()

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required")
    root = args.root.resolve()
    cache_root = root / "cache" / "huggingface"
    release = root / "release" / RELEASE_NAME
    state_dir = root / "state" / "quantization"
    release.mkdir(parents=True, exist_ok=True)
    state_dir.mkdir(parents=True, exist_ok=True)
    copy_source_metadata(release, cache_root / "metadata")
    write_legal_and_runtime_files(release)

    environment = {
        "source_model_id": SOURCE_ID,
        "source_revision": SOURCE_REVISION,
        "repo_id": REPO_ID,
        "orbitquant_version": orbitquant.__version__,
        "orbitquant_revision": ORBITQUANT_REVISION,
        "diffusers_version": __import__("diffusers").__version__,
        "transformers_version": __import__("transformers").__version__,
        "huggingface_hub_version": __import__("huggingface_hub").__version__,
        "torch": torch.__version__,
        "cuda": torch.version.cuda,
        "python": platform.python_version(),
        "hostname": platform.node(),
        "gpu": gpu_snapshot(),
    }
    write_json(root / "state" / "environment.json", environment)

    selected = [item for item in COMPONENTS if not args.component or item.name in args.component]
    for component in selected:
        state_path = state_dir / f"{component.name}.json"
        target_dir = release / component.name
        if state_path.is_file() and target_dir.is_dir():
            previous = read_json(state_path)
            if previous.get("status") == "complete":
                print(json.dumps({"component": component.name, "status": "already_complete"}))
                continue

        clean_cuda()
        started = time.perf_counter()
        rss_before = psutil.Process().memory_info().rss
        component_cache = cache_root / component.name
        model = load_quantized(component, component_cache)
        torch.cuda.synchronize()
        load_seconds = time.perf_counter() - started
        component_inventory = module_inventory(model)
        if component_inventory["orbitquant_module_count"] <= 0:
            raise RuntimeError(f"{component.name} produced no OrbitQuantLinear modules")
        if component_inventory["full_dequantized_cache_count"]:
            raise RuntimeError(f"{component.name} retained full dequantized caches")

        if target_dir.exists():
            shutil.rmtree(target_dir)
        save_started = time.perf_counter()
        model.save_pretrained(target_dir, safe_serialization=True, max_shard_size="4GB")
        save_seconds = time.perf_counter() - save_started
        hf_quantizer = getattr(model, "hf_quantizer", None)
        result = {
            "status": "complete",
            "component": component.name,
            "framework": component.framework,
            "class_name": component.class_name,
            "component_mode": "orbitquant_w4a4",
            "source_weight_bytes": source_weight_bytes(component, component_cache),
            "artifact_bytes": tree_bytes(target_dir),
            "load_and_quantize_seconds": load_seconds,
            "save_seconds": save_seconds,
            "wall_seconds": time.perf_counter() - started,
            "rss_before_bytes": rss_before,
            "rss_after_bytes": psutil.Process().memory_info().rss,
            "released_source_tensor_bytes": getattr(hf_quantizer, "released_source_tensor_bytes", None),
            "source_page_release_failures": getattr(hf_quantizer, "source_page_release_failures", None),
            "gpu": gpu_snapshot(),
            **component_inventory,
        }
        write_json(state_path, result)
        print(json.dumps({key: value for key, value in result.items() if not key.endswith("_modules")}))
        del model
        clean_cuda()
        if not args.keep_source_cache:
            shutil.rmtree(component_cache, ignore_errors=True)

    completed = []
    for component in COMPONENTS:
        path = state_dir / f"{component.name}.json"
        if path.is_file() and read_json(path).get("status") == "complete":
            completed.append(read_json(path))
    if len(completed) != len(COMPONENTS):
        print(json.dumps({"status": "partial", "completed_components": [item["component"] for item in completed]}))
        return 0

    totals = {
        "source_weight_bytes": sum(item["source_weight_bytes"] for item in completed),
        "artifact_bytes": sum(item["artifact_bytes"] for item in completed),
        "quantized_linear_weight_parameters": sum(
            item["quantized_linear_weight_parameters"] for item in completed
        ),
        "source_precision_linear_weight_parameters": sum(
            item["source_precision_linear_weight_parameters"] for item in completed
        ),
        "orbitquant_module_count": sum(item["orbitquant_module_count"] for item in completed),
        "adaln_int4_module_count": sum(item["adaln_int4_module_count"] for item in completed),
        "source_precision_linear_module_count": sum(
            item["source_precision_linear_module_count"] for item in completed
        ),
    }
    linear_total = (
        totals["quantized_linear_weight_parameters"]
        + totals["source_precision_linear_weight_parameters"]
    )
    totals["linear_parameter_coverage"] = (
        totals["quantized_linear_weight_parameters"] / linear_total if linear_total else 0.0
    )
    totals["release_bytes"] = tree_bytes(release)
    manifest = {
        "artifact_format": "orbitquant-multicomponent-v1",
        "source_model_id": SOURCE_ID,
        "source_revision": SOURCE_REVISION,
        "source_license": "krea-2-community-license-agreement",
        "repo_id": REPO_ID,
        "visibility": "public-ungated",
        "quant_method": "orbitquant",
        "recipe": "w4a4-universal",
        "weight_bits": 4,
        "activation_bits": 4,
        "w4a4_components": ["text_encoder", "transformer"],
        "source_precision_components": ["vae", "scheduler", "tokenizer"],
        "calibration_data": None,
        "orbitquant_version": orbitquant.__version__,
        "orbitquant_revision": ORBITQUANT_REVISION,
        "diffusers_version": DIFFUSERS_VERSION,
        "components": completed,
        "totals": totals,
    }
    write_json(release / "quantization_manifest.json", manifest)
    write_json(root / "state" / "quantization_complete.json", manifest)
    print(json.dumps({"status": "all_components_complete", "manifest": str(release / "quantization_manifest.json")}))
    return 0


if __name__ == "__main__":
    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
    raise SystemExit(main())