130 lines
4.7 KiB
Python
130 lines
4.7 KiB
Python
"""Utility functions for calibration data I/O and size estimation."""
|
|
import logging
|
|
import os
|
|
from typing import Dict, Tuple
|
|
|
|
import torch
|
|
|
|
logger = logging.getLogger("comfyui_gptq_calibration")
|
|
|
|
__all__ = [
|
|
"default_dual_output_paths",
|
|
"save_calibration",
|
|
"load_calibration",
|
|
"estimate_disk_size",
|
|
"human_size",
|
|
]
|
|
|
|
|
|
def default_dual_output_paths(base_path: str | None = None) -> Tuple[str, str]:
|
|
"""Derive ``<base>_positive.pt`` and ``<base>_negative.pt`` output paths.
|
|
|
|
If *base_path* is ``None``, uses the ComfyUI output directory.
|
|
If *base_path* ends with ``.pt``, the suffix is stripped and replaced.
|
|
Raises ``ValueError`` if the two paths resolve to the same location.
|
|
"""
|
|
try:
|
|
import folder_paths
|
|
default_dir = folder_paths.get_output_directory() if hasattr(folder_paths, "get_output_directory") else folder_paths.output_directory
|
|
except ImportError:
|
|
default_dir = "/tmp"
|
|
|
|
if base_path is None or not base_path.strip():
|
|
base = os.path.join(default_dir, "calibration")
|
|
else:
|
|
base = base_path
|
|
if base.endswith(".pt"):
|
|
base = base[:-3]
|
|
|
|
pos = f"{base}_positive.pt"
|
|
neg = f"{base}_negative.pt"
|
|
|
|
if os.path.realpath(pos) == os.path.realpath(neg):
|
|
raise ValueError(
|
|
f"output_path_positive and output_path_negative resolve to the same "
|
|
f"path: {os.path.realpath(pos)}. Use different paths."
|
|
)
|
|
return pos, neg
|
|
|
|
|
|
def save_calibration(data: dict, path: str) -> str:
|
|
"""Save calibration data to a .pt file.
|
|
|
|
The directory is created if it does not exist. The path is returned
|
|
so callers can chain a ``(path,)`` tuple back to ComfyUI.
|
|
"""
|
|
directory = os.path.dirname(os.path.abspath(path))
|
|
if directory and not os.path.isdir(directory):
|
|
os.makedirs(directory, exist_ok=True)
|
|
torch.save(data, path)
|
|
return path
|
|
|
|
|
|
def load_calibration(path: str) -> dict:
|
|
"""Load calibration data from a .pt file (CPU, ``weights_only=True``).
|
|
|
|
``weights_only=True`` is a security measure that prevents arbitrary code
|
|
execution when loading untrusted ``.pt`` files.
|
|
"""
|
|
return torch.load(path, map_location="cpu", weights_only=True)
|
|
|
|
|
|
def estimate_disk_size(num_layers: int,
|
|
avg_in_features: int,
|
|
hessian_block_size: int = 0,
|
|
collect_amax: bool = True,
|
|
dtype_bytes: int = 4,
|
|
hessian_format: str = "block",
|
|
dlr_rank: int = 0,
|
|
collect_channel_amax: bool = False,
|
|
collect_timestep_amax: bool = False,
|
|
timestep_buckets: int = 0) -> Dict[str, float]:
|
|
"""Estimate output file size.
|
|
|
|
``dtype_bytes`` defaults to 4 (float32). Returns a dict with keys:
|
|
``hessian_bytes``, ``amax_bytes``, ``channel_amax_bytes``,
|
|
``total_bytes``, ``total_gb`` (using ``1024**3``, not ``1e9``),
|
|
and ``per_layer_bytes``.
|
|
"""
|
|
if hessian_format in ("dlr", "dlr_nystrom") and dlr_rank > 0:
|
|
rank = min(dlr_rank, avg_in_features)
|
|
# DLR stores D (n floats) + U (n * rank floats)
|
|
per_layer = (avg_in_features + avg_in_features * rank) * dtype_bytes
|
|
elif hessian_block_size and hessian_block_size > 0:
|
|
block = hessian_block_size
|
|
num_blocks = max(1, (avg_in_features + block - 1) // block)
|
|
per_layer = num_blocks * block * block * dtype_bytes
|
|
else:
|
|
per_layer = avg_in_features * avg_in_features * dtype_bytes
|
|
|
|
hessian_bytes = num_layers * per_layer
|
|
amax_bytes = num_layers * dtype_bytes if collect_amax else 0
|
|
|
|
# ViDiT-Q channel amax: per-channel vector per layer + per-timestep-bucket tensor.
|
|
channel_amax_bytes = 0
|
|
if collect_channel_amax:
|
|
# channel_amax: in_features floats per layer
|
|
channel_amax_bytes += num_layers * avg_in_features * dtype_bytes
|
|
# channel_amax_by_timestep: num_buckets * in_features floats per layer
|
|
if collect_timestep_amax and timestep_buckets > 0:
|
|
channel_amax_bytes += num_layers * timestep_buckets * avg_in_features * dtype_bytes
|
|
|
|
total_bytes = hessian_bytes + amax_bytes + channel_amax_bytes
|
|
return {
|
|
"hessian_bytes": hessian_bytes,
|
|
"amax_bytes": amax_bytes,
|
|
"channel_amax_bytes": channel_amax_bytes,
|
|
"total_bytes": total_bytes,
|
|
"total_gb": total_bytes / (1024 ** 3),
|
|
"per_layer_bytes": per_layer,
|
|
}
|
|
|
|
|
|
def human_size(num_bytes: int) -> str:
|
|
"""Format a byte count as a human-readable string (B/KB/MB/GB/TB/PB)."""
|
|
for unit in ("B", "KB", "MB", "GB", "TB"):
|
|
if num_bytes < 1024.0:
|
|
return f"{num_bytes:.2f} {unit}"
|
|
num_bytes /= 1024.0
|
|
return f"{num_bytes:.2f} PB"
|