svdquant_quantize.py

"""Run the quantizer from inside ComfyUI instead of a terminal.

Same code path as ``quantize_krea2.py`` on the command line -- this imports `convert` rather
than reimplementing it, so there is one quantizer, not two that drift.

The honest caveats, which the node's DESCRIPTION repeats because people do not read module
docstrings:

* It blocks the queue. 54 s for a single-shot split, ~5.7 min with the refinement loop, on a
  3090. Nothing else runs meanwhile.
* It needs the GPU to itself, so it unloads whatever ComfyUI is holding first. Your next
  generation will pay a reload.
* It writes ~8 GB.
"""

from __future__ import annotations

import logging
import os
import shutil

import comfy.model_management
import comfy.utils
import folder_paths

from .quantize_krea2 import (
    SAMPLER_HINTS,
    convert,
    derive_out_path,
    resolve_format,
)
from .svdquant_diag import _CATEGORY

# A source has to be dequantized to bf16 before it can be requantized, so leave a margin over
# the ~8 GB output rather than exactly it.
_MIN_FREE_BYTES = 12 * 1024 ** 3


def _free_bytes(path: str) -> int:
    return shutil.disk_usage(os.path.dirname(os.path.abspath(path))).free


class Krea2SVDQuantQuantize:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "source_model": (folder_paths.get_filename_list("diffusion_models"), {
                    "tooltip": "The BF16 Krea 2 checkpoint to quantize (~24 GB). An "
                               "already-quantized file cannot be used as a source, except "
                               "FP8, which is unpacked back to BF16 first.",
                }),
                "format": (["svdq", "w4a4", "int8", "fp8"], {
                    "default": "svdq",
                    "tooltip": "svdq: 4-bit weights and activations plus a low-rank bf16 "
                               "correction branch. w4a4: the same without the branch - "
                               "smaller and ~9%% faster per step. int8: the most faithful "
                               "option and still ~2x fp8 on Ampere. fp8: storage only.",
                }),
                "rank": ("INT", {
                    "default": 64, "min": 8, "max": 1024, "step": 8,
                    "tooltip": "svdq only: size of the low-rank branch. Only pays off with "
                               "refine_iters > 0. Without a LoRA, 64 / 128 / 256 measure the "
                               "same, so 64 is enough. With a LoRA loaded, 256 wins clearly "
                               "and 64 loses most of its advantage - so 256 if you use LoRAs.",
                }),
                "rank_alloc": (["uniform", "gqa"], {
                    "default": "uniform",
                    "tooltip": "svdq only: how the rank budget is spread across the eight "
                               "projection types. Same file size either way. uniform gives "
                               "every layer the same rank. gqa moves the budget to attn.wk / "
                               "attn.wv, which absorb ~2x the quantization error at a third of "
                               "the branch cost because Krea 2 has only 12 kv heads. Measured "
                               "and it does not pay: LPIPS 0.3523 vs 0.3403 for uniform, 5 of "
                               "10 prompts better, no mean effect. It does halve the spread "
                               "across prompts and improve the worst one. Leave on uniform "
                               "unless you are re-testing that.",
                }),
                "refine_iters": ("INT", {
                    "default": 100, "min": 0, "max": 200,
                    "tooltip": "svdq only. 0 is a single-shot SVD split (~54s); 100 refines "
                               "the branch against the quantization error and early-stops "
                               "(~5.7min). Keep this on if rank > 16: refinement is what "
                               "makes rank behave. Without it, raising rank costs file size "
                               "and buys nothing measurable.",
                }),
                "groupsize": ("INT", {
                    "default": 256, "min": 32, "max": 1024, "step": 32,
                    "tooltip": "convrot rotation group size. Unused for fp8.",
                }),
                "variant": (["turbo", "base", "unknown"], {
                    "default": "unknown",
                    "tooltip": "Which Krea 2 release this is. Affects only the output "
                               "filename and the recorded metadata - quantization is "
                               "identical; what differs is the sampler settings afterwards.",
                }),
                "output_name": ("STRING", {
                    "default": "",
                    "tooltip": "Filename inside models/diffusion_models/. Leave empty to "
                               "derive it from the variant and format.",
                }),
                "overwrite": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "Off means an existing file of the same name is an error "
                               "rather than 8 GB written over your last run.",
                }),
            },
            # Optional so workflows saved before this input existed keep validating.
            "optional": {
                "act_stats": ("STRING", {
                    "default": "",
                    "tooltip": "svdq only: an activation-statistics file from the Capture "
                               "nodes (a bare filename is looked up in ComfyUI/output/). "
                               "Fits the low-rank branch against measured per-channel "
                               "activation energy instead of assuming it is uniform. Free at "
                               "inference and the best-measured setting here - LPIPS to BF16 "
                               "0.3378 to 0.2825 with no LoRA. Empty means the plain objective.",
                }),
            },
        }

    @classmethod
    def IS_CHANGED(cls, *args, **kwargs):
        # Writes an ~8 GB file as its side effect, so a cached summary would claim a file
        # exists that the user may have since deleted. Re-queueing is cheap to refuse
        # (overwrite=False still fails fast) and expensive to get wrong.
        return float("nan")

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("summary",)
    OUTPUT_TOOLTIPS = ("Where the checkpoint was written, and what went into it.",)
    OUTPUT_NODE = True
    FUNCTION = "run"
    CATEGORY = _CATEGORY
    TITLE = "Krea2 SVDQuant Quantize"
    DESCRIPTION = ("Builds a quantized Krea 2 checkpoint from a BF16 one, without leaving "
                   "ComfyUI. BLOCKS THE QUEUE while it runs (54s to ~6min), unloads any "
                   "loaded model to free the GPU, and writes ~8 GB. Load the result with the "
                   "Krea2 SVDQuant W4A4 Loader (svdq) or the stock UNETLoader (w4a4/int8/fp8).")

    def run(self, source_model, format, rank, rank_alloc, refine_iters, groupsize, variant,
            output_name, overwrite, act_stats=""):
        src = folder_paths.get_full_path_or_raise("diffusion_models", source_model)

        # `rank` always carries a value from the widget, so "was it set?" cannot be inferred
        # the way argparse does it. Non-svdq formats simply ignore it here rather than
        # erroring, which is the friendlier reading of a dropdown the user cannot un-set.
        fmt, rank = resolve_format(format, rank, rank_was_set=False)

        # Same validation the CLI does for --act-stats: a typed path that silently did nothing
        # would produce a checkpoint indistinguishable from a plain one.
        stats_path = act_stats.strip() or None
        if stats_path is not None:
            if format != "svdq":
                raise RuntimeError("act_stats only applies to format 'svdq': it weights the "
                                   "low-rank branch, and the other formats have no branch.")
            if not os.path.isabs(stats_path):
                stats_path = os.path.join(folder_paths.get_output_directory(), stats_path)
            if not os.path.isfile(stats_path):
                raise RuntimeError("act_stats file not found: {}".format(stats_path))

        if output_name.strip():
            name = output_name.strip()
            if not name.endswith(".safetensors"):
                name += ".safetensors"
            dst = os.path.join(os.path.dirname(src), name)
        else:
            dst, note = derive_out_path(src, format, rank, variant, rank_alloc, stats_path)
            if note:
                logging.info("[krea2-svdquant] %s", note)

        if os.path.exists(dst) and not overwrite:
            raise RuntimeError(
                "{} already exists. Enable 'overwrite', or set a different output_name."
                .format(dst))

        free = _free_bytes(dst)
        if free < _MIN_FREE_BYTES:
            raise RuntimeError(
                "only {:.1f} GB free on the drive holding {}; quantizing needs roughly {:.0f} "
                "GB of headroom.".format(free / 1024 ** 3, os.path.dirname(dst),
                                         _MIN_FREE_BYTES / 1024 ** 3))

        # Quantization wants the card to itself. Without this it competes with whatever the
        # last generation left resident and OOMs on the dequantize-to-bf16 step.
        comfy.model_management.unload_all_models()
        comfy.model_management.soft_empty_cache()

        pbar = comfy.utils.ProgressBar(1)
        state = {"total": 0}

        def progress(done, total, message):
            if total != state["total"]:
                state["total"] = total
                pbar.total = total
            pbar.update_absolute(done, total)

        logging.info("[krea2-svdquant] quantizing %s -> %s (format %s, rank %s/%s, "
                     "refine_iters %s, act_stats %s)", src, dst, fmt, rank, rank_alloc,
                     refine_iters if rank else 0, stats_path or "none")
        summary = convert(src, dst, fmt, groupsize, "cuda", rank, refine_iters,
                          variant=variant, progress_cb=progress,
                          rank_alloc=rank_alloc if rank else "uniform",
                          act_stats=stats_path)

        hint = SAMPLER_HINTS.get(variant)
        loader = ("Krea2 SVDQuant W4A4 Loader" if rank else "the stock UNETLoader")
        text = "\n".join(x for x in (
            summary, "Load it with {}.".format(loader), hint) if x)
        logging.info("[krea2-svdquant] %s", text)
        return {"ui": {"text": [text]}, "result": (text,)}


NODE_CLASS_MAPPINGS = {"Krea2SVDQuantQuantize": Krea2SVDQuantQuantize}
NODE_DISPLAY_NAME_MAPPINGS = {"Krea2SVDQuantQuantize": "Krea2 SVDQuant Quantize"}