#!/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())