Spaces:
Running on Zero
Running on Zero
| """Export toolkit-quantized modules into ComfyUI ``comfy_quant`` checkpoints — | |
| the inverse of toolkit/util/comfy_quant_import.py. | |
| Every quantized OstrisLinear whose backend has a comfy storage format emits | |
| its quantized tensors plus the ``<prefix>.comfy_quant`` uint8 JSON marker: | |
| - convrot8 (int8_tensorwise + convrot): weight int8 [out, in], fp32 | |
| weight_scale [out, 1] (comfy_kitchen's per-channel convention) | |
| - nvfp4: high-nibble-first packed fp4 pairs, e4m3 block scales re-swizzled | |
| to the cuBLAS 128x4 tile layout, fp32 weight_scale_2 per-tensor scale and | |
| optional AWQ pre_quant_scale | |
| - float8_e4m3fn: fp8_e4m3 weight + one fp32 per-tensor weight_scale | |
| - convrotcomfyw4a4: via convrot_quant.export_comfy_convrot_w4a4 | |
| Biases are NOT emitted here — they are ordinary parameters and flow through | |
| the regular state_dict path. | |
| """ | |
| import json | |
| from typing import Dict, List, Tuple | |
| import torch | |
| from toolkit.util.ostris_quant import OstrisLinear | |
| def comfy_quant_marker(conf: dict) -> torch.Tensor: | |
| return torch.tensor(list(json.dumps(conf).encode("utf-8")), dtype=torch.uint8) | |
| def export_comfy_quantized_layers( | |
| root: torch.nn.Module, | |
| ) -> Tuple[Dict[str, torch.Tensor], List[str], List[str]]: | |
| """Comfy-format state-dict entries for every quantized OstrisLinear in | |
| ``root``. Returns ``(entries, exported_names, unexportable_names)`` — | |
| entries are keyed by module path in root's layout (run | |
| convert_state_dict_on_save afterwards for the checkpoint layout); | |
| unexportable_names lists quantized modules whose backend has no comfy | |
| storage format (the caller decides whether to dequantize instead).""" | |
| from toolkit.util.convrot_quant import ( | |
| ConvRotComfyW4A4Quantizer, | |
| export_comfy_convrot_w4a4, | |
| ) | |
| from toolkit.util.nvfp4_quant import swap_nvfp4_nibbles, swizzle_nvfp4_scales | |
| entries: Dict[str, torch.Tensor] = {} | |
| exported: List[str] = [] | |
| unexportable: List[str] = [] | |
| for name, module in root.named_modules(): | |
| if not isinstance(module, OstrisLinear): | |
| continue | |
| if hasattr(module, "cr8_qdata"): | |
| rot = int(getattr(module, "cr8_rot_size", 1) or 1) | |
| conf = {"format": "int8_tensorwise"} | |
| if rot > 1: | |
| conf.update({"convrot": True, "convrot_groupsize": rot}) | |
| entries[f"{name}.weight"] = module.cr8_qdata.detach().cpu().contiguous() | |
| entries[f"{name}.weight_scale"] = ( | |
| module.cr8_scales.view(torch.float32) | |
| .detach() | |
| .cpu() | |
| .reshape(module.out_features, 1) | |
| .contiguous() | |
| ) | |
| entries[f"{name}.comfy_quant"] = comfy_quant_marker(conf) | |
| elif hasattr(module, "nv4_qdata"): | |
| entries[f"{name}.weight"] = swap_nvfp4_nibbles( | |
| module.nv4_qdata.detach().cpu() | |
| ) | |
| scales = module.nv4_scales.view(torch.float8_e4m3fn).detach().cpu() | |
| entries[f"{name}.weight_scale"] = swizzle_nvfp4_scales( | |
| scales.reshape(module.out_features, module.in_features // 16) | |
| ).view(torch.float8_e4m3fn) | |
| entries[f"{name}.weight_scale_2"] = ( | |
| module.nv4_pts.view(torch.float32).detach().cpu().reshape(()) | |
| ) | |
| if hasattr(module, "nv4_pre_scale"): | |
| entries[f"{name}.pre_quant_scale"] = ( | |
| module.nv4_pre_scale.view(torch.float32).detach().cpu().contiguous() | |
| ) | |
| entries[f"{name}.comfy_quant"] = comfy_quant_marker({"format": "nvfp4"}) | |
| elif hasattr(module, "f8_qdata"): | |
| entries[f"{name}.weight"] = module.f8_qdata.detach().cpu().contiguous() | |
| entries[f"{name}.weight_scale"] = ( | |
| module.f8_scale.view(torch.float32).detach().cpu().reshape(()) | |
| ) | |
| entries[f"{name}.comfy_quant"] = comfy_quant_marker( | |
| {"format": "float8_e4m3fn", "full_precision_matrix_mult": True} | |
| ) | |
| elif isinstance(module.ostris_quantizer, ConvRotComfyW4A4Quantizer): | |
| layer_entries = export_comfy_convrot_w4a4(module, f"{name}.") | |
| layer_entries.pop(f"{name}.bias", None) | |
| entries.update( | |
| {k: v.detach().cpu() if torch.is_tensor(v) else v for k, v in layer_entries.items()} | |
| ) | |
| else: | |
| unexportable.append(name) | |
| continue | |
| exported.append(name) | |
| return entries, exported, unexportable | |