533 lines
24 KiB
Python
533 lines
24 KiB
Python
#!/usr/bin/env python3
|
|
"""Quantize a safetensors file to NVFP4 GGUF without a Hugging Face model directory.
|
|
|
|
The learned rounding is the same convert_to_quant pass the other scripts use,
|
|
so a file you are already quantizing with convert_to_quant can be packed into
|
|
GGUF and compared tensor by tensor. Tensor names are kept. There is no
|
|
tokenizer and no architecture config; the GGUF is a weight container.
|
|
|
|
By default, tensors that convert_to_quant's base exclusion list would skip
|
|
(norms, biases, embeddings, lm_head, and the other names in that list) stay
|
|
in their original dtype. Pass --all-weights to quantize every 2D matrix whose
|
|
last dimension is a multiple of 64. A sibling bias is corrected the same way
|
|
as hf_to_nvfp4_gguf.py. Repeat --mmproj to copy projector GGUFs beside the
|
|
output and add a FROM line for each one.
|
|
|
|
--no-quant skips that rounding and writes the safetensors into a GGUF as
|
|
float. --outtype chooses f32, f16, bf16, or auto, which keeps each tensor's
|
|
stored dtype. --outtype nvfp4 repacks a ComfyUI NVFP4 safetensors file
|
|
(packed weight, weight_scale, and weight_scale_2) without fitting it again.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from safetensors import safe_open
|
|
|
|
from nvfp4gguf import (
|
|
DEFAULT_LLAMA_CPP,
|
|
TensorQuantizer,
|
|
add_quant_arguments,
|
|
bind_exclude,
|
|
ctq_skips_key,
|
|
ensure_cuda_interpreter,
|
|
import_gguf,
|
|
load_learned_converter,
|
|
looks_like_spark,
|
|
quantize_and_pack,
|
|
report_reconstruction,
|
|
require_quant_settings,
|
|
)
|
|
from nvfp4gguf.bundle import (
|
|
outtype as _outtype,
|
|
refuse_existing_bundle,
|
|
require_context_length,
|
|
warn_missing_gguf_metadata,
|
|
write_simple_modelfile,
|
|
)
|
|
|
|
log = logging.getLogger("safetensors_to_nvfp4_gguf")
|
|
|
|
|
|
def _checkpoint_file(directory: Path, name: str) -> Path:
|
|
digest = hashlib.sha1(name.encode("utf-8")).hexdigest()[:16]
|
|
safe = "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in name)[:80]
|
|
return directory / f"{safe}.{digest}.npz"
|
|
|
|
|
|
def _load_checkpoint(directory: Path, name: str, shape: tuple[int, ...], need_bias: bool):
|
|
path = _checkpoint_file(directory, name)
|
|
if not path.is_file():
|
|
return None
|
|
with np.load(path) as blob:
|
|
if tuple(int(v) for v in blob["shape"]) != tuple(shape):
|
|
return None
|
|
if need_bias and "bias" not in blob.files:
|
|
return None
|
|
raw = np.array(blob["raw"], copy=True)
|
|
scale = np.array(blob["scale"], copy=True)
|
|
bias = np.array(blob["bias"], copy=True) if "bias" in blob.files else np.zeros((0,), np.float32)
|
|
return raw, None if scale.size == 0 else scale, None if bias.size == 0 else bias
|
|
|
|
|
|
def _save_checkpoint(directory: Path, name: str, shape: tuple[int, ...], raw, scale, bias) -> None:
|
|
path = _checkpoint_file(directory, name)
|
|
temporary = path.with_name(path.name + ".tmp")
|
|
with temporary.open("wb") as handle:
|
|
np.savez_compressed(
|
|
handle,
|
|
shape=np.asarray(shape, dtype=np.int64),
|
|
raw=raw,
|
|
scale=np.zeros((0,), np.float32) if scale is None else np.asarray(scale, dtype=np.float32),
|
|
bias=np.zeros((0,), np.float32) if bias is None else np.asarray(bias, dtype=np.float32),
|
|
)
|
|
temporary.replace(path)
|
|
|
|
|
|
def _should_quantize(
|
|
key: str,
|
|
shape: tuple[int, ...],
|
|
all_weights: bool,
|
|
convert_to_quant_dir: Path | None,
|
|
exclude_re,
|
|
) -> bool:
|
|
if len(shape) != 2 or not key.endswith(".weight"):
|
|
return False
|
|
if shape[-1] % 64 != 0 or shape[-1] < 64:
|
|
return False
|
|
if exclude_re is not None and exclude_re.search(key):
|
|
return False
|
|
if not all_weights and ctq_skips_key(key, convert_to_quant_dir):
|
|
return False
|
|
return True
|
|
|
|
|
|
def _apply_bias(name: str, tensor: torch.Tensor, correction: np.ndarray | None) -> torch.Tensor:
|
|
if correction is None:
|
|
return tensor
|
|
values = tensor.detach().cpu().float()
|
|
flat = values.reshape(-1)
|
|
if tuple(flat.shape) != tuple(correction.shape):
|
|
log.warning(" bias %s shape %s does not match correction %s", name, tuple(flat.shape), tuple(correction.shape))
|
|
return tensor
|
|
flat = flat - torch.from_numpy(np.ascontiguousarray(correction, dtype=np.float32))
|
|
values = flat.reshape(values.shape).to(dtype=tensor.dtype)
|
|
log.info(" bias %s adjusted (mean %+.3e)", name, float(np.mean(correction)))
|
|
return values
|
|
|
|
|
|
def _float_file_type(gguf, outtype: str):
|
|
return {
|
|
"f32": gguf.LlamaFileType.ALL_F32,
|
|
"f16": gguf.LlamaFileType.MOSTLY_F16,
|
|
"bf16": gguf.LlamaFileType.MOSTLY_BF16,
|
|
}[outtype]
|
|
|
|
|
|
def _auto_file_type(gguf, dtypes: list[str]):
|
|
names = {dtype.upper() for dtype in dtypes}
|
|
if "BF16" in names:
|
|
return gguf.LlamaFileType.MOSTLY_BF16
|
|
if "F16" in names:
|
|
return gguf.LlamaFileType.MOSTLY_F16
|
|
if "F32" in names:
|
|
return gguf.LlamaFileType.ALL_F32
|
|
return gguf.LlamaFileType.GUESSED
|
|
|
|
|
|
def _cast_float(tensor: torch.Tensor, outtype: str) -> torch.Tensor:
|
|
if outtype == "auto" or not tensor.is_floating_point():
|
|
return tensor
|
|
dtype = {"f32": torch.float32, "f16": torch.float16, "bf16": torch.bfloat16}[outtype]
|
|
if tensor.dtype == dtype:
|
|
return tensor
|
|
return tensor.to(dtype)
|
|
|
|
|
|
def _default_output(source: Path, skip_quant: bool, outtype: str) -> Path:
|
|
if not skip_quant:
|
|
return source.with_name(source.stem + "-NVFP4.gguf")
|
|
if outtype == "auto":
|
|
return source.with_name(source.stem + ".gguf")
|
|
label = "NVFP4" if outtype == "nvfp4" else outtype.upper()
|
|
return source.with_name(f"{source.stem}-{label}.gguf")
|
|
|
|
|
|
def _add_original(writer, gguf, name: str, tensor: torch.Tensor) -> None:
|
|
tensor = tensor.detach().cpu().contiguous()
|
|
if tensor.dtype == torch.bfloat16:
|
|
raw = tensor.view(torch.uint16).numpy().view(np.uint8)
|
|
writer.add_tensor(name, np.ascontiguousarray(raw), raw_dtype=gguf.GGMLQuantizationType.BF16)
|
|
return
|
|
if tensor.dtype == torch.float16:
|
|
writer.add_tensor(name, tensor.numpy())
|
|
return
|
|
if tensor.dtype == torch.float32:
|
|
writer.add_tensor(name, tensor.numpy())
|
|
return
|
|
writer.add_tensor(name, tensor.float().numpy())
|
|
|
|
|
|
def _prepare_output(args) -> tuple[Path, Path, bool]:
|
|
source = Path(args.safetensors).resolve()
|
|
if not source.is_file():
|
|
raise SystemExit(f"Not a file: {source}")
|
|
skip_quant = bool(args.no_quant or args.outtype == "nvfp4")
|
|
output = Path(args.output).resolve() if args.output else _default_output(source, skip_quant, args.outtype)
|
|
return source, output, skip_quant
|
|
|
|
|
|
def _comfy_formats(handle) -> dict[str, str]:
|
|
"""Map a weight prefix to the format named in its .comfy_quant sidecar."""
|
|
found: dict[str, str] = {}
|
|
suffix = ".comfy_quant"
|
|
for key in handle.keys():
|
|
if not key.endswith(suffix):
|
|
continue
|
|
raw = handle.get_tensor(key).detach().cpu().numpy().tobytes()
|
|
try:
|
|
payload = json.loads(raw.decode("utf-8"))
|
|
except (UnicodeError, json.JSONDecodeError):
|
|
continue
|
|
fmt = payload.get("format") if isinstance(payload, dict) else None
|
|
if isinstance(fmt, str):
|
|
found[key[: -len(suffix)]] = fmt
|
|
return found
|
|
|
|
|
|
def _pack_comfy_nvfp4(weight: torch.Tensor, scale: torch.Tensor) -> np.ndarray:
|
|
"""Repack ComfyUI / ModelOpt NVFP4 into ggml super-blocks.
|
|
|
|
The code bytes are already nibble-packed the way ggml stores them. Each
|
|
group of four 16-value blocks becomes 4 UE4M3 scale bytes followed by
|
|
those 32 code bytes.
|
|
"""
|
|
if scale.ndim != 2 or weight.ndim != 2:
|
|
raise ValueError(f"NVFP4 weight {tuple(weight.shape)} and scale {tuple(scale.shape)} must be matrices")
|
|
rows, n_blocks = scale.shape
|
|
if weight.shape != (rows, n_blocks * 8):
|
|
raise ValueError(f"NVFP4 weight {tuple(weight.shape)} does not match scale {tuple(scale.shape)}")
|
|
if n_blocks % 4 != 0:
|
|
raise ValueError(f"NVFP4 scale count {n_blocks} is not a multiple of 4")
|
|
n_super = n_blocks // 4
|
|
codes = np.ascontiguousarray(weight.detach().cpu().contiguous().view(torch.uint8).numpy()).reshape(rows, n_super, 32)
|
|
scales = np.ascontiguousarray(scale.detach().cpu().contiguous().view(torch.uint8).numpy()).reshape(rows, n_super, 4)
|
|
scales = scales & np.uint8(0x7F)
|
|
return np.ascontiguousarray(np.concatenate([scales, codes], axis=-1).reshape(rows, n_super * 36))
|
|
|
|
|
|
def _scale_is_trivial(scale: torch.Tensor) -> bool:
|
|
return scale.numel() <= 1 and abs(float(scale.detach().float().sum()) - 1.0) < 1e-6
|
|
|
|
|
|
def _float8_to_bf16(tensor: torch.Tensor, scale: torch.Tensor | None) -> torch.Tensor:
|
|
values = tensor.detach().float()
|
|
if scale is not None:
|
|
factor = scale.detach().float()
|
|
if factor.numel() == 1:
|
|
values = values * factor.reshape(())
|
|
elif tuple(factor.shape) == tuple(values.shape):
|
|
values = values * factor
|
|
else:
|
|
raise ValueError(f"float8 scale {tuple(factor.shape)} does not match weight {tuple(values.shape)}")
|
|
return values.to(torch.bfloat16)
|
|
|
|
|
|
def convert_packed_nvfp4(args) -> None:
|
|
"""Copy ComfyUI NVFP4 blocks into a GGUF. Float tensors are stored as stored."""
|
|
logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
|
|
source, output, _skip = _prepare_output(args)
|
|
overwrite = bool(args.force or args.overwrite)
|
|
if output.exists() and not overwrite:
|
|
raise SystemExit(f"{output} already exists. Pass --overwrite to replace it.")
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
if args.mmproj:
|
|
modelfile = output.parent / "Modelfile"
|
|
if modelfile.exists() and not overwrite:
|
|
raise SystemExit(f"{modelfile} already exists. Pass --overwrite to replace it.")
|
|
refuse_existing_bundle(output, overwrite, args.mmproj, skip_gguf=True)
|
|
|
|
gguf = import_gguf(args.llama_cpp)
|
|
with safe_open(source, framework="pt", device="cpu") as handle:
|
|
keys = list(handle.keys())
|
|
formats = _comfy_formats(handle)
|
|
packed = [name for name, fmt in formats.items() if fmt == "nvfp4"]
|
|
if not packed:
|
|
raise SystemExit(
|
|
f"{source} has no ComfyUI NVFP4 weights to repack. "
|
|
"Use --no-quant with --outtype auto, f16, bf16, or f32 to write float weights without quantizing them."
|
|
)
|
|
present = set(keys)
|
|
for name in packed:
|
|
if name + ".weight" not in present or name + ".weight_scale" not in present:
|
|
raise SystemExit(f"{name} is marked nvfp4 but is missing weight or weight_scale")
|
|
log.info("Repacking %d NVFP4 tensors from %s. Learned rounding is skipped.", len(packed), source.name)
|
|
writer = gguf.GGUFWriter(output, "safetensors", use_temp_file=not args.no_temp_file)
|
|
writer.add_file_type(int(gguf.LlamaFileType.MOSTLY_NVFP4))
|
|
writer.add_string("general.name", source.stem)
|
|
writer.add_string("general.description", "NVFP4 weights repacked from safetensors without refitting")
|
|
writer.add_uint32("general.quantization_version", 2)
|
|
context_length = require_context_length(args.context_length)
|
|
if context_length is not None:
|
|
writer.add_context_length(context_length)
|
|
log.info("context length %d", context_length)
|
|
|
|
skip = {name + ".comfy_quant" for name in formats}
|
|
for name, fmt in formats.items():
|
|
if fmt in ("nvfp4", "float8_e4m3fn"):
|
|
skip.add(name + ".weight_scale")
|
|
skip.add(name + ".weight_scale_2")
|
|
promoted = 0
|
|
with safe_open(source, framework="pt", device="cpu") as handle:
|
|
for key in keys:
|
|
if key in skip:
|
|
continue
|
|
if key.endswith(".weight") and formats.get(key[: -len(".weight")]) == "nvfp4":
|
|
base = key[: -len(".weight")]
|
|
raw = _pack_comfy_nvfp4(handle.get_tensor(key), handle.get_tensor(base + ".weight_scale"))
|
|
writer.add_tensor(key, raw, raw_dtype=gguf.GGMLQuantizationType.NVFP4)
|
|
scale2_name = base + ".weight_scale_2"
|
|
if scale2_name in present:
|
|
scale2 = handle.get_tensor(scale2_name)
|
|
if not _scale_is_trivial(scale2):
|
|
writer.add_tensor(base + ".scale", scale2.detach().float().numpy().reshape(-1).astype(np.float32))
|
|
continue
|
|
if key.endswith(".weight") and formats.get(key[: -len(".weight")]) == "float8_e4m3fn":
|
|
base = key[: -len(".weight")]
|
|
scale_name = base + ".weight_scale"
|
|
scale = handle.get_tensor(scale_name) if scale_name in present else None
|
|
_add_original(writer, gguf, key, _float8_to_bf16(handle.get_tensor(key), scale))
|
|
promoted += 1
|
|
log.info("Stored %s as BF16 (source is float8_e4m3fn)", key)
|
|
continue
|
|
_add_original(writer, gguf, key, handle.get_tensor(key))
|
|
if promoted:
|
|
log.info("%d float8 tensors were stored as BF16. GGUF has no float8 weight type.", promoted)
|
|
log.info("Writing %s", output)
|
|
warn_missing_gguf_metadata(writer)
|
|
writer.write_header_to_file()
|
|
writer.write_kv_data_to_file()
|
|
writer.write_tensors_to_file(progress=True)
|
|
writer.close()
|
|
if args.mmproj:
|
|
write_simple_modelfile(output, overwrite, args.mmproj)
|
|
log.info("Wrote %s", output)
|
|
|
|
|
|
def convert_float(args) -> None:
|
|
"""Write the safetensors into a GGUF without learned rounding."""
|
|
logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
|
|
if args.outtype == "nvfp4":
|
|
convert_packed_nvfp4(args)
|
|
return
|
|
source, output, _skip = _prepare_output(args)
|
|
overwrite = bool(args.force or args.overwrite)
|
|
if output.exists() and not overwrite:
|
|
raise SystemExit(f"{output} already exists. Pass --overwrite to replace it.")
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
if args.mmproj:
|
|
modelfile = output.parent / "Modelfile"
|
|
if modelfile.exists() and not overwrite:
|
|
raise SystemExit(f"{modelfile} already exists. Pass --overwrite to replace it.")
|
|
refuse_existing_bundle(output, overwrite, args.mmproj, skip_gguf=True)
|
|
|
|
gguf = import_gguf(args.llama_cpp)
|
|
with safe_open(source, framework="pt", device="cpu") as handle:
|
|
keys = list(handle.keys())
|
|
dtypes = [handle.get_slice(key).get_dtype() for key in keys]
|
|
file_type = _float_file_type(gguf, args.outtype) if args.outtype != "auto" else _auto_file_type(gguf, dtypes)
|
|
log.info("Writing float GGUF (%s). NVFP4 quantization is skipped.", args.outtype)
|
|
writer = gguf.GGUFWriter(output, "safetensors", use_temp_file=not args.no_temp_file)
|
|
writer.add_file_type(int(file_type))
|
|
writer.add_string("general.name", source.stem)
|
|
writer.add_string("general.description", "Weights converted from safetensors without quantization")
|
|
context_length = require_context_length(args.context_length)
|
|
if context_length is not None:
|
|
writer.add_context_length(context_length)
|
|
log.info("context length %d", context_length)
|
|
with safe_open(source, framework="pt", device="cpu") as handle:
|
|
for key in keys:
|
|
_add_original(writer, gguf, key, _cast_float(handle.get_tensor(key), args.outtype))
|
|
log.info("Writing %s", output)
|
|
warn_missing_gguf_metadata(writer)
|
|
writer.write_header_to_file()
|
|
writer.write_kv_data_to_file()
|
|
writer.write_tensors_to_file(progress=True)
|
|
writer.close()
|
|
if args.mmproj:
|
|
write_simple_modelfile(output, overwrite, args.mmproj)
|
|
log.info("Wrote %s", output)
|
|
|
|
|
|
def convert(args) -> None:
|
|
logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
|
|
if args.no_quant or args.outtype == "nvfp4":
|
|
convert_float(args)
|
|
return
|
|
ensure_cuda_interpreter(args.device)
|
|
source, output, _skip = _prepare_output(args)
|
|
overwrite = bool(args.force or args.overwrite)
|
|
if output.exists() and not overwrite:
|
|
raise SystemExit(f"{output} already exists. Pass --overwrite to replace it.")
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
if args.mmproj:
|
|
modelfile = output.parent / "Modelfile"
|
|
if modelfile.exists() and not overwrite:
|
|
raise SystemExit(f"{modelfile} already exists. Pass --overwrite to replace it.")
|
|
refuse_existing_bundle(output, overwrite, args.mmproj, skip_gguf=True)
|
|
|
|
gguf = import_gguf(args.llama_cpp)
|
|
with safe_open(source, framework="pt", device="cpu") as handle:
|
|
keys = list(handle.keys())
|
|
shapes = {key: tuple(handle.get_slice(key).get_shape()) for key in keys}
|
|
bind_exclude(args, "spark2_5" if looks_like_spark(keys) else None)
|
|
chosen = [
|
|
key for key in keys
|
|
if _should_quantize(key, shapes[key], args.all_weights, args.convert_to_quant, args.exclude_re)
|
|
]
|
|
log.info("%d of %d tensors queued for NVFP4", len(chosen), len(keys))
|
|
if not chosen:
|
|
raise SystemExit("No 2D .weight tensor has a last dimension divisible by 64.")
|
|
|
|
work = Path(args.checkpoint_dir) if args.checkpoint_dir else output.with_name(output.name + ".nvfp4-work")
|
|
if args.fresh and work.exists():
|
|
shutil.rmtree(work)
|
|
work.mkdir(parents=True, exist_ok=True)
|
|
require_quant_settings(work, args)
|
|
quantizer = TensorQuantizer(args, gguf, load_learned_converter(args.convert_to_quant)[1])
|
|
writer = gguf.GGUFWriter(output, "safetensors", use_temp_file=not args.no_temp_file)
|
|
writer.add_file_type(int(gguf.LlamaFileType.MOSTLY_NVFP4))
|
|
writer.add_string("general.name", source.stem)
|
|
writer.add_string("general.description", "NVFP4 weights converted from safetensors for comparison with convert_to_quant")
|
|
writer.add_uint32("general.quantization_version", 2)
|
|
context_length = require_context_length(args.context_length)
|
|
if context_length is not None:
|
|
writer.add_context_length(context_length)
|
|
log.info("context length %d", context_length)
|
|
|
|
done = 0
|
|
chosen_set = set(chosen)
|
|
pending_bias: dict[str, torch.Tensor] = {}
|
|
corrections: dict[str, np.ndarray] = {}
|
|
|
|
def _write_bias(key: str, tensor: torch.Tensor) -> None:
|
|
weight_name = key[: -len(".bias")] + ".weight"
|
|
correction = corrections.get(key) if weight_name in chosen_set and not args.no_bias_correction else None
|
|
_add_original(writer, gguf, key, _apply_bias(key, tensor, correction))
|
|
|
|
with safe_open(source, framework="pt", device="cpu") as handle:
|
|
for key in keys:
|
|
tensor = handle.get_tensor(key)
|
|
weight_name = key[: -len(".bias")] + ".weight" if key.endswith(".bias") else ""
|
|
if weight_name in chosen_set and not args.no_bias_correction and key not in corrections:
|
|
pending_bias[key] = tensor
|
|
continue
|
|
if key not in chosen_set:
|
|
if len(shapes[key]) == 2 and key.endswith(".weight") and shapes[key][-1] % 64 != 0:
|
|
log.info("Keeping %s %s (last dimension is not a multiple of 64)", key, shapes[key])
|
|
if key.endswith(".bias") and weight_name in chosen_set:
|
|
_write_bias(key, tensor)
|
|
else:
|
|
_add_original(writer, gguf, key, tensor)
|
|
continue
|
|
done += 1
|
|
bias_name = key[: -len(".weight")] + ".bias"
|
|
need_bias = not args.no_bias_correction and bias_name in shapes and len(shapes[key]) == 2
|
|
cached = _load_checkpoint(work, key, shapes[key], need_bias)
|
|
log.info(
|
|
"NVFP4 [%d/%d] %s %s%s",
|
|
done, len(chosen), key, shapes[key], " (checkpoint)" if cached else "",
|
|
)
|
|
if cached is None:
|
|
weight = tensor.float().numpy()
|
|
raw, scale = quantize_and_pack(quantizer, weight, False, key)
|
|
report_reconstruction(quantizer, weight, raw, scale)
|
|
bias = None
|
|
if need_bias:
|
|
per_tensor = None if scale is None else float(np.reshape(scale, -1)[0])
|
|
runtime = quantizer.runtime_dequant(raw, per_tensor)
|
|
bias = quantizer.bias_correction(weight, runtime)
|
|
_save_checkpoint(work, key, shapes[key], raw, scale, bias)
|
|
del weight
|
|
else:
|
|
raw, scale, bias = cached
|
|
if bias is not None:
|
|
corrections[bias_name] = bias
|
|
pending = pending_bias.pop(bias_name, None)
|
|
if pending is not None:
|
|
_write_bias(bias_name, pending)
|
|
writer.add_tensor(key, raw, raw_dtype=gguf.GGMLQuantizationType.NVFP4)
|
|
if scale is not None:
|
|
writer.add_tensor(key[: -len(".weight")] + ".scale", scale)
|
|
del tensor, raw
|
|
for key, tensor in pending_bias.items():
|
|
_write_bias(key, tensor)
|
|
|
|
log.info("Writing %s", output)
|
|
warn_missing_gguf_metadata(writer)
|
|
writer.write_header_to_file()
|
|
writer.write_kv_data_to_file()
|
|
writer.write_tensors_to_file(progress=True)
|
|
writer.close()
|
|
if not args.keep_checkpoints:
|
|
shutil.rmtree(work, ignore_errors=True)
|
|
if args.mmproj:
|
|
write_simple_modelfile(output, overwrite, args.mmproj)
|
|
log.info("Wrote %s", output)
|
|
|
|
|
|
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Quantize a safetensors file to an NVFP4 GGUF weight container.")
|
|
parser.add_argument("safetensors", nargs="?", help="Input .safetensors file")
|
|
parser.add_argument("-o", "--output", help="Output GGUF path. Default: <stem>-NVFP4.gguf, or <stem>.gguf with --no-quant and auto")
|
|
parser.add_argument("--no-quant", action="store_true",
|
|
help="Skip learned rounding. Writes a float GGUF, or repacks ComfyUI NVFP4 weights with --outtype nvfp4")
|
|
parser.add_argument("--outtype", default="auto", type=_outtype, choices=("auto", "f32", "f16", "bf16", "nvfp4"),
|
|
help="Type used when learned rounding is skipped. nvfp4 repacks ComfyUI NVFP4 weights. auto keeps each tensor's stored dtype (default: auto)")
|
|
parser.add_argument("--all-weights", action="store_true",
|
|
help="Quantize every eligible 2D weight. By default, convert_to_quant's exclusion list is skipped too.")
|
|
parser.add_argument("--llama-cpp", type=Path, default=DEFAULT_LLAMA_CPP,
|
|
help="llama.cpp checkout used for the GGUF writer (default: ~/llama.cpp)")
|
|
parser.add_argument("--overwrite", action="store_true", help="Replace an existing GGUF, Modelfile, or copied projector")
|
|
parser.add_argument("--force", action="store_true", help="Same as --overwrite")
|
|
parser.add_argument("--mmproj", action="append", type=Path, default=None,
|
|
help="Vision projector GGUF. Copied next to the output under its own name and added as another FROM line. Repeat for more than one projector.")
|
|
parser.add_argument("--context-length", type=int, default=None,
|
|
help="Max context stored in the GGUF. A safetensors file has no context length to copy.")
|
|
parser.add_argument("--checkpoint-dir", type=Path, default=None, help="Per-tensor resume directory")
|
|
parser.add_argument("--fresh", action="store_true")
|
|
parser.add_argument("--keep-checkpoints", action="store_true")
|
|
parser.add_argument("--no-temp-file", action="store_true")
|
|
parser.add_argument("--verbose", action="store_true")
|
|
parser.add_argument("--self-test", action="store_true", help="Check NVFP4 packing against llama.cpp's dequantizer and exit")
|
|
add_quant_arguments(parser)
|
|
args = parser.parse_args(argv)
|
|
if args.self_test:
|
|
if args.device == "auto":
|
|
args.device = "cpu"
|
|
return args
|
|
if not args.safetensors:
|
|
parser.error("safetensors file is required")
|
|
return args
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> None:
|
|
args = parse_args(argv)
|
|
if args.self_test:
|
|
from hf_to_nvfp4_gguf import self_test
|
|
self_test(args)
|
|
return
|
|
convert(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|