standalone_inference.py

#!/usr/bin/env python3
"""Standalone inference for quantized RUM FLUX.2 Klein 4B checkpoints.

This file intentionally contains the RUM architecture and comfy-kitchen
checkpoint loader so users do not need to clone the RUM or quantizer projects.
It still requires Python packages and the upstream FLUX.2 pipeline components.

Install:
    pip install torch diffusers==0.37.1 transformers safetensors \
        huggingface-hub pillow comfy-kitchen==0.2.21

Lower-quality Qwen-only trial:
    python standalone_inference.py \
        --checkpoint rum-flux2-int8-convrot-adaround-realcal-hq.safetensors \
        --base-model models/FLUX.2-klein-base-4B \
        --qwen-only \
        --prompt "1girl, kisaki (blue archive), holding baozi, eating, indoors"

Full RUM dual-text inference:
    python standalone_inference.py \
        --checkpoint rum-flux2-int8-convrot-adaround-realcal-hq.safetensors \
        --base-model models/FLUX.2-klein-base-4B \
        --teacher-model models/waiNSFWIllustrious_v140.safetensors \
        --prompt "1girl, kisaki (blue archive), holding baozi, eating, indoors"
"""

from __future__ import annotations

import argparse
import gc
import json
import platform
import time
import warnings
from importlib.metadata import version
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from safetensors import safe_open

try:
    from comfy_kitchen.tensor import (
        QuantizedTensor,
        TensorCoreConvRotW4A4Layout,
        TensorWiseINT8Layout,
    )
except ImportError as exc:  # pragma: no cover - import failure is environment-specific
    raise SystemExit(
        'Missing comfy-kitchen. Install the pinned backend with: '
        'pip install "comfy-kitchen==0.2.21"'
    ) from exc

from diffusers import Flux2KleinPipeline, StableDiffusionXLPipeline
from diffusers.models._modeling_parallel import (
    ContextParallelInput,
    ContextParallelOutput,
)
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
from diffusers.utils import apply_lora_scale


COMFY_KITCHEN_VERSION = "0.2.21"
DEFAULT_BASE_MODEL = "models/FLUX.2-klein-base-4B"
DEFAULT_TEACHER_MODEL = "models/waiNSFWIllustrious_v140.safetensors"
DEFAULT_PROMPT = (
    "1girl, kisaki (blue archive), general, holding baozi, eating, indoors, "
    "momoko (momopoco), newest"
)


