Add files using upload-large-folder tool

README.md

@@ -0,0 +1,136 @@
+---
+license: apache-2.0
+base_model:
+- Wan-AI/Wan2.2-T2V-A14B
+library_name: diffusers
+tags:
+- video_generation
+- NVFP4
+- Sparse_Attention
+- Wan
+---
+# 🎬 Wan2.2-NVFP4-Sparse
+
+> **An extremely efficient Wan 2.2 14B variant: NVFP4 Quantization-Aware Step Distillation with Sparse Attention for Blackwell Architecture**
+
+[![GitHub](https://img.shields.io/badge/GitHub-ModelTC/LightX2V-blue)](https://github.com/ModelTC/LightX2V)
+[![HuggingFace](https://img.shields.io/badge/HuggingFace-lightx2v-yellow)](https://huggingface.co/lightx2v/)
+
+## 📋 Table of Contents
+
+- [✨ Features](#-features)
+- [🚀 Quick Start](#-quick-start)
+- [🎬 Generation Results](#-generation-results)
+- [⚡ Performance Comparison](#-performance-comparison)
+- [⚠️ Notes](#️-notes)
+- [🤝 Community](#-community)
+
+## ✨ Features
+
+- **⚡ 4-Step Inference**: Two high-noise expert steps followed by two low-noise expert steps, enabling extremely fast Wan2.2 MoE generation on a single Blackwell GPU.
+- **🎯 NVFP4 Quantization**: Quantization-aware step distillation reduces memory traffic and compute cost while targeting Blackwell architecture.
+- **🧩 Sparse Attention**: Accelerates the costly O(n²) self-attention workload with sparse attention, reducing end-to-end latency for high-resolution video generation.
+- **🔧 LightX2V Integration**: Recommended runtime stack for stable deployment and best performance.
+- **🚀 High-Quality Generation**: Preserves the visual quality of Wan2.2-T2V-14B while dramatically improving inference speed.
+
+## 🚀 Quick Start
+
+We strongly recommend using the official LightX2V Docker image for the cleanest environment and best reproducibility.
+
+### Option A: Docker Recommended
+
+```bash
+# 1. Pull LightX2V Docker image
+docker pull lightx2v/lightx2v:26052801-cu130-5090
+
+# 2. Run inference
+bash scripts/wan22/distill/run_wan22_moe_t2v_extreme.sh
+```
+
+### Option B: Manual Installation
+
+If Docker is not available, install the environment manually:
+
+```bash
+# 1. Install LightX2V
+git clone https://github.com/ModelTC/LightX2V.git
+cd LightX2V
+uv pip install -v .
+
+# 2. Install NVFP4 Kernel
+pip install scikit_build_core uv
+git clone https://github.com/NVIDIA/cutlass.git
+cd lightx2v_kernel
+
+MAX_JOBS=$(nproc) CMAKE_BUILD_PARALLEL_LEVEL=$(nproc) \
+uv build --wheel \
+  -Cbuild-dir=build . \
+  -Ccmake.define.CUTLASS_PATH=/path/to/cutlass \
+  --verbose --color=always --no-build-isolation
+
+pip install dist/*whl --force-reinstall --no-deps
+
+# 3. Run inference
+bash scripts/wan22/distill/run_wan22_moe_t2v_extreme.sh
+```
+
+Script: [run_wan22_moe_t2v_extreme.sh](https://github.com/ModelTC/LightX2V/blob/main/scripts/wan22/distill/run_wan22_moe_t2v_extreme.sh)
+
+## 🎬 Generation Results
+
+<div style="background: #f8fafc; border: 1px solid #e2e8f0; border-radius: 8px; padding: 16px; margin: 16px 0;">
+<p style="font-style: italic; color: #475569; margin: 0; padding: 12px; background: white; border-radius: 6px; border-left: 4px solid #3b82f6;">
+"Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage"
+</p>
+</div>
+
+
+| Resolution | Wan2.2-T2V-14B | Wan2.2-NVFP4-Sparse |
+| --- | --- | --- |
+| 480p | <video controls style="width: 260px; height: 180px; border-radius: 6px; object-fit: cover;" src="https://cdn-uploads.huggingface.co/production/uploads/658e760cccbc1e2cc78b4258/WTHhrzx7XR4S1Ys_6Kzx4.mp4"></video> | <video controls style="width: 260px; height: 180px; border-radius: 6px; object-fit: cover;" src="https://cdn-uploads.huggingface.co/production/uploads/658e760cccbc1e2cc78b4258/zorpw7gm9At0J2kCmvkDr.mp4"></video> |
+| 720p | <video controls style="width: 260px; height: 180px; border-radius: 6px; object-fit: cover;" src="https://cdn-uploads.huggingface.co/production/uploads/658e760cccbc1e2cc78b4258/vkiyKj7CJA-r0yTz7TEum.mp4"></video> | <video controls style="width: 260px; height: 180px; border-radius: 6px; object-fit: cover;" src="https://cdn-uploads.huggingface.co/production/uploads/658e760cccbc1e2cc78b4258/TuECbzvW5jI9NHG6GLvIR.mp4"></video> |
+
+
+## ⚡ Performance Comparison
+
+**Test Environment**: RTX 5090 Single GPU | LightX2V Framework | End-to-End Latency
+
+| Resolution | Wan2.2-T2V-14B | Wan2.2-NVFP4-Sparse | Speedup |
+| --- | ---: | ---: | ---: |
+| 480p | 734s | 14.15s | 51.9x |
+| 720p | 2668s | 45s | 59.3x |
+
+## ⚠️ Notes
+
+### System Requirements
+
+- **Required Hardware**: NVIDIA RTX 50-series GPUs or other Blackwell architecture GPUs.
+- **Recommended Runtime**: `lightx2v/lightx2v:26052801-cu130-5090`.
+
+### Dependencies
+
+- Prepare Wan2.2 T5 / VAE components following the standard LightX2V Wan2.2 model structure.
+- Use Blackwell + NVFP4 kernels for optimal speed and memory efficiency.
+
+### Performance Tips
+
+- Use the provided extreme inference script for the 4-step high-noise / low-noise expert schedule.
+- Sparse attention is most beneficial at higher resolutions where self-attention dominates latency.
+- Enable CPU offload only when GPU memory is limited, since offload can reduce throughput.
+
+## 🤝 Community
+
+- **🐛 Issues**: [GitHub Issues](https://github.com/ModelTC/LightX2V/issues)
+- **🤗 Models**: [HuggingFace Hub](https://huggingface.co/lightx2v/)
+- **📖 Documentation**: [LightX2V Docs](https://github.com/ModelTC/LightX2V)
+
+---
+
+<div align="center">
+
+**If you find this project helpful, please give us a ⭐ on [GitHub](https://github.com/ModelTC/LightX2V)**
+
+For questions or issues, please open an issue on [LightX2V](https://github.com/ModelTC/LightX2V/issues) or contact lvchengtao0319@gmail.com.
+
+</div>
+

Wan2.2-I2V-A14B_NVFP4_Sparse_high.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c77b92616d678fa8373462973a3d2e3ceefd2fbdb3ccf33454e74f4c79fcf07a
+size 8379833592

Wan2.2-I2V-A14B_NVFP4_Sparse_high_comfy.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:998934b07f124111b83c3bfcdf9efc2ff67ff531879806e54a1453c14066a5db
+size 8379918824

Wan2.2-I2V-A14B_NVFP4_Sparse_low.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d76b53286dd9b68f5d956941037c0a6fd5a1fe0fd05eb97f8bf8abf9915da562
+size 8379833592

Wan2.2-I2V-A14B_NVFP4_Sparse_low_comfy.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a26076d43590941a6526ed52d2b682a9340c81f650ada0751a90e5bdcde5efb7
+size 8379918824

Wan2.2-T2V-A14B_NVFP4_Sparse_high.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1da94b1604f4bb181d1319f311245e39e33e499b3a0a5bf2fafb117b08750ac7
+size 8379014392

Wan2.2-T2V-A14B_NVFP4_Sparse_low.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9feda8a7f4791c6e4f0dd8ab383fa67d429441dfb006b97a1e3b5400e74bc232
+size 8379014392

Wan2.2-T2V-A14B_NVFP4_Sparse_low_comfy.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:765476d100853683f0b122be824f87121a5bb3c8cb1c5838c4dbe9eccd76a43b
+oid sha256:cdccd3005e0050a700a9dec435ab3ef37923294ba03b3132b53f94460f377f6d
 size 8379099624

analysis.md

@@ -0,0 +1,138 @@
+# Wan2.2 NVFP4 Sparse to ComfyUI Conversion Analysis
+
+## Sources checked
+
+- Kijai Hugging Face repo: https://huggingface.co/Kijai/WanVideo_comfy_nvfp4
+- ComfyUI Wan2.2 workflow docs: https://docs.comfy.org/tutorials/video/wan/wan2_2
+- ComfyUI Wan2.2 examples: https://comfyanonymous.github.io/ComfyUI_examples/wan22/
+- ComfyUI mixed precision loader reference:
+  https://huggingface.co/mhnakif/comfy/blob/main/comfy/ops.py
+- ComfyUI quant op reference:
+  https://huggingface.co/mhnakif/comfy/blob/main/comfy/quant_ops.py
+- Comfy Kitchen hardware/backend reference:
+  https://github.com/Comfy-Org/comfy-kitchen
+- Local ComfyUI source checkout used for verification:
+  `Comfy-Org/ComfyUI` commit
+  `5aa71b9bc28809a16596bb9fa3d0a6300d8e3f0e`
+
+## Script provenance
+
+`convert_lightx2v_nvfp4_to_comfy.py` is a local conversion script written for
+this directory. It is not copied from one upstream script. The implementation is
+derived from these upstream pages and ComfyUI source conventions:
+
+- Kijai model page:
+  https://huggingface.co/Kijai/WanVideo_comfy_nvfp4
+  - This page gives the actual LightX2V NVFP4 to Comfy NVFP4 conversion rule:
+    nibble-swap packed U8 weights, keep `weight_scale`, set
+    `weight_scale_2 = alpha * input_global_scale`, and set
+    `input_scale = 1 / input_global_scale`.
+- ComfyUI quantized loader:
+  https://github.com/Comfy-Org/ComfyUI/blob/5aa71b9bc28809a16596bb9fa3d0a6300d8e3f0e/comfy/ops.py#L1058-L1091
+  - This loader reads `{layer}.comfy_quant`, branches on `format == "nvfp4"`,
+    then requires `{layer}.weight_scale_2` and `{layer}.weight_scale`.
+- ComfyUI quant algorithm registry:
+  https://github.com/Comfy-Org/ComfyUI/blob/5aa71b9bc28809a16596bb9fa3d0a6300d8e3f0e/comfy/quant_ops.py#L190-L205
+  - This defines the `nvfp4` storage dtype as `torch.uint8` and the parameter
+    set as `weight_scale`, `weight_scale_2`, and `input_scale`.
+- ComfyUI quantization metadata handling:
+  https://github.com/Comfy-Org/ComfyUI/blob/5aa71b9bc28809a16596bb9fa3d0a6300d8e3f0e/comfy/utils.py#L1360-L1421
+  - This shows that `_quantization_metadata.layers` is converted into
+    `{layer}.comfy_quant` JSON byte tensors and that the presence of
+    `.comfy_quant` enables mixed quantized ops.
+- ComfyUI native NVFP4 hardware gate:
+  https://github.com/Comfy-Org/ComfyUI/blob/5aa71b9bc28809a16596bb9fa3d0a6300d8e3f0e/comfy/model_management.py#L1877-L1885
+  - This returns true only for NVIDIA GPUs with compute capability major
+    version >= 10, which is why H100 can validate/load files but is not expected
+    to use native Blackwell NVFP4 tensor-core compute.
+
+## Format findings
+
+The original files in this directory are LightX2V NVFP4 Sparse safetensors. Each
+file has 400 quantized Linear layers with these LightX2V-side tensors:
+
+- `{layer}.weight`: packed NVFP4 values in `torch.uint8`
+- `{layer}.weight_scale`: FP8 E4M3 block scale tensor
+- `{layer}.alpha`: scalar post-matmul rescaler
+- `{layer}.input_global_scale`: scalar input scale convention
+
+Kijai's model card says the ComfyUI conversion is still NVFP4 and uses the same
+datatype, but changes conventions:
+
+- Swap the high/low nibbles in each packed `uint8` weight byte.
+- Keep `{layer}.weight_scale` as-is.
+- Convert `{layer}.alpha * {layer}.input_global_scale` into
+  `{layer}.weight_scale_2`.
+- Convert `1 / {layer}.input_global_scale` into `{layer}.input_scale`.
+
+ComfyUI's mixed precision loader expects a `{layer}.comfy_quant` tensor
+containing JSON bytes. For NVFP4 it then loads:
+
+- `{layer}.weight`
+- `{layer}.weight_scale_2`
+- `{layer}.weight_scale`
+- optional registered parameters such as `{layer}.input_scale`
+
+The converted file metadata also includes `_quantization_metadata` with one
+`nvfp4` layer entry per quantized layer so ComfyUI can select mixed precision
+operations for the model.
+
+## H100 note
+
+The conversion itself does not require a Blackwell GPU; it is a safetensors
+layout conversion. However, Comfy Kitchen documents `TensorCoreNVFP4Layout` as
+requiring SM >= 10.0 / Blackwell for native NVFP4 tensor-core acceleration. H100
+is Hopper, so ComfyUI may disable native NVFP4 compute and run a fallback path.
+
+## Script
+
+The conversion script is:
+
+```bash
+python convert_lightx2v_nvfp4_to_comfy.py
+```
+
+Useful options:
+
+```bash
+python convert_lightx2v_nvfp4_to_comfy.py --dry-run
+python convert_lightx2v_nvfp4_to_comfy.py --overwrite
+python convert_lightx2v_nvfp4_to_comfy.py input.safetensors --output-dir /path/to/out
+```
+
+The script writes `<original_stem>_comfy.safetensors` and uses a temporary file
+before renaming into place.
+
+## Converted outputs
+
+- `Wan2.2-I2V-A14B_NVFP4_Sparse_high_comfy.safetensors`
+- `Wan2.2-I2V-A14B_NVFP4_Sparse_low_comfy.safetensors`
+- `Wan2.2-T2V-A14B_NVFP4_Sparse_high_comfy.safetensors`
+- `Wan2.2-T2V-A14B_NVFP4_Sparse_low_comfy.safetensors`
+
+For ComfyUI native workflows, place these diffusion model files under:
+
+```text
+ComfyUI/models/diffusion_models/
+```
+
+The Wan2.2 14B workflows still need the normal text encoder and VAE files in
+their ComfyUI locations.
+
+## Verification performed
+
+For each converted file:
+
+- Tensor count is 2695.
+- `_quantization_metadata` contains 400 quantized layers.
+- `alpha` count is 0.
+- `input_global_scale` count is 0.
+- `input_scale` count is 400.
+- `weight_scale` count is 400.
+- `weight_scale_2` count is 400.
+- `comfy_quant` count is 400.
+- `{layer}.comfy_quant` decodes to `{"format": "nvfp4"}`.
+- A sampled `blocks.0.cross_attn.k.weight` block equals the expected nibble
+  swap from the original.
+- The sampled `weight_scale_2` equals `alpha * input_global_scale`.
+- The sampled `input_scale` equals `1 / input_global_scale`.

convert_lightx2v_nvfp4_to_comfy.py

@@ -0,0 +1,172 @@
+#!/usr/bin/env python3
+"""Convert LightX2V Wan2.2 NVFP4 safetensors to ComfyUI NVFP4 format."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+
+COMFY_QUANT_CONF = {"format": "nvfp4"}
+QUANT_LAYER_CONF = {"format": "nvfp4", "full_precision_matrix_mult": False}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Convert LightX2V Wan2.2 NVFP4 Sparse safetensors into ComfyUI's "
+            "native NVFP4 safetensors convention."
+        )
+    )
+    parser.add_argument(
+        "inputs",
+        nargs="*",
+        type=Path,
+        help="Input .safetensors files. Defaults to non-_comfy safetensors in the current directory.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=Path("."),
+        help="Directory for converted files. Defaults to the current directory.",
+    )
+    parser.add_argument(
+        "--suffix",
+        default="_comfy",
+        help="Suffix appended before .safetensors. Defaults to _comfy.",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Overwrite existing converted files.",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Inspect inputs and print planned outputs without writing files.",
+    )
+    return parser.parse_args()
+
+
+def default_inputs() -> list[Path]:
+    return sorted(
+        p
+        for p in Path(".").glob("*.safetensors")
+        if not p.name.endswith("_comfy.safetensors")
+    )
+
+
+def output_path(input_path: Path, output_dir: Path, suffix: str) -> Path:
+    return output_dir / f"{input_path.stem}{suffix}{input_path.suffix}"
+
+
+def is_quant_layer(prefix: str, keys: set[str]) -> bool:
+    return (
+        f"{prefix}.weight" in keys
+        and f"{prefix}.weight_scale" in keys
+        and f"{prefix}.alpha" in keys
+        and f"{prefix}.input_global_scale" in keys
+    )
+
+
+def quant_layer_prefixes(keys: list[str]) -> list[str]:
+    key_set = set(keys)
+    prefixes: list[str] = []
+    for key in keys:
+        if key.endswith(".weight"):
+            prefix = key[: -len(".weight")]
+            if is_quant_layer(prefix, key_set):
+                prefixes.append(prefix)
+    return prefixes
+
+
+def swap_fp4_nibbles(weight: torch.Tensor) -> torch.Tensor:
+    if weight.dtype != torch.uint8:
+        raise TypeError(f"Expected uint8 packed NVFP4 weight, got {weight.dtype}")
+    low = torch.bitwise_left_shift(torch.bitwise_and(weight, 0x0F), 4)
+    high = torch.bitwise_right_shift(weight, 4)
+    return torch.bitwise_or(low, high).contiguous()
+
+
+def comfy_quant_tensor() -> torch.Tensor:
+    payload = json.dumps(COMFY_QUANT_CONF).encode("utf-8")
+    return torch.tensor(list(payload), dtype=torch.uint8)
+
+
+def convert_one(input_path: Path, output_path_: Path, overwrite: bool, dry_run: bool) -> None:
+    if output_path_.exists() and not overwrite and not dry_run:
+        raise FileExistsError(f"{output_path_} exists; pass --overwrite to replace it")
+
+    with safe_open(input_path, framework="pt", device="cpu") as sf:
+        keys = list(sf.keys())
+        quant_prefixes = quant_layer_prefixes(keys)
+        quant_prefix_set = set(quant_prefixes)
+
+        if dry_run:
+            print(
+                f"{input_path} -> {output_path_}: "
+                f"{len(keys)} input tensors, {len(quant_prefixes)} quantized layers"
+            )
+            return
+
+        tensors: dict[str, torch.Tensor] = {}
+        for key in keys:
+            if key.endswith(".alpha") or key.endswith(".input_global_scale"):
+                prefix = key.rsplit(".", 1)[0]
+                if prefix in quant_prefix_set:
+                    continue
+
+            if key.endswith(".weight"):
+                prefix = key[: -len(".weight")]
+                tensor = sf.get_tensor(key)
+                if prefix in quant_prefix_set:
+                    tensors[key] = swap_fp4_nibbles(tensor)
+                    alpha = sf.get_tensor(f"{prefix}.alpha").to(torch.float32)
+                    input_global_scale = sf.get_tensor(f"{prefix}.input_global_scale").to(torch.float32)
+                    tensors[f"{prefix}.weight_scale_2"] = (alpha * input_global_scale).contiguous()
+                    tensors[f"{prefix}.input_scale"] = (1.0 / input_global_scale).contiguous()
+                    tensors[f"{prefix}.comfy_quant"] = comfy_quant_tensor()
+                else:
+                    tensors[key] = tensor.contiguous() if not tensor.is_contiguous() else tensor
+                continue
+
+            tensor = sf.get_tensor(key)
+            tensors[key] = tensor.contiguous() if not tensor.is_contiguous() else tensor
+
+    metadata = {
+        "format": "pt",
+        "_quantization_metadata": json.dumps(
+            {"layers": {prefix: QUANT_LAYER_CONF for prefix in quant_prefixes}},
+            sort_keys=True,
+        ),
+    }
+    tmp_path = output_path_.with_name(f"{output_path_.name}.tmp")
+    if tmp_path.exists():
+        tmp_path.unlink()
+    save_file(tensors, tmp_path, metadata=metadata)
+    tmp_path.replace(output_path_)
+    print(f"Wrote {output_path_} ({len(tensors)} tensors, {len(quant_prefixes)} quantized layers)")
+
+
+def main() -> None:
+    args = parse_args()
+    inputs = args.inputs or default_inputs()
+    if not inputs:
+        raise SystemExit("No input .safetensors files found")
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    for input_path in inputs:
+        input_path = input_path.resolve()
+        if input_path.suffix != ".safetensors":
+            raise ValueError(f"Expected a .safetensors input, got {input_path}")
+        out = output_path(input_path, args.output_dir, args.suffix).resolve()
+        convert_one(input_path, out, overwrite=args.overwrite, dry_run=args.dry_run)
+
+
+if __name__ == "__main__":
+    main()

verify_comfy_nvfp4.py

@@ -0,0 +1,167 @@
+#!/usr/bin/env python3
+"""Validate converted Wan2.2 NVFP4 safetensors against ComfyUI expectations."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+
+
+REQUIRED_QUANT_KEYS = (
+    "weight",
+    "weight_scale",
+    "weight_scale_2",
+    "input_scale",
+    "comfy_quant",
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "files",
+        nargs="*",
+        type=Path,
+        help="Converted *_comfy.safetensors files. Defaults to all *_comfy.safetensors in cwd.",
+    )
+    parser.add_argument(
+        "--source-dir",
+        type=Path,
+        default=Path("."),
+        help="Directory containing the original non-_comfy safetensors for numeric relation checks.",
+    )
+    return parser.parse_args()
+
+
+def default_files() -> list[Path]:
+    return sorted(Path(".").glob("*_comfy.safetensors"))
+
+
+def source_for(converted: Path, source_dir: Path) -> Path | None:
+    source_name = converted.name.replace("_comfy.safetensors", ".safetensors")
+    candidate = source_dir / source_name
+    return candidate if candidate.exists() else None
+
+
+def layer_prefixes_from_metadata(metadata: dict[str, str]) -> list[str]:
+    raw = metadata.get("_quantization_metadata")
+    if raw is None:
+        raise AssertionError("missing _quantization_metadata")
+    parsed = json.loads(raw)
+    layers = parsed.get("layers")
+    if not isinstance(layers, dict):
+        raise AssertionError("_quantization_metadata.layers is not an object")
+    for prefix, conf in layers.items():
+        if conf.get("format") != "nvfp4":
+            raise AssertionError(f"{prefix}: expected nvfp4 format, got {conf!r}")
+    return sorted(layers)
+
+
+def check_converted_file(path: Path) -> tuple[int, int]:
+    with safe_open(path, framework="pt", device="cpu") as sf:
+        keys = set(sf.keys())
+        metadata = sf.metadata() or {}
+        prefixes = layer_prefixes_from_metadata(metadata)
+
+        if len(prefixes) != 400:
+            raise AssertionError(f"{path.name}: expected 400 quantized layers, got {len(prefixes)}")
+
+        for prefix in prefixes:
+            for suffix in REQUIRED_QUANT_KEYS:
+                key = f"{prefix}.{suffix}"
+                if key not in keys:
+                    raise AssertionError(f"{path.name}: missing {key}")
+
+            if f"{prefix}.alpha" in keys:
+                raise AssertionError(f"{path.name}: unexpected LightX2V key {prefix}.alpha")
+            if f"{prefix}.input_global_scale" in keys:
+                raise AssertionError(f"{path.name}: unexpected LightX2V key {prefix}.input_global_scale")
+
+            weight = sf.get_tensor(f"{prefix}.weight")
+            weight_scale = sf.get_tensor(f"{prefix}.weight_scale")
+            weight_scale_2 = sf.get_tensor(f"{prefix}.weight_scale_2")
+            input_scale = sf.get_tensor(f"{prefix}.input_scale")
+            comfy_quant = sf.get_tensor(f"{prefix}.comfy_quant")
+
+            if weight.dtype != torch.uint8:
+                raise AssertionError(f"{path.name}: {prefix}.weight dtype is {weight.dtype}, expected uint8")
+            if weight_scale.dtype != torch.float8_e4m3fn:
+                raise AssertionError(
+                    f"{path.name}: {prefix}.weight_scale dtype is {weight_scale.dtype}, expected float8_e4m3fn"
+                )
+            if weight_scale_2.dtype != torch.float32 or weight_scale_2.shape != torch.Size([]):
+                raise AssertionError(f"{path.name}: {prefix}.weight_scale_2 must be scalar float32")
+            if input_scale.dtype != torch.float32 or input_scale.shape != torch.Size([]):
+                raise AssertionError(f"{path.name}: {prefix}.input_scale must be scalar float32")
+
+            marker = json.loads(bytes(comfy_quant.tolist()).decode("utf-8"))
+            if marker != {"format": "nvfp4"}:
+                raise AssertionError(f"{path.name}: {prefix}.comfy_quant is {marker!r}")
+
+        if sum(key.endswith(".input_scale") for key in keys) != 400:
+            raise AssertionError(f"{path.name}: expected 400 input_scale tensors")
+        if sum(key.endswith(".weight_scale_2") for key in keys) != 400:
+            raise AssertionError(f"{path.name}: expected 400 weight_scale_2 tensors")
+        if sum(key.endswith(".comfy_quant") for key in keys) != 400:
+            raise AssertionError(f"{path.name}: expected 400 comfy_quant tensors")
+        if sum(key.endswith(".alpha") for key in keys) != 0:
+            raise AssertionError(f"{path.name}: alpha keys remain after conversion")
+        if sum(key.endswith(".input_global_scale") for key in keys) != 0:
+            raise AssertionError(f"{path.name}: input_global_scale keys remain after conversion")
+
+        return len(keys), len(prefixes)
+
+
+def check_against_source(converted: Path, source: Path) -> None:
+    sample_prefixes = (
+        "blocks.0.cross_attn.k",
+        "blocks.0.ffn.0",
+        "blocks.39.self_attn.v",
+    )
+    with safe_open(source, framework="pt", device="cpu") as src, safe_open(
+        converted, framework="pt", device="cpu"
+    ) as dst:
+        for prefix in sample_prefixes:
+            src_weight = src.get_tensor(f"{prefix}.weight")[:16, :16]
+            dst_weight = dst.get_tensor(f"{prefix}.weight")[:16, :16]
+            expected = torch.bitwise_or(
+                torch.bitwise_left_shift(torch.bitwise_and(src_weight, 0x0F), 4),
+                torch.bitwise_right_shift(src_weight, 4),
+            )
+            if not torch.equal(dst_weight, expected):
+                raise AssertionError(f"{converted.name}: {prefix}.weight nibble swap mismatch")
+
+            alpha = src.get_tensor(f"{prefix}.alpha").to(torch.float32)
+            input_global_scale = src.get_tensor(f"{prefix}.input_global_scale").to(torch.float32)
+            weight_scale_2 = dst.get_tensor(f"{prefix}.weight_scale_2").to(torch.float32)
+            input_scale = dst.get_tensor(f"{prefix}.input_scale").to(torch.float32)
+
+            if not torch.allclose(weight_scale_2, alpha * input_global_scale):
+                raise AssertionError(f"{converted.name}: {prefix}.weight_scale_2 mismatch")
+            if not torch.allclose(input_scale, 1.0 / input_global_scale):
+                raise AssertionError(f"{converted.name}: {prefix}.input_scale mismatch")
+
+
+def main() -> None:
+    args = parse_args()
+    files = args.files or default_files()
+    if not files:
+        raise SystemExit("No *_comfy.safetensors files found")
+
+    for path in files:
+        tensor_count, layer_count = check_converted_file(path)
+        source = source_for(path, args.source_dir)
+        if source is not None:
+            check_against_source(path, source)
+            relation = "source relation checks passed"
+        else:
+            relation = "source file not found; skipped relation checks"
+        print(f"OK {path.name}: {tensor_count} tensors, {layer_count} nvfp4 layers, {relation}")
+
+
+if __name__ == "__main__":
+    main()

verify_comfy_source_compat.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+"""Check converted NVFP4 files against a local ComfyUI source checkout."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("comfyui", type=Path, help="Path to a ComfyUI source checkout")
+    parser.add_argument(
+        "files",
+        nargs="*",
+        type=Path,
+        help="Converted *_comfy.safetensors files. Defaults to all *_comfy.safetensors in cwd.",
+    )
+    return parser.parse_args()
+
+
+def read_subset(path: Path, limit_layers: int = 3) -> tuple[dict[str, torch.Tensor], dict[str, str]]:
+    state_dict: dict[str, torch.Tensor] = {}
+    with safe_open(path, framework="pt", device="cpu") as sf:
+        metadata = sf.metadata() or {}
+        layers = sorted(json.loads(metadata["_quantization_metadata"])["layers"])[:limit_layers]
+        for layer in layers:
+            for suffix in ("weight", "weight_scale", "weight_scale_2", "input_scale", "comfy_quant"):
+                key = f"{layer}.{suffix}"
+                state_dict[key] = sf.get_tensor(key)
+        return state_dict, metadata
+
+
+def main() -> None:
+    args = parse_args()
+    sys.path.insert(0, str(args.comfyui.resolve()))
+
+    import comfy.quant_ops as quant_ops
+    import comfy.utils as comfy_utils
+
+    nvfp4 = quant_ops.QUANT_ALGOS.get("nvfp4")
+    assert nvfp4 is not None, "ComfyUI source has no nvfp4 quant algo"
+    assert nvfp4["storage_t"] == torch.uint8, nvfp4
+    assert {"weight_scale", "weight_scale_2", "input_scale"} <= nvfp4["parameters"], nvfp4
+    assert nvfp4["comfy_tensor_layout"] == "TensorCoreNVFP4Layout", nvfp4
+
+    files = args.files or sorted(Path(".").glob("*_comfy.safetensors"))
+    for path in files:
+        state_dict, metadata = read_subset(path)
+        converted, _ = comfy_utils.convert_old_quants(dict(state_dict), metadata=metadata)
+        comfy_quant_keys = sorted(k for k in converted if k.endswith(".comfy_quant"))
+        if not comfy_quant_keys:
+            raise AssertionError(f"{path.name}: ComfyUI did not expose comfy_quant keys")
+        for key in comfy_quant_keys:
+            marker = json.loads(bytes(converted[key].tolist()).decode("utf-8"))
+            if marker.get("format") != "nvfp4":
+                raise AssertionError(f"{path.name}: {key} marker is {marker!r}")
+        print(f"OK {path.name}: ComfyUI source recognizes nvfp4 metadata and comfy_quant keys")
+
+
+if __name__ == "__main__":
+    main()