"""Local scaled-FP8 checkpoint support for the Krea backends."""
from __future__ import annotations
from dataclasses import dataclass
import json
from pathlib import Path
from typing import Any
import torch
from torch import nn
from torch.nn import functional as F
from safetensors import safe_open
from diffusion_cli.config import ModelSource
from diffusion_cli.errors import DiffusionCliError
SUPPORTED_FP8_FORMAT = "float8_e4m3fn"
FP8_DTYPES = {torch.float8_e4m3fn}
@dataclass(frozen=True)
class QuantizedLayerSpec:
"""Quantization settings for one parameterized layer."""
format: str
full_precision_matrix_mult: bool = False
@dataclass(frozen=True)
class QuantizationManifest:
"""Quantization settings indexed by module path."""
layers: dict[str, QuantizedLayerSpec]
def _decodeJson(value: Any, label: str) -> Any:
"""Decode JSON metadata stored as a string or uint8 tensor."""
if isinstance(value, torch.Tensor):
value = bytes(value.detach().cpu().tolist())
if isinstance(value, bytes):
try:
value = value.decode("utf-8")
except UnicodeDecodeError as exc:
raise DiffusionCliError(
f"Invalid quantization metadata in {label}"
) from exc
if isinstance(value, str):
try:
return json.loads(value)
except json.JSONDecodeError as exc:
raise DiffusionCliError(
f"Invalid quantization metadata in {label}: {exc}"
) from exc
if isinstance(value, dict):
return value
raise DiffusionCliError(f"Invalid quantization metadata in {label}")
def _normalizeLayerName(name: str) -> str:
"""Normalize ComfyUI and Transformers module prefixes."""
name = name.removesuffix(".weight")
name = name.removesuffix(".comfy_quant")
changed = True
while changed:
changed = False
for prefix in (
"model.diffusion_model.",
"model.",
"diffusion_model.",
):
if name.startswith(prefix):
name = name.removeprefix(prefix)
changed = True
break
return name
def _manifestEntries(value: Any) -> dict[str, Any]:
"""Find the layer mapping in the observed ComfyUI metadata shapes."""
if not isinstance(value, dict):
raise DiffusionCliError("Quantization metadata must be a JSON object")
for key in ("layers", "quantization", "modules"):
nested = value.get(key)
if isinstance(nested, dict):
return nested
return value
def _specFromValue(value: Any, label: str) -> QuantizedLayerSpec:
"""Parse one quantized-layer metadata object."""
if isinstance(value, str):
format_name = value
full_precision = False
elif isinstance(value, dict):
format_name = value.get("format") or value.get("quant_format")
full_precision = bool(value.get("full_precision_matrix_mult", False))
else:
raise DiffusionCliError(f"Invalid quantization metadata for {label}")
if not isinstance(format_name, str) or not format_name:
raise DiffusionCliError(
f"Missing quantization format for {label}"
)
return QuantizedLayerSpec(format_name, full_precision)
def manifestFromMetadata(value: Any, label: str = "checkpoint") -> QuantizationManifest:
"""Normalize a JSON quantization metadata object into a manifest."""
entries = _manifestEntries(_decodeJson(value, label))
layers: dict[str, QuantizedLayerSpec] = {}
for raw_name, raw_spec in entries.items():
if not isinstance(raw_name, str):
raise DiffusionCliError(f"Invalid quantization layer in {label}")
name = _normalizeLayerName(raw_name)
if name.startswith("_"):
continue
layers[name] = _specFromValue(raw_spec, name)
return QuantizationManifest(layers)
def mergeManifests(*manifests: QuantizationManifest) -> QuantizationManifest:
"""Merge manifests while rejecting inconsistent layer declarations."""
layers: dict[str, QuantizedLayerSpec] = {}
for manifest in manifests:
for name, spec in manifest.layers.items():
previous = layers.get(name)
if previous is not None and previous != spec:
raise DiffusionCliError(
f"Inconsistent quantization metadata for {name}"
)
layers[name] = spec
return QuantizationManifest(layers)
def manifestFromStateDict(
state_dict: dict[str, torch.Tensor],
) -> QuantizationManifest:
"""Read per-layer ``comfy_quant`` tensors from a state dict."""
manifests = []
for key, value in state_dict.items():
if key.endswith(".comfy_quant"):
manifests.append(manifestFromMetadata(value, key))
return mergeManifests(*manifests)
def loadQuantizationManifest(source: ModelSource) -> QuantizationManifest:
"""Read file-level and per-layer Krea quantization metadata locally."""
manifests: list[QuantizationManifest] = []
with safe_open(source.path, framework="pt", device="cpu") as tensors:
metadata = tensors.metadata() or {}
raw_metadata = metadata.get("_quantization_metadata")
if raw_metadata is not None:
manifests.append(manifestFromMetadata(raw_metadata, str(source.path)))
for key in tensors.keys():
if source.checkpoint_prefix is not None:
if not key.startswith(source.checkpoint_prefix):
continue
if key.endswith(".comfy_quant"):
decoded = _decodeJson(tensors.get_tensor(key), key)
if isinstance(decoded, dict) and "format" in decoded:
manifests.append(
QuantizationManifest(
{
_normalizeLayerName(key): _specFromValue(
decoded,
_normalizeLayerName(key),
)
}
)
)
else:
manifests.append(manifestFromMetadata(decoded, key))
return mergeManifests(*manifests)
def validateQuantizedStateDict(
state_dict: dict[str, torch.Tensor],
manifest: QuantizationManifest | None = None,
) -> QuantizationManifest:
"""Validate scaled-FP8 weights, scalar scales, and supported metadata."""
if manifest is None:
manifest = manifestFromStateDict(state_dict)
for name, spec in manifest.layers.items():
if spec.format != SUPPORTED_FP8_FORMAT:
raise DiffusionCliError(
f"Unsupported quantization format in {name}: {spec.format}"
)
weight_key = f"{name}.weight"
scale_key = f"{name}.weight_scale"
weight = next(
(
state_dict.get(candidate)
for candidate in (
weight_key,
f"model.{weight_key}",
f"diffusion_model.{weight_key}",
)
if state_dict.get(candidate) is not None
),
None,
)
if weight is None:
raise DiffusionCliError(f"Quantized layer {name} is missing weight")
if weight.dtype not in FP8_DTYPES:
raise DiffusionCliError(
f"Unsupported quantized weight dtype in {name}: {weight.dtype}"
)
scale = next(
(
state_dict.get(candidate)
for candidate in (
scale_key,
f"model.{scale_key}",
f"diffusion_model.{scale_key}",
)
if state_dict.get(candidate) is not None
),
None,
)
if scale is None:
raise DiffusionCliError(
f"Krea layer {name} is missing weight_scale"
)
if scale.numel() != 1:
raise DiffusionCliError(
"Expected scalar weight scale, got shape "
f"{tuple(scale.shape)}"
)
if scale.dtype != torch.float32:
raise DiffusionCliError(
f"Expected F32 weight scale in {name}, got {scale.dtype}"
)
bias = next(
(
state_dict.get(candidate)
for candidate in (f"{name}.bias", f"model.{name}.bias")
if state_dict.get(candidate) is not None
),
None,
)
if bias is not None and (
bias.dtype in FP8_DTYPES or not bias.dtype.is_floating_point
):
raise DiffusionCliError(f"Quantized bias is unsupported in {name}")
for key, value in state_dict.items():
if key.endswith(".weight") and (
value.dtype in {torch.int8, torch.uint8, torch.int32, torch.int64}
):
raise DiffusionCliError(
f"Unsupported integer quantized weight in {_normalizeLayerName(key)}"
)
if key.endswith(".weight") and value.dtype in FP8_DTYPES:
name = _normalizeLayerName(key)
if name not in manifest.layers:
raise DiffusionCliError(
f"FP8 layer {name} has no quantization metadata"
)
return manifest
class ScaledFp8Linear(nn.Module):
"""A correctness-first linear layer for scalar-scaled FP8 weights."""
def __init__(
self,
in_features: int,
out_features: int,
*,
bias: bool = True,
device=None,
dtype=torch.bfloat16,
) -> None:
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.weight = nn.Parameter(
torch.empty(
out_features,
in_features,
device=device,
dtype=torch.float8_e4m3fn,
),
requires_grad=False,
)
self.register_buffer(
"weight_scale",
torch.ones((), device=device, dtype=torch.float32),
)
if bias:
self.bias = nn.Parameter(
torch.zeros(out_features, device=device, dtype=dtype),
requires_grad=False,
)
else:
self.register_parameter("bias", None)
def forward(self, input: torch.Tensor) -> torch.Tensor:
"""Dequantize this layer's weight for one matrix operation."""
compute_dtype = input.dtype
weight = self.weight.to(device=input.device, dtype=compute_dtype)
scale = self.weight_scale.to(device=input.device, dtype=compute_dtype)
weight = weight * scale.reshape(1)
bias = None if self.bias is None else self.bias.to(input.dtype)
return F.linear(input, weight, bias)
def toDevice(self, device, dtype=None) -> None:
"""Move the layer without expanding its FP8 storage weight."""
self.weight.data = self.weight.data.to(device=device)
self.weight_scale.data = self.weight_scale.data.to(device=device)
if self.bias is not None:
self.bias.data = self.bias.data.to(device=device)
if dtype is not None and self.bias.dtype.is_floating_point:
self.bias.data = self.bias.data.to(dtype=dtype)
def moveModulePreservingQuantization(
module: nn.Module,
device,
runtime_dtype: torch.dtype,
) -> nn.Module:
"""Move a module while retaining FP8 and F32 storage dtypes."""
for child in module.modules():
if isinstance(child, ScaledFp8Linear):
child.toDevice(device, runtime_dtype)
else:
for parameter in child.parameters(recurse=False):
if parameter.dtype in FP8_DTYPES:
parameter.data = parameter.data.to(device=device)
else:
parameter.data = parameter.data.to(
device=device,
dtype=(
runtime_dtype
if parameter.dtype.is_floating_point
and parameter.dtype != torch.float32
else parameter.dtype
),
)
for name, buffer in child.named_buffers(recurse=False):
if buffer is not None:
buffer = buffer.to(device=device)
child._buffers[name] = buffer
return module