class RUMFlux2Transformer(Flux2Transformer2DModel):
    """FLUX.2 Klein Transformer with RUM's second 2048-wide text projection."""

    _supports_gradient_checkpointing = True
    _no_split_modules = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
    _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
    _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
    _cp_plan = {
        "": {
            "hidden_states": ContextParallelInput(
                split_dim=1, expected_dims=3, split_output=False
            ),
            "encoder_hidden_states": ContextParallelInput(
                split_dim=1, expected_dims=3, split_output=False
            ),
            "img_ids": ContextParallelInput(
                split_dim=1, expected_dims=3, split_output=False
            ),
            "txt_ids": ContextParallelInput(
                split_dim=1, expected_dims=3, split_output=False
            ),
        },
        "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
    }

    def __init__(
        self,
        patch_size: int = 1,
        in_channels: int = 128,
        out_channels: int | None = None,
        num_layers: int = 8,
        num_single_layers: int = 48,
        attention_head_dim: int = 128,
        num_attention_heads: int = 48,
        joint_attention_dim: int = 15360,
        timestep_guidance_channels: int = 256,
        mlp_ratio: float = 3.0,
        axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32),
        rope_theta: int = 2000,
        eps: float = 1e-6,
        guidance_embeds: bool = True,
    ) -> None:
        super().__init__(
            patch_size=patch_size,
            in_channels=in_channels,
            out_channels=out_channels,
            num_layers=num_layers,
            num_single_layers=num_single_layers,
            attention_head_dim=attention_head_dim,
            num_attention_heads=num_attention_heads,
            joint_attention_dim=joint_attention_dim,
            timestep_guidance_channels=timestep_guidance_channels,
            mlp_ratio=mlp_ratio,
            axes_dims_rope=axes_dims_rope,
            rope_theta=rope_theta,
            eps=eps,
            guidance_embeds=guidance_embeds,
        )
        self.context_embedder_2 = nn.Linear(2048, self.inner_dim, bias=False)

    @apply_lora_scale("joint_attention_kwargs")
    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor | None = None,
        timestep: torch.LongTensor | None = None,
        img_ids: torch.Tensor | None = None,
        txt_ids: torch.Tensor | None = None,
        guidance: torch.Tensor | None = None,
        joint_attention_kwargs: dict[str, Any] | None = None,
        return_dict: bool = True,
    ) -> torch.Tensor | Transformer2DModelOutput:
        if encoder_hidden_states is None:
            raise ValueError("encoder_hidden_states is required")
        num_txt_tokens = encoder_hidden_states.shape[1]
        timestep = timestep.to(hidden_states.dtype) * 1000
        if guidance is not None:
            guidance = guidance.to(hidden_states.dtype) * 1000
        temb = self.time_guidance_embed(timestep, guidance)
        double_stream_mod_img = self.double_stream_modulation_img(temb)
        double_stream_mod_txt = self.double_stream_modulation_txt(temb)
        single_stream_mod = self.single_stream_modulation(temb)
        hidden_states = self.x_embedder(hidden_states)

        if encoder_hidden_states.shape[1] > 200:
            qwen = encoder_hidden_states[:, :200]
            clip = encoder_hidden_states[:, 200:, :2048]
            encoder_hidden_states = torch.cat(
                [self.context_embedder(qwen), self.context_embedder_2(clip)], dim=1
            )
        else:
            encoder_hidden_states = self.context_embedder(encoder_hidden_states)

        if img_ids.ndim == 3:
            img_ids = img_ids[0]
        if txt_ids.ndim == 3:
            txt_ids = txt_ids[0]
        image_rotary_emb = self.pos_embed(img_ids)
        text_rotary_emb = self.pos_embed(txt_ids)
        concat_rotary_emb = (
            torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
            torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
        )

        for block in self.transformer_blocks:
            encoder_hidden_states, hidden_states = block(
                hidden_states=hidden_states,
                encoder_hidden_states=encoder_hidden_states,
                temb_mod_img=double_stream_mod_img,
                temb_mod_txt=double_stream_mod_txt,
                image_rotary_emb=concat_rotary_emb,
                joint_attention_kwargs=joint_attention_kwargs,
            )
        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
        for block in self.single_transformer_blocks:
            hidden_states = block(
                hidden_states=hidden_states,
                encoder_hidden_states=None,
                temb_mod=single_stream_mod,
                image_rotary_emb=concat_rotary_emb,
                joint_attention_kwargs=joint_attention_kwargs,
            )
        hidden_states = hidden_states[:, num_txt_tokens:, ...]
        hidden_states = self.norm_out(hidden_states, temb)
        output = self.proj_out(hidden_states)
        if not return_dict:
            return (output,)
        return Transformer2DModelOutput(sample=output)


class NativeINT8Linear(nn.Module):
    """Native comfy-kitchen W8A8 Linear restored from serialized tensors."""

    def __init__(
        self,
        qweight: torch.Tensor,
        weight_scale: torch.Tensor,
        *,
        bias: torch.Tensor | None,
        convrot_group_size: int | None,
    ) -> None:
        super().__init__()
        if qweight.ndim != 2 or qweight.dtype != torch.int8:
            raise ValueError("INT8 qweight must be a two-dimensional int8 tensor")
        if weight_scale.shape not in (
            torch.Size([]),
            torch.Size([qweight.shape[0]]),
            torch.Size([qweight.shape[0], 1]),
        ):
            raise ValueError("INT8 scale must be scalar or contain one value per row")
        scale = (
            weight_scale.reshape(qweight.shape[0], 1)
            if weight_scale.numel() == qweight.shape[0]
            else weight_scale
        )
        params = TensorWiseINT8Layout.Params(
            scale=scale.contiguous(),
            orig_dtype=torch.bfloat16,
            orig_shape=tuple(qweight.shape),
            is_weight=True,
            convrot=convrot_group_size is not None,
            convrot_groupsize=convrot_group_size or 256,
        )
        self.in_features = qweight.shape[1]
        self.out_features = qweight.shape[0]
        self.register_buffer(
            "weight",
            QuantizedTensor(qweight.contiguous(), "TensorWiseINT8Layout", params),
        )
        self.register_buffer("bias", bias.contiguous() if bias is not None else None)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        bias = self.bias.to(inputs.dtype) if self.bias is not None else None
        return F.linear(inputs, self.weight, bias)


class NativeINT4ConvRotLinear(nn.Module):
    """Native comfy-kitchen packed ConvRot W4A4 Linear."""

    def __init__(
        self,
        qweight: torch.Tensor,
        weight_scale: torch.Tensor,
        *,
        orig_shape: tuple[int, int],
        bias: torch.Tensor | None,
        convrot_group_size: int,
        quant_group_size: int,
        linear_dtype: str,
    ) -> None:
        super().__init__()
        expected_shape = (orig_shape[0], orig_shape[1] // 2)
        if qweight.dtype != torch.int8 or tuple(qweight.shape) != expected_shape:
            raise ValueError(
                f"packed INT4 qweight shape must be {expected_shape}, got {tuple(qweight.shape)}"
            )
        params = TensorCoreConvRotW4A4Layout.Params(
            scale=weight_scale.reshape(orig_shape[0]).contiguous(),
            orig_dtype=torch.bfloat16,
            orig_shape=orig_shape,
            convrot_groupsize=convrot_group_size,
            quant_group_size=quant_group_size,
            linear_dtype=linear_dtype,
        )
        self.in_features = orig_shape[1]
        self.out_features = orig_shape[0]
        self.register_buffer(
            "weight",
            QuantizedTensor(
                qweight.contiguous(), "TensorCoreConvRotW4A4Layout", params
            ),
        )
        self.register_buffer("bias", bias.contiguous() if bias is not None else None)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        bias = self.bias.to(inputs.dtype) if self.bias is not None else None
        return F.linear(inputs, self.weight, bias)


def _set_submodule(model: nn.Module, name: str, module: nn.Module) -> None:
    parent_name, _, child_name = name.rpartition(".")
    parent = model.get_submodule(parent_name) if parent_name else model
    setattr(parent, child_name, module)


def load_quantized_checkpoint(
    model: nn.Module, checkpoint_path: str | Path
) -> dict[str, int]:
    """Restore original, INT8, or mixed INT4/INT8 RUM checkpoint tensors."""
    if version("comfy-kitchen") != COMFY_KITCHEN_VERSION:
        raise RuntimeError(
            f"This checkpoint requires comfy-kitchen {COMFY_KITCHEN_VERSION}; "
            f"found {version('comfy-kitchen')}"
        )
    with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle:
        keys = set(handle.keys())
        raw_metadata = (handle.metadata() or {}).get("_quantization_metadata")
        layers = json.loads(raw_metadata).get("layers", {}) if raw_metadata else {}
        quantized_names = set()
        counts = {"int8": 0, "int4": 0}
        for name, metadata in layers.items():
            original = model.get_submodule(name)
            if not isinstance(original, nn.Linear):
                raise TypeError(f"checkpoint layer {name!r} is not an nn.Linear")
            weight_key = f"{name}.weight"
            scale_key = f"{name}.weight_scale"
            if weight_key not in keys or scale_key not in keys:
                raise KeyError(f"incomplete quantized layer {name!r}")
            bias_key = f"{name}.bias"
            bias = handle.get_tensor(bias_key) if bias_key in keys else None
            layer_format = metadata.get("format")
            if layer_format == "int8_tensorwise":
                replacement = NativeINT8Linear(
                    handle.get_tensor(weight_key),
                    handle.get_tensor(scale_key),
                    bias=bias,
                    convrot_group_size=(
                        int(metadata["convrot_groupsize"])
                        if metadata.get("convrot")
                        else None
                    ),
                )
                counts["int8"] += 1
            elif layer_format == "convrot_w4a4":
                raw_shape = metadata.get("orig_shape")
                if not raw_shape or len(raw_shape) != 2:
                    raise ValueError(f"INT4 layer {name!r} is missing orig_shape")
                replacement = NativeINT4ConvRotLinear(
                    handle.get_tensor(weight_key),
                    handle.get_tensor(scale_key),
                    orig_shape=(int(raw_shape[0]), int(raw_shape[1])),
                    bias=bias,
                    convrot_group_size=int(metadata["convrot_groupsize"]),
                    quant_group_size=int(metadata.get("quant_group_size", 64)),
                    linear_dtype=str(metadata.get("linear_dtype", "int4")),
                )
                counts["int4"] += 1
            else:
                raise ValueError(f"unsupported quantized format {layer_format!r}")
            _set_submodule(model, name, replacement)
            quantized_names.add(name)

        excluded = {
            f"{name}.{suffix}"
            for name in quantized_names
            for suffix in ("weight", "weight_scale", "bias")
        }
        state = {
            key: handle.get_tensor(key)
            for key in keys
            if key not in excluded
            and not key.endswith((".weight_scale", ".comfy_quant"))
        }
    incompatible = model.load_state_dict(state, strict=False, assign=True)
    allowed_missing = tuple(f"{name}." for name in quantized_names)
    unexpected_missing = [
        key for key in incompatible.missing_keys if not key.startswith(allowed_missing)
    ]
    if unexpected_missing or incompatible.unexpected_keys:
        raise RuntimeError(
            "checkpoint/model mismatch: "
            f"missing={unexpected_missing}, unexpected={incompatible.unexpected_keys}"
        )
    return {
        "quantized_linears": len(quantized_names),
        "int8_linears": counts["int8"],
        "int4_linears": counts["int4"],
    }


def combine_prompt_embeddings(
    qwen_embeddings: torch.Tensor,
    clip_embeddings: torch.Tensor | None = None,
) -> torch.Tensor:
    """Return Qwen-only conditioning or RUM's Qwen + padded SDXL CLIP tokens."""
    if qwen_embeddings.ndim != 3 or qwen_embeddings.shape[-1] != 7680:
        raise ValueError("Qwen embeddings must have shape (batch, tokens, 7680)")
    if clip_embeddings is None:
        return qwen_embeddings
    if clip_embeddings.ndim != 3 or clip_embeddings.shape[-1] != 2048:
        raise ValueError("SDXL CLIP embeddings must have shape (batch, tokens, 2048)")
    if clip_embeddings.shape[0] != qwen_embeddings.shape[0]:
        raise ValueError("Qwen and CLIP embedding batch sizes must match")
    clip = F.pad(
        clip_embeddings.to(
            device=qwen_embeddings.device, dtype=qwen_embeddings.dtype
        ),
        (0, 7680 - 2048),
    )
    return torch.cat([qwen_embeddings, clip], dim=1)


def resolve_checkpoint(checkpoint: str) -> Path:
    """Resolve a local checkpoint without performing implicit downloads."""
    local = Path(checkpoint).expanduser()
    if local.is_file():
        return local.resolve()
    raise FileNotFoundError(
        f"checkpoint {checkpoint!r} was not found. Download it first and pass "
        "the local .safetensors path with --checkpoint"
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--checkpoint",
        required=True,
        help="Local quantized RUM safetensors path.",
    )
    parser.add_argument("--base-model", default=DEFAULT_BASE_MODEL)
    parser.add_argument(
        "--teacher-model",
        default=DEFAULT_TEACHER_MODEL,
        help=(
            "Local compatible SDXL/WAI checkpoint or pipeline directory used "
            "for full RUM dual-text conditioning."
        ),
    )
    parser.add_argument(
        "--qwen-only",
        action="store_true",
        help="Skip the SDXL teacher for a lower-quality dependency-light trial.",
    )
    parser.add_argument("--task", choices=("t2i", "edit"), default="t2i")
    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
    parser.add_argument("--negative-prompt", default="")
    parser.add_argument("--image", help="Input image path required for edit mode.")
    parser.add_argument("--output", default="rum_output.png")
    parser.add_argument("--width", type=int, default=960)
    parser.add_argument("--height", type=int, default=1152)
    parser.add_argument("--steps", type=int, default=20)
    parser.add_argument("--cfg", type=float, default=5.0)
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--device", default="cuda")
    parser.add_argument(
        "--offload",
        action="store_true",
        help=(
            "Use Diffusers model CPU offload and release the SDXL teacher from "
            "CUDA after prompt encoding. Reduces VRAM at the cost of transfers."
        ),
    )
    return parser


def warn_if_quality_steps_are_too_low(task: str, steps: int) -> None:
    recommended = 20 if task == "t2i" else 10
    if steps < recommended:
        warnings.warn(
            f"{steps} step(s) is below the recommended {recommended} for {task} "
            "and is not suitable for judging image quality. RUM is not a "
            "step-distilled model.",
            UserWarning,
            stacklevel=2,
        )


def validate_request(args: argparse.Namespace) -> None:
    if args.task == "edit" and not args.image:
        raise ValueError("--image is required when --task edit")
    if args.image and not Path(args.image).expanduser().is_file():
        raise FileNotFoundError(args.image)
    if args.width <= 0 or args.height <= 0 or args.steps <= 0:
        raise ValueError("width, height, and steps must be positive")
    warn_if_quality_steps_are_too_low(args.task, args.steps)
    if not Path(args.base_model).expanduser().is_dir():
        raise FileNotFoundError(
            f"base model directory not found: {args.base_model}. Download it first."
        )
    if not args.qwen_only and not Path(args.teacher_model).expanduser().exists():
        raise FileNotFoundError(
            f"teacher model not found: {args.teacher_model}. Download it first "
            "or pass --qwen-only for a lower-quality trial."
        )
    if args.device != "cuda":
        raise ValueError("native RUM INT8/INT4 inference currently requires --device cuda")
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required but torch.cuda.is_available() is false")


def _load_teacher(args: argparse.Namespace, device: torch.device):
    if args.qwen_only:
        warnings.warn(
            "--qwen-only enabled: this convenience mode does not reproduce full "
            "RUM dual-text quality.",
            stacklevel=2,
        )
        return None
    source_path = Path(args.teacher_model).expanduser().resolve()
    common = {"torch_dtype": torch.float16, "local_files_only": True}
    if source_path.is_file():
        teacher = StableDiffusionXLPipeline.from_single_file(
            str(source_path), **common
        )
    else:
        teacher = StableDiffusionXLPipeline.from_pretrained(
            str(source_path), **common
        )
    if not args.offload:
        teacher.text_encoder.to(device)
        teacher.text_encoder_2.to(device)
    teacher.unet = None
    teacher.vae = None
    gc.collect()
    return teacher


@torch.inference_mode()
def _encode_prompt(
    pipeline: Flux2KleinPipeline,
    teacher,
    prompt: str,
    device: torch.device,
) -> torch.Tensor:
    qwen, _ = pipeline.encode_prompt(
        prompt=prompt,
        device=device,
        max_sequence_length=200,
        text_encoder_out_layers=(10, 20, 30),
    )
    if teacher is None:
        return combine_prompt_embeddings(qwen.to(torch.bfloat16))
    clip, *_ = teacher.encode_prompt(prompt, device=device)
    return combine_prompt_embeddings(qwen.to(torch.bfloat16), clip)


def load_pipeline(args: argparse.Namespace):
    checkpoint = resolve_checkpoint(args.checkpoint)
    base_model = str(Path(args.base_model).expanduser().resolve())
    config = RUMFlux2Transformer.load_config(
        base_model,
        subfolder="transformer",
        local_files_only=True,
    )
    with torch.device("meta"):
        transformer = RUMFlux2Transformer.from_config(config)
    quantization = load_quantized_checkpoint(transformer, checkpoint)
    pipeline = Flux2KleinPipeline.from_pretrained(
        base_model,
        transformer=transformer,
        torch_dtype=torch.bfloat16,
        local_files_only=True,
    )
    device = torch.device(args.device)
    if args.offload:
        pipeline.enable_model_cpu_offload(device=device)
    else:
        pipeline.to(device)
    pipeline.set_progress_bar_config(desc="RUM")
    teacher = _load_teacher(args, device)
    return pipeline, teacher, checkpoint, quantization


def _cuda_memory() -> dict[str, int]:
    return {
        "allocated_bytes": torch.cuda.memory_allocated(),
        "reserved_bytes": torch.cuda.memory_reserved(),
    }


def _finish_cuda_stage(started: float) -> dict[str, Any]:
    torch.cuda.synchronize()
    current = _cuda_memory()
    return {
        "seconds": time.perf_counter() - started,
        "end_allocated_bytes": current["allocated_bytes"],
        "end_reserved_bytes": current["reserved_bytes"],
        "peak_allocated_bytes": torch.cuda.max_memory_allocated(),
        "peak_reserved_bytes": torch.cuda.max_memory_reserved(),
    }


@torch.inference_mode()
def run(args: argparse.Namespace) -> dict[str, Any]:
    validate_request(args)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    load_started = time.perf_counter()
    pipeline, teacher, checkpoint, quantization = load_pipeline(args)
    load_profile = _finish_cuda_stage(load_started)
    load_seconds = load_profile["seconds"]
    device = torch.device(args.device)

    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    encode_started = time.perf_counter()
    if args.offload and teacher is not None:
        teacher.text_encoder.to(device)
        teacher.text_encoder_2.to(device)
    prompt_embeddings = _encode_prompt(pipeline, teacher, args.prompt, device)
    negative_embeddings = _encode_prompt(
        pipeline, teacher, args.negative_prompt, device
    )
    if args.offload:
        if teacher is not None:
            teacher.text_encoder.to("cpu")
            teacher.text_encoder_2.to("cpu")
        pipeline.maybe_free_model_hooks()
        torch.cuda.synchronize()
        gc.collect()
        torch.cuda.empty_cache()
    encode_profile = _finish_cuda_stage(encode_started)

    generation_args: dict[str, Any] = {
        "prompt_embeds": prompt_embeddings,
        "negative_prompt_embeds": negative_embeddings,
        "generator": torch.Generator(device="cpu").manual_seed(args.seed),
        "num_inference_steps": args.steps,
        "guidance_scale": args.cfg,
    }
    if args.task == "edit":
        with Image.open(args.image) as source:
            generation_args["image"] = source.convert("RGB")
    else:
        generation_args.update(width=args.width, height=args.height)

    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    started = time.perf_counter()
    denoise_profile: dict[str, Any] = {}

    def record_denoise_peak(_pipeline, step: int, _timestep, callback_kwargs):
        if step == args.steps - 1:
            torch.cuda.synchronize()
            denoise_profile.update(
                {
                    "seconds": time.perf_counter() - started,
                    "peak_allocated_bytes": torch.cuda.max_memory_allocated(),
                    "peak_reserved_bytes": torch.cuda.max_memory_reserved(),
                }
            )
            torch.cuda.reset_peak_memory_stats()
        return callback_kwargs

    generation_args["callback_on_step_end"] = record_denoise_peak
    image = pipeline(**generation_args).images[0]
    torch.cuda.synchronize()
    generation_seconds = time.perf_counter() - started
    decode_peak = {
        "seconds": generation_seconds - denoise_profile.get("seconds", 0.0),
        "peak_allocated_bytes": torch.cuda.max_memory_allocated(),
        "peak_reserved_bytes": torch.cuda.max_memory_reserved(),
    }
    generation_peak_allocated = max(
        denoise_profile.get("peak_allocated_bytes", 0),
        decode_peak["peak_allocated_bytes"],
    )
    generation_peak_reserved = max(
        denoise_profile.get("peak_reserved_bytes", 0),
        decode_peak["peak_reserved_bytes"],
    )
    if args.offload:
        pipeline.maybe_free_model_hooks()
        torch.cuda.synchronize()
        gc.collect()
        torch.cuda.empty_cache()
    after_generation = _cuda_memory()
    output = Path(args.output).expanduser()
    output.parent.mkdir(parents=True, exist_ok=True)
    image.save(output)
    properties = torch.cuda.get_device_properties(device)
    report = {
        "checkpoint": str(checkpoint),
        "base_model": args.base_model,
        "teacher_model": None if args.qwen_only else args.teacher_model,
        "text_conditioning": "qwen+sdxl" if teacher is not None else "qwen-only",
        "task": args.task,
        "prompt": args.prompt,
        "negative_prompt": args.negative_prompt,
        "seed": args.seed,
        "steps": args.steps,
        "cfg": args.cfg,
        "output": str(output.resolve()),
        "load_seconds": load_seconds,
        "generation_seconds": generation_seconds,
        "offload": args.offload,
        "peak_allocated_bytes": generation_peak_allocated,
        "peak_reserved_bytes": generation_peak_reserved,
        "memory_profile": {
            "load": load_profile,
            "prompt_encoding": encode_profile,
            "denoising": denoise_profile,
            "vae_decode_and_postprocess": decode_peak,
            "after_generation_and_offload": after_generation,
        },
        "quantization": quantization,
        "environment": {
            "python": platform.python_version(),
            "torch": torch.__version__,
            "cuda": torch.version.cuda,
            "gpu": properties.name,
            "comfy_kitchen": version("comfy-kitchen"),
        },
    }
    report_path = output.with_suffix(output.suffix + ".json")
    report_path.write_text(
        json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    return report


def main() -> None:
    args = build_parser().parse_args()
    report = run(args)
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()