diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 00000000..dd84ea78 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,38 @@ +--- +name: Bug report +about: Create a report to help us improve +title: '' +labels: '' +assignees: '' + +--- + +**Describe the bug** +A clear and concise description of what the bug is. + +**To Reproduce** +Steps to reproduce the behavior: +1. Go to '...' +2. Click on '....' +3. Scroll down to '....' +4. See error + +**Expected behavior** +A clear and concise description of what you expected to happen. + +**Screenshots** +If applicable, add screenshots to help explain your problem. + +**Desktop (please complete the following information):** + - OS: [e.g. iOS] + - Browser [e.g. chrome, safari] + - Version [e.g. 22] + +**Smartphone (please complete the following information):** + - Device: [e.g. iPhone6] + - OS: [e.g. iOS8.1] + - Browser [e.g. stock browser, safari] + - Version [e.g. 22] + +**Additional context** +Add any other context about the problem here. diff --git a/.github/ISSUE_TEMPLATE/request-model-support.md b/.github/ISSUE_TEMPLATE/request-model-support.md new file mode 100644 index 00000000..835e8098 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/request-model-support.md @@ -0,0 +1,18 @@ +--- +name: Request model support +about: Request support for a new model architecture +title: '' +labels: enhancement +assignees: '' +type: Task + +--- + +### Model Details + +* **Model Name:** +* **Hugging Face / Model Source Link:** + +### Additional Context + +Provide any links to reference implementations, relevant pull requests, or technical details about this architecture, if you have one. diff --git a/.github/skills/add-model-architecture/SKILL.md b/.github/skills/add-model-architecture/SKILL.md new file mode 100644 index 00000000..59272d97 --- /dev/null +++ b/.github/skills/add-model-architecture/SKILL.md @@ -0,0 +1,83 @@ +--- +name: add-model-architecture +description: Add safe, tested support for a new ComfyUI model architecture and its GGUF conversion path. +--- + +# Add a Model Architecture + +Use this skill when adding or investigating a new diffusion, text, or vision +model architecture for ComfyUI-GGUF. + +## Goal + +Deliver end-to-end support: a source checkpoint is recognized, converted with +correct precision choices, accepted by this node's loader, and recognized by +the target ComfyUI installation. Do not label a converter-only match as +supported. + +## Discovery + +1. Identify the model role: diffusion model, text encoder, vision encoder, or + multimodal projector. +2. Obtain an authoritative checkpoint key listing with shapes and dtypes. Use + a safetensors header when possible; do not download model weights merely to + infer names. +3. Verify the current ComfyUI source can detect the model and instantiate its + runtime class. Record the minimum compatible ComfyUI revision if it is new. +4. Compare the key layout against every existing `Model*` class in + `tools/convert.py`. Reuse an existing architecture only when its detection, + precision rules, and runtime behavior all apply. + +## Conversion Support + +1. Add a focused `Model` subclass in `tools/convert.py`. +2. Set `arch` to the ComfyUI/GGUF architecture identifier expected at load + time. +3. Define `keys_detect` with multiple stable, distinctive keys. Use alternate + key sets only for known checkpoint export variants. +4. Add `keys_banned` when a similarly named incompatible checkpoint format + exists, such as a Diffusers export with incompatible fused projections. +5. Classify tensors before quantization: + - `keys_hiprec`: must remain FP32 due to numerical sensitivity, buffers, or + ComfyUI runtime requirements. + - `keys_noquant`: retain source FP16/BF16 because native low-bit execution + is unsafe or slower. + - `keys_ignore`: omit conversion-only state that is not model weight data. +6. Add the class to `arch_list`. Confirm `handle_tensors` preserves original + shapes and that all quantized dimensions satisfy the selected GGML block + size. + +## Loading Support + +1. Add the architecture to `IMG_ARCH_LIST`, `TXT_ARCH_LIST`, or + `VIS_TYPE_LIST` in `loader.py`, as appropriate. +2. Add a key mapper, tokenizer loader, or detection marker only when a real + naming or ComfyUI-detection mismatch requires it. Keep such transformations + deterministic and covered by a test. +3. Check standard and Dynamic VRAM loading. Dynamic loading preserves GGML + storage through `quant_ops.py`; static loading uses `ops.py`. +4. Do not enable `_K` diffusion quantization as a performance optimization: + this repository currently expands standard GGML quants before PyTorch + compute. Prefer `Q8_CR` for supported native INT8 Linear inference. + +## Tests and Documentation + +1. Add a synthetic test to `tests/test_targeted_quantization.py` that verifies + architecture detection from the distinctive keys. +2. Convert a minimal state dict and assert `general.architecture`, selected + GGML tensor types, and required FP32/FP16 exceptions. +3. Run `python -m unittest tests.test_targeted_quantization`. +4. Validate one real checkpoint in ComfyUI with a fixed workflow and inspect + loader logs for tensor types and unexpected-key failures. +5. Update `README.md` with supported model variants, minimum ComfyUI version, + conversion command, and quantization limitations. + +## Completion Checklist + +- [ ] ComfyUI support is present and its minimum version is documented. +- [ ] Detection uses distinctive keys and rejects incompatible formats. +- [ ] Sensitive tensors have explicit precision treatment. +- [ ] Static and Dynamic VRAM loaders accept the generated GGUF. +- [ ] Synthetic conversion tests pass. +- [ ] A real model loads and produces output in ComfyUI. +- [ ] User-facing documentation makes no unmeasured performance claim. diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 00000000..9b8932f7 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,19 @@ +name: Publish to Comfy Registry + +on: + workflow_dispatch: + release: + types: [published] + +jobs: + publish-node: + name: Publish Custom Node + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@v4 + + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@main + with: + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/.gitignore b/.gitignore index 44fe82fd..7f8e6d53 100644 --- a/.gitignore +++ b/.gitignore @@ -165,3 +165,8 @@ cython_debug/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ +.idea/comfyui-gguf-reboot.iml +.idea/inspectionProfiles/profiles_settings.xml +.idea/modules.xml +.idea/pyProjectModel.xml +.idea/vcs.xml diff --git a/.idea/.gitignore b/.idea/.gitignore new file mode 100644 index 00000000..30cf57ed --- /dev/null +++ b/.idea/.gitignore @@ -0,0 +1,10 @@ +# Default ignored files +/shelf/ +/workspace.xml +# Editor-based HTTP Client requests +/httpRequests/ +# Ignored default folder with query files +/queries/ +# Datasource local storage ignored files +/dataSources/ +/dataSources.local.xml diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..f3e8abce --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,56 @@ +# Repository Guide for Agents + +## Purpose + +ComfyUI-GGUF loads GGUF-encoded diffusion, text, and vision models into +ComfyUI. It also converts supported checkpoint layouts to GGUF. Preserve +ComfyUI compatibility and model output quality over file-size reductions. + +## Repository Map + +| Area | Responsibility | +| --- | --- | +| `loader.py` | Reads GGUF metadata and tensors, maps text/vision checkpoints, and selects loader behavior. | +| `ops.py` | Defines runtime tensor wrappers and on-the-fly dequantization or native INT8 execution. | +| `dequant.py` | PyTorch implementations of GGML block dequantizers. | +| `quant_ops.py` | Dynamic-VRAM `GGMLLayout` integration. | +| `tools/convert.py` | Detects checkpoint architectures and writes GGUF files. | +| `nodes.py` | ComfyUI node definitions and loading/conversion entry points. | +| `tests/test_targeted_quantization.py` | Unit and integration coverage for conversion and loader detection. | + +## Working Rules + +- Treat `general.architecture`, tensor names, tensor shapes, dtypes, and GGML + quantization types as compatibility contracts. Reject unsupported inputs with + clear errors rather than guessing. +- Add a model architecture in `tools/convert.py` only after confirming its + checkpoint key layout and that the installed ComfyUI can detect and run it. + A converter-only match is not usable model support. +- Protect non-Linear, numerically sensitive, and architecture-specific tensors + with `keys_hiprec` or `keys_noquant`. Do not quantize Conv2d weights merely + because they are two-dimensional after reshaping. +- Standard GGML quants are dequantized before PyTorch compute in this project. + Do not describe them as native low-bit inference. `Q8_CR` is the supported + native INT8 Linear path. +- Keep static and Dynamic VRAM behavior aligned. A new quantization type must + be supported by both `dequant.py` and `quant_ops.py`, or be rejected. +- Keep changes focused. Do not alter user-owned working-tree changes, generated + files, or model assets. + +## Validation + +Run the focused suite from the repository root when dependencies are available: + +```powershell +python -m unittest tests.test_targeted_quantization +``` + +For a new architecture, add a minimal synthetic checkpoint test that verifies +detection, intended protected-tensor precision, and GGUF metadata. Validate a +real checkpoint in ComfyUI before advertising support. + +## Documentation + +Update `README.md` when user-visible model support, conversion options, or +quantization behavior changes. Keep performance statements qualified by the +actual runtime path and hardware; do not publish unmeasured speed claims. diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md new file mode 100644 index 00000000..7b0dfb98 --- /dev/null +++ b/ARCHITECTURE.md @@ -0,0 +1,78 @@ +# ComfyUI-GGUF Architecture & Technical Details + +This document covers the engineering details behind ComfyUI-GGUF, including the custom `Q8_CR` native INT8 layout, target-size quantization fallback algorithms, dynamic patch behavior, and fallback caches for non-native adapter patches. + +## Native Weight-Only Quantization (Q8_CR) + +The converter supports one custom global quantization mode tailored for DiT/transformer UNets called `Q8_CR`. This is an INT8 weight-only format designed to reduce GGUF model storage and VRAM pressure while preserving the fast native INT8 Linear operations available on supported NVIDIA GPUs. It prevents severe VRAM bottlenecks when ComfyUI needs to offload weights to CPU memory. + +During conversion, the quantization pipeline performs the following sequence: +1. Selects eligible 2-D Linear weights while intentionally excluding one-dimensional tensors, small tensors, architecture-designated sensitive tensors, and Conv2d weights to maintain FP32/FP16 precision where required. +2. Applies the compatible ConvRot/Hadamard rotation to each eligible weight matrix. +3. Quantizes the rotated weights to INT8 using an FP32 scale for every output row. +4. Stores the INT8 payload, row scales, and ConvRot metadata directly in the compiled GGUF file. + +`Q8_CR` conversion accepts `--quantization-device auto`, `cpu`, or `cuda`. The `auto` flag prioritizes CUDA. If a matrix cannot fit in free VRAM, the converter logs a CPU fallback for that specific matrix without altering the overall output format. + +During load time, the GGUF loader reads the specialized metadata and passes the raw INT8 weights and row scales to ComfyUI's `TensorWiseINT8Layout`. On CUDA systems, ComfyUI natively executes the INT8/ConvRot Linear path directly without expanding the weight matrix to FP16. + +### Q8_CR Platform Support + +Q8_CR execution operates through ComfyUI's `comfy_kitchen` layout backend: +* NVIDIA CUDA triggers ComfyUI's optimized native INT8 backend automatically. +* Linux environments utilize the eager backend when CUDA is unavailable. +* Non-CUDA systems rely on the `comfy_kitchen` eager backend fallback. +* CPU Q8_CR loading and inference are fully supported but execute slower than hardware-accelerated CUDA passes. + +### Maintainer Recommendation for NVIDIA RTX 30-Series + +For Krea 2 and Ideogram 4 models on RTX 30-series architecture, `Q8_CR` offers significant benefits: +* Fast native INT8 operations routed through ComfyUI's ConvRot backend. +* Convenient CPU offload and memory-mapped model storage via the GGUF container. +* High image fidelity expected from 8-bit quantization while shielding sensitive tensors. +* Reduced VRAM pressure during complex multi-model spatial workflows. + +## Target-Size Quantization Algorithm + +Developers can utilize `tools/convert.py --max-size-mb ` to mandate the best supported mixed quantization below a strict output size ceiling. + +The fallback logic operates as follows: +1. Core 2-D Linear weights default to native INT8 ConvRot (`Q8_CR`), preserving all protected tensors in FP32. +2. Matrices closest to the model's core center drop to `Q5_0`. +3. If further reduction is required, those core matrices drop to `Q4_0`. +4. If the target size is still unmet after all core matrices reach `Q4_0`, standard 1-D tensors are reduced to BF16 (protected architecture tensors strictly remain FP32). + +`Q4_0` acts as the hard floor for core quantization. If a specified target falls below this theoretical minimum, the converter throws an error reporting the minimum achievable size. + +## LoRAs and Fused GGUF Exports + +**Load LoRA (GGUF)** integrates standard GGUF adapters directly into ComfyUI's patch mechanism. It natively parses `general.type=adapter` and `adapter.type=lora` alongside `.lora_a`/`.lora_b` tensors in F32, F16, BF16, or Q8_0. + +Imported GGUF LoRAs retain normal dynamic-patch behavior to ensure maximum compatibility. Because of this, an active LoRA prevents `Q8_CR` Linear layers from utilizing their native INT8 fast path. + +For `Q4_CR_W4A4`, compatible standard LoRA and LoKr patches keep the packed +INT4 base and add an exact BF16 low-rank output correction on every platform. +They do not fuse or cache a full patched weight. + +Other patch forms that cannot use either compatible LoRA/LoKr route fall back to a compute-dtype cache: +1. It dequantizes the target and applies the patch in the compute dtype. +2. It caches the resulting patched floating-point weight while the model remains loaded. +3. It evicts the derived caches when the model or patch layout changes. + +If a patched INT4 layer exhausts CUDA memory during execution, the layer retries in +system memory: the packed base is dequantized on the CPU, the adapter is applied +there, and only the completed output is copied back to the execution device. This +reduces the transient CUDA peak but cannot avoid the final output allocation needed +by the following GPU layer. + +The fallback cache remains floating-point because re-quantizing a patched matrix to INT4 can erase small deltas. Unsupported patch forms and CUDA OOM retries retain the existing floating-point fallback behavior. + +### Performance Diagnostics + +Performance logging is off unless `COMFYUI_GGUF_PERF_LOG` is set to `1`, a +truthy value, or a log path. It synchronizes CUDA before and after every +quantized Linear to produce per-layer wall-clock timings. This intentionally +serializes work and is unsuitable for normal inference or comparative +throughput benchmarks. + +For fixed adapter combinations, developers should merge adapters statically during export. Running `tools/convert.py --lora path/to/adapter.safetensors` fuses the parameters prior to quantization. Using the **Targeted Quantization (GGUF)** node's `streamed` input flag reads, fuses, quantizes, and stages one tensor block at a time to aggressively minimize peak RAM consumption. diff --git a/DESIGN.md b/DESIGN.md new file mode 100644 index 00000000..9a376d67 --- /dev/null +++ b/DESIGN.md @@ -0,0 +1,23 @@ +# Design + + + +## Conversion Dashboard + +The conversion dashboard is an operate-mode local workbench for long-running, +resource-intensive model conversion. + +- **Scene:** An operator is managing a local GPU workstation, often in a dim + environment, so the app uses a charcoal field rather than a generic light + admin surface. +- **Hierarchy:** A large condensed title anchors the tool. The left work area + is reserved for one new conversion, while the right queue exposes live + status and raw converter output without changing pages. +- **Color:** Off-white text on dark green-charcoal surfaces; electric lime + identifies the active action and running work. Success, cancellation, and + failure retain distinct high-contrast status colors. +- **Controls:** Square, bordered fields and flat status bands make paths, + settings, and console output feel like equipment labels rather than + decorative cards. +- **Responsive behavior:** The queue moves below the form below 900px, and + form controls become one column below 620px. diff --git a/PRODUCT.md b/PRODUCT.md new file mode 100644 index 00000000..1e1476d4 --- /dev/null +++ b/PRODUCT.md @@ -0,0 +1,34 @@ +# Product + + + +## Platform + +web + +## Users + +ComfyUI users who need to convert local diffusion-model checkpoints into GGUF files and want to monitor several conversions without repeatedly composing terminal commands. + +## Product Purpose + +ComfyUI-GGUF loads GGUF-encoded diffusion, text, and vision models into ComfyUI and converts supported local checkpoints to GGUF. The conversion dashboard makes the existing converter easier to operate while keeping files on the user's computer. + +## Operating Context + +Users run the dashboard locally from the ComfyUI-GGUF checkout, provide filesystem paths to large checkpoint files, select a quantization strategy, and wait for a memory- and GPU-intensive conversion to finish. + +## Capabilities and Constraints + +The dashboard invokes `tools/convert.py` with the active Python interpreter. It must stay dependency-free, bind only to localhost, accept filesystem paths instead of uploading model files, and run conversions serially to avoid resource contention. + +## Evidence on Hand + +The repository provides `tools/convert.py`, supported quantization types, target-size conversion, device selection, streamed safetensors conversion, and documented CLI examples in `README.md`. No product-specific visual assets are available. + +## Product Principles + +- Preserve the converter's behavior and error messages rather than reimplementing conversion. +- Keep large model files local and visible as paths. +- Make queued, running, failed, and completed work unambiguous. +- Prefer a small, dependable operator tool over a complex deployment. diff --git a/README.md b/README.md index 6915927e..9af3054b 100644 --- a/README.md +++ b/README.md @@ -1,49 +1,78 @@ + # ComfyUI-GGUF -GGUF Quantization support for native ComfyUI models -This is currently very much WIP. These custom nodes provide support for model files stored in the GGUF format popularized by [llama.cpp](https://github.com/ggerganov/llama.cpp). +GGUF Quantization support for native ComfyUI models including the custom Q8_CR. + +> [!NOTE] +> This is a fork of the original nodes, updated to support loading Ideogram 4 GGUFs and Krea 2 GGUFs. +> To use this maintained fork, clone `https://github.com/molbal/ComfyUI-GGUF`. -While quantization wasn't feasible for regular UNET models (conv2d), transformer/DiT models such as flux seem less affected by quantization. This allows running it in much lower bits per weight variable bitrate quants on low-end GPUs. For further VRAM savings, a node to load a quantized version of the T5 text encoder is also included. +While quantization was previously unfeasible for regular UNET models (conv2d), transformer/DiT models such as flux are less affected by quantization. This allows running them in lower bits per weight variable bitrate quants on GPUs with less VRAM. -![Comfy_Flux1_dev_Q4_0_GGUF_1024](https://github.com/user-attachments/assets/70d16d97-c522-4ef4-9435-633f128644c8) +More details on how to use it, pre-converted models, and sample workflows are here: [Documentation](https://molbal.github.io/gguf/ecosystem/using-the-custom-nodes.html) -Note: The "Force/Set CLIP Device" is **NOT** part of this node pack. Do not install it if you only have one GPU. Do not set it to cuda:0 then complain about OOM errors if you do not undestand what it is for. There is not need to copy the workflow above, just use your own workflow and replace the stock "Load Diffusion Model" with the "Unet Loader (GGUF)" node. +For technical details on the custom `Q8_CR` and `Q4_CR` formats , memory-mapped loading, please see [ARCHITECTURE.md](ARCHITECTURE.md). ## Installation > [!IMPORTANT] -> Make sure your ComfyUI is on a recent-enough version to support custom ops when loading the UNET-only. - -To install the custom node normally, git clone this repository into your custom nodes folder (`ComfyUI/custom_nodes`) and install the only dependency for inference (`pip install --upgrade gguf`) - -``` -git clone https://github.com/city96/ComfyUI-GGUF -``` +> Make sure your ComfyUI is on v0.27.0 or later. -To install the custom node on a standalone ComfyUI release, open a CMD inside the "ComfyUI_windows_portable" folder (where your `run_nvidia_gpu.bat` file is) and use the following commands: +To install the custom node normally, git clone this repository into your custom nodes folder (`ComfyUI/custom_nodes`) and restart ComfyUI. +```bash +git clone https://github.com/molbal/ComfyUI-GGUF ``` -git clone https://github.com/city96/ComfyUI-GGUF ComfyUI/custom_nodes/ComfyUI-GGUF -.\python_embeded\python.exe -s -m pip install -r .\ComfyUI\custom_nodes\ComfyUI-GGUF\requirements.txt -``` - -On MacOS sequoia, torch 2.4.1 seems to be required, as 2.6.X nightly versions cause a "M1 buffer is not large enough" error. See [this issue](https://github.com/city96/ComfyUI-GGUF/issues/107) for more information/workarounds. - -## Usage - -Simply use the GGUF Unet loader found under the `bootleg` category. Place the .gguf model files in your `ComfyUI/models/unet` folder. - -LoRA loading is experimental but it should work with just the built-in LoRA loader node(s). - -Pre-quantized models: - -- [flux1-dev GGUF](https://huggingface.co/city96/FLUX.1-dev-gguf) -- [flux1-schnell GGUF](https://huggingface.co/city96/FLUX.1-schnell-gguf) -- [stable-diffusion-3.5-large GGUF](https://huggingface.co/city96/stable-diffusion-3.5-large-gguf) -- [stable-diffusion-3.5-large-turbo GGUF](https://huggingface.co/city96/stable-diffusion-3.5-large-turbo-gguf) - -Initial support for quantizing T5 has also been added recently, these can be used using the various `*CLIPLoader (gguf)` nodes which can be used inplace of the regular ones. For the CLIP model, use whatever model you were using before for CLIP. The loader can handle both types of files - `gguf` and regular `safetensors`/`bin`. - -- [t5_v1.1-xxl GGUF](https://huggingface.co/city96/t5-v1_1-xxl-encoder-gguf) - -See the instructions in the [tools](https://github.com/city96/ComfyUI-GGUF/tree/main/tools) folder for how to create your own quants. + + + +## Usage + +Simply use the GGUF Unet loader found under the `bootleg` category. Place the .gguf model files in your `ComfyUI/models/unet` folder. + +Pre-quantized models (🍴 icon on ones added by this fork): + +- [flux1-dev GGUF](https://huggingface.co/city96/FLUX.1-dev-gguf) +- [flux1-schnell GGUF](https://huggingface.co/city96/FLUX.1-schnell-gguf) +- [stable-diffusion-3.5-large GGUF](https://huggingface.co/city96/stable-diffusion-3.5-large-gguf) +- [stable-diffusion-3.5-large-turbo GGUF](https://huggingface.co/city96/stable-diffusion-3.5-large-turbo-gguf) +- [Krea 2 (Both Turbo and Raw)](https://huggingface.co/molbal/krea2-gguf) 🍴 +- [Ideogram 4](https://huggingface.co/molbal/ideogram-4-gguf) 🍴 +- [MiniMax H3](https://huggingface.co/molbal/MiniMax-H3-GGUF) 🍴 +- [MiniMax Music3](https://huggingface.co/molbal/Minimax-Music3-GGUF) 🍴 +- [LTX 2.5](https://huggingface.co/molbal/LTX-2.5-GGUF) 🍴 + + +> [!IMPORTANT] > Please note, that this fork does not support _K quants on diffusion models, only on text encoders. They may or may not load, but inference speed may be very slow. There may be other forks, or other custom nodes with better support for these quantization types. + +Initial support for quantizing T5 has also been added recently, these can be used using the various `*CLIPLoader (gguf)` nodes which can be used inplace of the regular ones. For the CLIP model, use whatever model you were using before for CLIP. The loader can handle both types of files - `gguf` and regular `safetensors`/`bin`. + +- [t5_v1.1-xxl GGUF](https://huggingface.co/city96/t5-v1_1-xxl-encoder-gguf) +- [Qwen3-VL-4B-Instruct-GGUF](https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF) 🍴 +- [Qwen3-VL-32B-Instruct-GGUF](https://huggingface.co/unsloth/Qwen3-VL-32B-Instruct-GGUF) 🍴 +- [Qwen3-VL-32B-Instruct-MiniMax-H3 pruned GGUFs](https://huggingface.co/nif0/Qwen3-VL-32B-Instruct-MiniMax-H3-GGUF) 🍴 +- [Qwen3.5 GGUF](https://huggingface.co/unsloth/Qwen3.5-4B-GGUF) text encoders (0.8B, 2B, 4B, 9B, and 27B) with a ComfyUI build containing Qwen3.5 TE support. Place the matching `mmproj-*.gguf` beside the text encoder for image conditioning; text-only workflows do not need it. 🍴 +- [Gemma 4 GGUF](https://huggingface.co/unsloth/gemma-4-E4B-it-qat-GGUF) text encoders (E2B, E4B, 12B, and 31B) with ComfyUI v0.30.0 or later. 🍴 + +## Converting Models (Krea 2, Ideogram 4, MiniMax H3, MiniMax Music 3) + +This node pack includes a GGUF converter. It has 3 possible interfaces that you can use: +- a python file you can call directly +- a web interface +- a custom node + +Each option is documented here: [Quantizing models](https://molbal.github.io/gguf/ecosystem/quantizing-models.html) + +## Supported Conversion Formats + +| Format | Storage / execution | Recommended use | +|--------|---------------------------|-------------------------------------------------------------| +| F16 | FP16 GGUF | Maximum compatibility with half-precision storage. | +| BF16 | BF16 GGUF | Preserve BF16 source models where the target supports BF16. | +| Q8_0 | Standard GGML 8-bit | Excellent general-quality 8-bit GGUF. | +| Q5_1 | Standard GGML 5-bit | Lower storage with a quality-oriented 5-bit format. | +| Q5_0 | Standard GGML 5-bit | Lower storage alternative to `Q5_1`. | +| Q4_1 | Standard GGML 4-bit | Smaller files when VRAM or RAM is constrained. | +| Q4_0 | Standard GGML 4-bit | Smallest supported format for constrained setups. | +| Q8_CR | Per-row INT8 ConvRot | Maintainer recommendation for NVIDIA RTX 30-series systems. | +| Q4_CR | Experimental INT4 ConvRot | Maintainer recommendation for NVIDIA RTX 30-series systems. | \ No newline at end of file diff --git a/__init__.py b/__init__.py index a03726e3..9cd3727c 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,5 @@ +WEB_DIRECTORY = "./web" + # only import if running as a custom node try: import comfy.utils @@ -6,4 +8,4 @@ else: from .nodes import NODE_CLASS_MAPPINGS NODE_DISPLAY_NAME_MAPPINGS = {k:v.TITLE for k,v in NODE_CLASS_MAPPINGS.items()} - __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] diff --git a/dequant.py b/dequant.py index 78f5f266..67d6b1f0 100644 --- a/dequant.py +++ b/dequant.py @@ -1,7 +1,15 @@ # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) import gguf +import numpy as np import torch from tqdm import tqdm +from gguf.quants import ( + IQ2_S as _IQ2_S, + IQ2_XS as _IQ2_XS, + IQ2_XXS as _IQ2_XXS, + IQ3_S as _IQ3_S, + IQ3_XXS as _IQ3_XXS, +) TORCH_COMPATIBLE_QTYPES = (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16) @@ -18,6 +26,9 @@ def dequantize_tensor(tensor, dtype=None, dequant_dtype=None): if qtype in TORCH_COMPATIBLE_QTYPES: return tensor.to(dtype) + elif qtype == gguf.GGMLQuantizationType.BF16: + tensor = torch.Tensor(tensor.data.view(torch.bfloat16).reshape(oshape)) + return tensor if dtype is None or dtype == torch.bfloat16 else tensor.to(dtype) elif qtype in dequantize_functions: dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype return dequantize(tensor.data, qtype, oshape, dtype=dequant_dtype).to(dtype) @@ -240,6 +251,21 @@ def dequantize_blocks_Q2_K(blocks, block_size, type_size, dtype=None): # IQ quants KVALUES = torch.tensor([-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113], dtype=torch.int8) +def _get_iq_grid(iq_cls): + iq_cls.init_grid() + return torch.from_numpy(np.array(iq_cls.grid).squeeze().copy()) + +def _get_iq_ksigns(iq_cls): + iq_cls.init_grid() + return torch.from_numpy(np.frombuffer(iq_cls.ksigns, dtype=np.uint8).copy()) + +GRID_IQ3_S = _get_iq_grid(_IQ3_S) +GRID_IQ3_XXS = _get_iq_grid(_IQ3_XXS) +GRID_IQ2_S = _get_iq_grid(_IQ2_S) +GRID_IQ2_XS = _get_iq_grid(_IQ2_XS) +GRID_IQ2_XXS = _get_iq_grid(_IQ2_XXS) +KSIGNS_IQ2_XXS = _get_iq_ksigns(_IQ2_XXS) + def dequantize_blocks_IQ4_NL(blocks, block_size, type_size, dtype=None): n_blocks = blocks.shape[0] @@ -284,6 +310,156 @@ def dequantize_blocks_IQ4_XS(blocks, block_size, type_size, dtype=None): return (dl * qs).reshape((n_blocks, -1)) +def dequantize_blocks_IQ3_S(blocks, block_size, type_size, dtype=None): + n_blocks = blocks.shape[0] + + d, qs, qh, signs, scales = split_block_dims(blocks, 2, 64, 8, 32) + d = d.view(torch.float16).to(dtype) + + scales = scales.view(torch.uint8) + scales = torch.stack([scales & 0xF, scales >> 4], dim=-1).reshape((n_blocks, 8)) + db = d * (1 + 2 * scales.to(dtype)) + db = db.reshape((n_blocks, 8, 1, 1)) + + shifts = torch.arange(8, device=d.device, dtype=torch.uint8).reshape((1, 1, 8)) + signs = (signs.unsqueeze(-1) >> shifts) & 1 + signs = torch.where( + signs == 0, + torch.ones(1, dtype=dtype, device=d.device), + torch.full((1,), -1.0, dtype=dtype, device=d.device), + ) + signs = signs.reshape((n_blocks, 8, 8, 4)) + + qh_bits = (qh.unsqueeze(-1) >> shifts) & 1 + qh_bits = qh_bits.reshape((n_blocks, 64)) + qs = qs.to(torch.int16) | (qh_bits.to(torch.int16) << 8) + + grid = GRID_IQ3_S.to(dtype=dtype, device=d.device) + grid_val = grid[qs.to(torch.long)].reshape((n_blocks, 8, 8, 4)) + return (db * grid_val * signs).reshape((n_blocks, QK_K)) + +def dequantize_blocks_IQ3_XXS(blocks, block_size, type_size, dtype=None): + n_blocks = blocks.shape[0] + + d, qs, scales, _ = split_block_dims(blocks, 2, 64, 32) + d = d.view(torch.float16).to(dtype) + + scales = scales.reshape((n_blocks, 8, 4)).to(torch.int32) + scales = scales[:, :, 0] | scales[:, :, 1] << 8 | scales[:, :, 2] << 16 | scales[:, :, 3] << 24 + + db = d * (0.5 + ((scales >> 28) & 0xF).to(dtype)) * 0.5 + db = db.reshape((n_blocks, 8, 1, 1)) + + shifts = torch.tensor([0, 7, 14, 21], device=d.device, dtype=torch.int32).reshape((1, 1, 4)) + sign_indices = (scales.reshape((n_blocks, 8, 1)) >> shifts) & 0x7F + sign_bytes = KSIGNS_IQ2_XXS.to(d.device)[sign_indices.to(torch.long)] + + shifts_bits = torch.arange(8, device=d.device, dtype=torch.uint8).reshape((1, 1, 1, 8)) + signs = (sign_bytes.unsqueeze(-1) >> shifts_bits) & 1 + signs = torch.where( + signs == 0, + torch.ones(1, dtype=dtype, device=d.device), + torch.full((1,), -1.0, dtype=dtype, device=d.device), + ) + signs = signs.reshape((n_blocks, 8, 4, 8)) + + grid = GRID_IQ3_XXS.to(dtype=dtype, device=d.device) + grid_val = grid[qs.to(torch.long)].reshape((n_blocks, 8, 4, 8)) + return (db * grid_val * signs).reshape((n_blocks, QK_K)) + +def dequantize_blocks_IQ2_S(blocks, block_size, type_size, dtype=None): + n_blocks = blocks.shape[0] + + d, qs, signs, qh, scales = split_block_dims(blocks, 2, 32, 32, 8) + d = d.view(torch.float16).to(dtype) + + scales = scales.view(torch.uint8) + scales = torch.stack([scales & 0xF, scales >> 4], dim=-1).reshape((n_blocks, 16)) + db = d * (0.5 + scales.to(dtype)) * 0.25 + db = db.reshape((n_blocks, 16, 1, 1)) + + shifts = torch.arange(8, device=d.device, dtype=torch.uint8).reshape((1, 1, 8)) + signs = (signs.unsqueeze(-1) >> shifts) & 1 + signs = torch.where( + signs == 0, + torch.ones(1, dtype=dtype, device=d.device), + torch.full((1,), -1.0, dtype=dtype, device=d.device), + ) + signs = signs.reshape((n_blocks, 16, 2, 8)) + + qh_shifts = torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape((1, 1, 4)) + qh_bits = (qh.view(torch.uint8).reshape((n_blocks, 8, 1)) >> qh_shifts) & 3 + qh_bits = qh_bits.reshape((n_blocks, 32)) + qs = qs.view(torch.uint8).to(torch.int32) + indices = qs | (qh_bits.to(torch.int32) << 8) + + grid = GRID_IQ2_S.to(dtype=dtype, device=d.device) + grid_val = grid[indices.to(torch.long)].reshape((n_blocks, 16, 2, 8)) + return (db * grid_val * signs).reshape((n_blocks, QK_K)) + +def dequantize_blocks_IQ2_XS(blocks, block_size, type_size, dtype=None): + n_blocks = blocks.shape[0] + + d, qs, scales = split_block_dims(blocks, 2, 2 * QK_K // 8) + d = d.view(torch.float16).to(dtype) + + qs = qs.contiguous().reshape(n_blocks, 32, 2).to(torch.int32) + qs = qs[:, :, 0] | (qs[:, :, 1] << 8) + + shifts_sc = torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape(1, 1, 2) + sc = (scales.unsqueeze(-1) >> shifts_sc) & 0x0F + db = d.reshape(n_blocks, 1) * (0.5 + sc.reshape(n_blocks, 16).to(dtype)) * 0.25 + db = db.reshape(n_blocks, 16, 1, 1) + + sign_bytes = KSIGNS_IQ2_XXS.to(d.device)[(qs >> 9).to(torch.long)] + shifts_bits = torch.arange(8, device=d.device, dtype=torch.uint8).reshape(1, 1, 8) + signs = (sign_bytes.unsqueeze(-1) >> shifts_bits) & 1 + signs = torch.where( + signs == 0, + torch.ones(1, dtype=dtype, device=d.device), + torch.full((1,), -1.0, dtype=dtype, device=d.device), + ) + signs = signs.reshape(n_blocks, 16, 2, 8) + + grid = GRID_IQ2_XS.to(dtype=dtype, device=d.device) + grid_values = grid[(qs & 511).to(torch.long)] + grid_values = grid_values.reshape(n_blocks, 16, 2, 8) + return (db * grid_values * signs).reshape(n_blocks, QK_K) + +def dequantize_blocks_IQ2_XXS(blocks, block_size, type_size, dtype=None): + n_blocks = blocks.shape[0] + + d, qs = split_block_dims(blocks, 2) + d = d.view(torch.float16).to(dtype) + + u32 = qs.reshape((n_blocks, 16, 4)).to(torch.int32) + u32 = u32[:, :, 0] | (u32[:, :, 1] << 8) | (u32[:, :, 2] << 16) | (u32[:, :, 3] << 24) + u32 = u32.reshape((n_blocks, 8, 2)) + + q0 = u32[:, :, 0] + q1 = u32[:, :, 1] + + db = d * (0.5 + ((q1 >> 28) & 0xF).to(dtype)) * 0.25 + db = db.reshape((n_blocks, 8, 1, 1)) + + shifts = torch.tensor([0, 7, 14, 21], device=d.device, dtype=torch.int32).reshape((1, 1, 4)) + sign_indices = (q1.unsqueeze(-1) >> shifts) & 0x7F + sign_bytes = KSIGNS_IQ2_XXS.to(d.device)[sign_indices.to(torch.long)] + + shifts_bits = torch.arange(8, device=d.device, dtype=torch.uint8).reshape((1, 1, 1, 8)) + signs = (sign_bytes.unsqueeze(-1) >> shifts_bits) & 1 + signs = torch.where( + signs == 0, + torch.ones(1, dtype=dtype, device=d.device), + torch.full((1,), -1.0, dtype=dtype, device=d.device), + ) + signs = signs.reshape((n_blocks, 8, 4, 8)) + + indices = q0.contiguous().view(torch.uint8) + grid = GRID_IQ2_XXS.to(dtype=dtype, device=d.device) + grid_val = grid[indices.to(torch.long)].reshape((n_blocks, 8, 4, 8)) + return (db * grid_val * signs).reshape((n_blocks, QK_K)) + dequantize_functions = { gguf.GGMLQuantizationType.BF16: dequantize_blocks_BF16, gguf.GGMLQuantizationType.Q8_0: dequantize_blocks_Q8_0, @@ -298,4 +474,9 @@ def dequantize_blocks_IQ4_XS(blocks, block_size, type_size, dtype=None): gguf.GGMLQuantizationType.Q2_K: dequantize_blocks_Q2_K, gguf.GGMLQuantizationType.IQ4_NL: dequantize_blocks_IQ4_NL, gguf.GGMLQuantizationType.IQ4_XS: dequantize_blocks_IQ4_XS, + gguf.GGMLQuantizationType.IQ3_S: dequantize_blocks_IQ3_S, + gguf.GGMLQuantizationType.IQ3_XXS: dequantize_blocks_IQ3_XXS, + gguf.GGMLQuantizationType.IQ2_S: dequantize_blocks_IQ2_S, + gguf.GGMLQuantizationType.IQ2_XS: dequantize_blocks_IQ2_XS, + gguf.GGMLQuantizationType.IQ2_XXS: dequantize_blocks_IQ2_XXS, } diff --git a/docs/k-quant-inference.md b/docs/k-quant-inference.md new file mode 100644 index 00000000..74f0764f --- /dev/null +++ b/docs/k-quant-inference.md @@ -0,0 +1,112 @@ +# `_K` Quantization During Inference + +## Conclusion + +`_K` quants reduce GGUF storage and can reduce the resident compressed-weight +footprint. In this repository, they do **not** execute as native low-bit matrix +multiplication: standard GGML weights are expanded to the requested floating +compute dtype before each affected operation. Therefore `_K` diffusion models +can be substantially slower than `Q8_CR` and may be slower than simpler GGML +formats despite their smaller files. + +`_K` remains reasonable for text encoders when the compressed file or +CPU/offload footprint is the primary constraint and the one-time text encoding +latency is acceptable. It should not be presented as an inference-speed +optimization in this node. + +## What the Formats Save + +`_K` types encode 256 weights per super-block with additional per-sub-block +scales/minima. That improves quality at a given storage budget, but it makes +unpacking more complex. The GGML block definitions give these payload sizes: + +| Tensor type | Bytes per 256 weights | Payload bits/weight | Storage comparison | +| --- | ---: | ---: | --- | +| `Q2_K` | 84 | 2.625 | Much smaller than FP16 | +| `Q3_K` | 110 | 3.438 | Much smaller than FP16 | +| `Q4_K` | 144 | 4.500 | Same payload density as `Q4_0` | +| `Q5_K` | 176 | 5.500 | Same payload density as `Q5_0` | +| `Q6_K` | 210 | 6.563 | Smaller than `Q8_0` | +| `Q8_0` | 272 | 8.500 | Higher-quality conventional GGML baseline | +| FP16 | 512 | 16.000 | Uncompressed compute-weight baseline | + +The `_S`, `_M`, and `_L` suffixes in distribution filenames commonly describe +how a model mixes quantization choices across tensors; they are not a separate +single tensor encoding that this loader can execute differently. + +## Why This Node Can Be Slow + +The current runtime path is explicit: + +1. `GGMLLayer.cast_bias_weight()` in [`ops.py`](../ops.py) calls + `get_weight()` for every weighted operation. +2. `get_weight()` calls `dequantize_tensor()`. +3. `dequant.py` decodes `_K` blocks through PyTorch tensor operations: unpacking + bit fields, expanding scales/minima, and materializing floating-point + weights. +4. `torch.nn.functional.linear()` or the matching PyTorch operation then runs + on the expanded weight. + +Dynamic VRAM uses `GGMLLayout.dequantize()` in +[`quant_ops.py`](../quant_ops.py), which follows the same materialization model +for standard GGML types. The compressed data may remain mmap-backed until +needed, but the operation still needs a full floating-point temporary. + +In contrast, `Q8_CR` is converted to ComfyUI's +`TensorWiseINT8Layout`; [`get_gguf_q8_ops()`](../ops.py) retains INT8 weights +and uses ComfyUI's native INT8 Linear route. That avoids the generic GGML +dequantize-then-FP16/BF16-matmul sequence for eligible Linear layers. + +## Practical Implications + +| Scenario | `_K` result in this node | +| --- | --- | +| GGUF disk size / mmap-backed source weights | Lower, according to its payload bits per weight. | +| System RAM or VRAM while a layer is not resident | Often lower, especially with offload or Dynamic VRAM. | +| Peak working memory for an executing layer | Still requires a floating temporary; the largest layer is accounted for in `ops.py`. | +| Diffusion denoising latency | Usually unfavorable: every sampling step revisits many layers and repeats unpacking. | +| Text-encoder latency | May be acceptable because encoding occurs once per prompt, but measure it. | +| Output quality at a size budget | Often better than legacy quants of comparable payload size, but architecture- and model-dependent. | +| Native CUDA low-bit throughput | Not available through the standard `_K` path in this repository. | + +The actual outcome also depends on GPU, CPU, PCIe bandwidth, batch/sequence +size, ComfyUI offload policy, and whether the model is compute- or +transfer-bound. There is no defensible universal tokens/s or seconds/step +multiplier without a benchmark on the target workflow. + +## Recommended Choices + +- **Diffusion models on NVIDIA:** Prefer `Q8_CR` for eligible transformer/DiT + Linear weights when it fits the quality and compatibility target. +- **Portable diffusion GGUF:** Prefer the documented standard formats + (`Q8_0`, `Q5_0`, `Q4_0`) and choose file size versus output quality. Do not + select `_K` expecting faster samples. +- **Text encoders under memory pressure:** `_K` can be useful if its measured + prompt-encoding latency is acceptable. `Q4_K_M`/`Q5_K_M` are quality/storage + candidates, not speed recommendations. + +## Benchmark Protocol + +Compare only one variable at a time: use the same source model, ComfyUI +revision, workflow, seed, prompt, resolution, sampling steps, scheduler, +offload mode, and device placement. + +1. Warm up the workflow once to exclude compilation and initial allocation. +2. Run at least five measured generations for each quantization. +3. Record median wall-clock seconds per denoise step, total generation time, + text-encoder time, peak allocated/reserved VRAM, and process RAM. +4. Repeat once with full model residency and once with the intended offload or + Dynamic VRAM policy; transfer-bound behavior can reverse a result. +5. Compare outputs at the same seed for visible degradation before accepting a + smaller format. +6. Include loader logs with tensor type counts and document the hardware, + PyTorch, ComfyUI, and node revision with the result. + +## Sources + +- [GGUF specification](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md): + GGUF is mmap-compatible and stores model metadata and tensors. +- [GGML quantization reference implementation](https://github.com/ggml-org/llama.cpp/blob/master/ggml/src/ggml-quants.c): + defines the legacy and `_K` super-block encodings and their dequantization. +- Local implementation: [`dequant.py`](../dequant.py), + [`ops.py`](../ops.py), and [`quant_ops.py`](../quant_ops.py). diff --git a/editor.md b/editor.md new file mode 100644 index 00000000..870afe8e --- /dev/null +++ b/editor.md @@ -0,0 +1,60 @@ +# Prompt Canvas Editor + +A single-file browser editor for building Ideogram-style structured JSON prompts with canvas-based bounding boxes. The app lives entirely in `ui.html`; there is no build step, package manager, or local server requirement. + +## Features + +- Set canvas width and height with sliders or by double-clicking the displayed values. +- Draw, move, resize, delete, and edit bounding boxes directly on the canvas. +- Cycle through selected boxes with the `<-` and `->` controls when boxes overlap. +- Edit global prompt fields, including high-level description, aesthetics, lighting, medium, style/photo mode, background, and color palette. +- Edit per-box mode, description, optional text content, and per-box color palette. +- Generate formatted JSON from the current canvas and form state. +- Paste existing prompt JSON into the JSON box and load it back into the editable canvas. + +## Usage + +Open `ui.html` directly in a modern browser. + +The Tailwind design system is loaded from the Tailwind CDN, so the page needs internet access for styling. The editor logic itself is plain HTML, CSS, and JavaScript. + +## Basic Workflow + +1. Set the canvas size. +2. Draw boxes on the canvas by clicking and dragging. +3. Select a box and edit its properties in the right panel. +4. Fill in the global prompt settings. +5. Click `Generate JSON` to write the prompt JSON into the textarea. +6. Copy or save the generated JSON wherever your workflow needs it. + +To edit an existing prompt, paste the JSON into the textarea and click `Load JSON`. The editor will rebuild the canvas boxes and form fields from the prompt. + +## JSON Shape + +The editor expects prompt JSON in this general form: + +```json +{ + "high_level_description": "", + "style_description": { + "aesthetics": "", + "lighting": "", + "medium": "", + "art_style": "", + "color_palette": [] + }, + "compositional_deconstruction": { + "background": "", + "elements": [ + { + "type": "obj", + "bbox": [0, 0, 1000, 1000], + "desc": "", + "color_palette": [] + } + ] + } +} +``` + +Bounding boxes use normalized coordinates from `0` to `1000` in `[y1, x1, y2, x2]` order. The editor converts those coordinates to the current canvas size when loading JSON, then converts them back to normalized coordinates when generating JSON. \ No newline at end of file diff --git a/icon.png b/icon.png new file mode 100644 index 00000000..3cc47f8c Binary files /dev/null and b/icon.png differ diff --git a/loader.py b/loader.py index 7cefb113..d13591aa 100644 --- a/loader.py +++ b/loader.py @@ -1,506 +1,1094 @@ -# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) -import warnings -import logging -import torch -import gguf -import re -import os - -from .ops import GGMLTensor -from .dequant import is_quantized, dequantize_tensor - -IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "hidream", "cosmos", "ltxv", "hyvid", "wan", "lumina2", "qwen_image"} -TXT_ARCH_LIST = {"t5", "t5encoder", "llama", "qwen2vl", "qwen3", "qwen3vl", "gemma3"} -VIS_TYPE_LIST = {"clip-vision", "mmproj"} - -def get_orig_shape(reader, tensor_name): - field_key = f"comfy.gguf.orig_shape.{tensor_name}" - field = reader.get_field(field_key) - if field is None: - return None - # Has original shape metadata, so we try to decode it. - if len(field.types) != 2 or field.types[0] != gguf.GGUFValueType.ARRAY or field.types[1] != gguf.GGUFValueType.INT32: - raise TypeError(f"Bad original shape metadata for {field_key}: Expected ARRAY of INT32, got {field.types}") - return torch.Size(tuple(int(field.parts[part_idx][0]) for part_idx in field.data)) - -def get_field(reader, field_name, field_type): - field = reader.get_field(field_name) - if field is None: - return None - elif field_type == str: - # extra check here as this is used for checking arch string - if len(field.types) != 1 or field.types[0] != gguf.GGUFValueType.STRING: - raise TypeError(f"Bad type for GGUF {field_name} key: expected string, got {field.types!r}") - return str(field.parts[field.data[-1]], encoding="utf-8") - elif field_type in [int, float, bool]: - return field_type(field.parts[field.data[-1]].item()) - else: - raise TypeError(f"Unknown field type {field_type}") - -def get_list_field(reader, field_name, field_type): - field = reader.get_field(field_name) - if field is None: - return None - elif field_type == str: - return tuple(str(field.parts[part_idx], encoding="utf-8") for part_idx in field.data) - elif field_type in [int, float, bool]: - return tuple(field_type(field.parts[part_idx][0]) for part_idx in field.data) - else: - raise TypeError(f"Unknown field type {field_type}") - -def get_gguf_metadata(reader): - """Extract all simple metadata fields like safetensors""" - metadata = {} - for field_name in reader.fields: - try: - field = reader.get_field(field_name) - if len(field.types) == 1: # Simple scalar fields only - if field.types[0] == gguf.GGUFValueType.STRING: - metadata[field_name] = str(field.parts[field.data[-1]], "utf-8") - elif field.types[0] == gguf.GGUFValueType.INT32: - metadata[field_name] = int(field.parts[field.data[-1]]) - elif field.types[0] == gguf.GGUFValueType.F32: - metadata[field_name] = float(field.parts[field.data[-1]]) - elif field.types[0] == gguf.GGUFValueType.BOOL: - metadata[field_name] = bool(field.parts[field.data[-1]]) - except: - continue - return metadata - -def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", is_text_model=False): - """ - Read state dict as fake tensors - """ - reader = gguf.GGUFReader(path) - - # filter and strip prefix - has_prefix = False - if handle_prefix is not None: - prefix_len = len(handle_prefix) - tensor_names = set(tensor.name for tensor in reader.tensors) - has_prefix = any(s.startswith(handle_prefix) for s in tensor_names) - - tensors = [] - for tensor in reader.tensors: - sd_key = tensor_name = tensor.name - if has_prefix: - if not tensor_name.startswith(handle_prefix): - continue - sd_key = tensor_name[prefix_len:] - tensors.append((sd_key, tensor)) - - # detect and verify architecture - compat = None - arch_str = get_field(reader, "general.architecture", str) - type_str = get_field(reader, "general.type", str) - if arch_str in [None, "pig", "cow"]: - if is_text_model: - raise ValueError(f"This gguf file is incompatible with llama.cpp!\nConsider using safetensors or a compatible gguf file\n({path})") - compat = "sd.cpp" if arch_str is None else arch_str - # import here to avoid changes to convert.py breaking regular models - from .tools.convert import detect_arch - try: - arch_str = detect_arch(set(val[0] for val in tensors)).arch - except Exception as e: - raise ValueError(f"This model is not currently supported - ({e})") - elif arch_str not in TXT_ARCH_LIST and is_text_model: - if type_str not in VIS_TYPE_LIST: - raise ValueError(f"Unexpected text model architecture type in GGUF file: {arch_str!r}") - elif arch_str not in IMG_ARCH_LIST and not is_text_model: - raise ValueError(f"Unexpected architecture type in GGUF file: {arch_str!r}") - - if compat: - logging.warning(f"Warning: This gguf model file is loaded in compatibility mode '{compat}' [arch:{arch_str}]") - - # main loading loop - state_dict = {} - qtype_dict = {} - for sd_key, tensor in tensors: - tensor_name = tensor.name - # torch_tensor = torch.from_numpy(tensor.data) # mmap - - # NOTE: line above replaced with this block to avoid persistent numpy warning about mmap - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", message="The given NumPy array is not writable") - torch_tensor = torch.from_numpy(tensor.data) # mmap - - shape = get_orig_shape(reader, tensor_name) - if shape is None: - shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) - # Workaround for stable-diffusion.cpp SDXL detection. - if compat == "sd.cpp" and arch_str == "sdxl": - if any([tensor_name.endswith(x) for x in (".proj_in.weight", ".proj_out.weight")]): - while len(shape) > 2 and shape[-1] == 1: - shape = shape[:-1] - - # add to state dict - if tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: - torch_tensor = torch_tensor.view(*shape) - state_dict[sd_key] = GGMLTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape) - - # 1D tensors shouldn't be quantized, this is a fix for BF16 - if len(shape) <= 1 and tensor.tensor_type == gguf.GGMLQuantizationType.BF16: - state_dict[sd_key] = dequantize_tensor(state_dict[sd_key], dtype=torch.float32) - - # keep track of loaded tensor types - tensor_type_str = getattr(tensor.tensor_type, "name", repr(tensor.tensor_type)) - qtype_dict[tensor_type_str] = qtype_dict.get(tensor_type_str, 0) + 1 - - # print loaded tensor type counts - logging.info("gguf qtypes: " + ", ".join(f"{k} ({v})" for k, v in qtype_dict.items())) - - # mark largest tensor for vram estimation - qsd = {k:v for k,v in state_dict.items() if is_quantized(v)} - if len(qsd) > 0: - max_key = max(qsd.keys(), key=lambda k: qsd[k].numel()) - state_dict[max_key].is_largest_weight = True - - # extra info to return - extra = { - "arch_str": arch_str, - "metadata": get_gguf_metadata(reader) - } - return (state_dict, extra) - -# for remapping llama.cpp -> original key names -T5_SD_MAP = { - "enc.": "encoder.", - ".blk.": ".block.", - "token_embd": "shared", - "output_norm": "final_layer_norm", - "attn_q": "layer.0.SelfAttention.q", - "attn_k": "layer.0.SelfAttention.k", - "attn_v": "layer.0.SelfAttention.v", - "attn_o": "layer.0.SelfAttention.o", - "attn_norm": "layer.0.layer_norm", - "attn_rel_b": "layer.0.SelfAttention.relative_attention_bias", - "ffn_up": "layer.1.DenseReluDense.wi_1", - "ffn_down": "layer.1.DenseReluDense.wo", - "ffn_gate": "layer.1.DenseReluDense.wi_0", - "ffn_norm": "layer.1.layer_norm", -} - -LLAMA_SD_MAP = { - "blk.": "model.layers.", - "attn_norm": "input_layernorm", - "attn_q_norm.": "self_attn.q_norm.", - "attn_k_norm.": "self_attn.k_norm.", - "attn_v_norm.": "self_attn.v_norm.", - "attn_q": "self_attn.q_proj", - "attn_k": "self_attn.k_proj", - "attn_v": "self_attn.v_proj", - "attn_output": "self_attn.o_proj", - "ffn_up": "mlp.up_proj", - "ffn_down": "mlp.down_proj", - "ffn_gate": "mlp.gate_proj", - "ffn_norm": "post_attention_layernorm", - "token_embd": "model.embed_tokens", - "output_norm": "model.norm", - "output.weight": "lm_head.weight", -} - -GEMMA3_SD_MAP = LLAMA_SD_MAP.copy() -GEMMA3_SD_MAP.update({ - "ffn_norm": "pre_feedforward_layernorm", - "post_ffw_norm": "post_feedforward_layernorm", - "post_attention_norm": "post_attention_layernorm", -}) - -CLIP_VISION_SD_MAP = { - "mm.": "visual.merger.mlp.", - "v.post_ln.": "visual.merger.ln_q.", - "v.patch_embd": "visual.patch_embed.proj", - "v.blk.": "visual.blocks.", - "ffn_up": "mlp.up_proj", - "ffn_down": "mlp.down_proj", - "ffn_gate": "mlp.gate_proj", - "attn_out.": "attn.proj.", - "ln1.": "norm1.", - "ln2.": "norm2.", -} - -def sd_map_replace(raw_sd, key_map): - sd = {} - for k,v in raw_sd.items(): - for s,d in key_map.items(): - k = k.replace(s,d) - sd[k] = v - return sd - -def llama_permute(raw_sd, n_head, n_head_kv): - # Reverse version of LlamaModel.permute in llama.cpp convert script - sd = {} - permute = lambda x,h: x.reshape(h, x.shape[0] // h // 2, 2, *x.shape[1:]).swapaxes(1, 2).reshape(x.shape) - for k,v in raw_sd.items(): - if k.endswith(("q_proj.weight", "q_proj.bias")): - v.data = permute(v.data, n_head) - if k.endswith(("k_proj.weight", "k_proj.bias")): - v.data = permute(v.data, n_head_kv) - sd[k] = v - return sd - -def gemma3_norm_corrections(sd): - # Reverse change from Gemma3Model modify_tensors in llama.cpp convert script - norm_patterns = [ - "input_layernorm.weight", - "post_attention_layernorm.weight", - "pre_feedforward_layernorm.weight", - "post_feedforward_layernorm.weight", - "self_attn.q_norm.weight", - "self_attn.k_norm.weight", - "model.norm.weight" - ] - corrected = 0 - for key in list(sd.keys()): - if any(p in key for p in norm_patterns): - if is_quantized(sd[key]): - sd[key] = dequantize_tensor(sd[key], dtype=torch.float32) - 1.0 - else: - sd[key] = sd[key].float() - 1.0 - corrected += 1 - #logging.info(f"Gemma3: Applied -1 norm correction to {corrected} tensors") - return sd - -def strip_quant_suffix(name): - pattern = r"[-_]?(?:ud-)?i?q[0-9]_[a-z0-9_\-]{1,8}$" - match = re.search(pattern, name, re.IGNORECASE) - if match: - name = name[:match.start()] - return name - -def gguf_mmproj_loader(path): - # Reverse version of Qwen2VLVisionModel.modify_tensors - logging.info("Attenpting to find mmproj file for text encoder...") - - # get name to match w/o quant suffix - tenc_fname = os.path.basename(path) - tenc = os.path.splitext(tenc_fname)[0].lower() - tenc = strip_quant_suffix(tenc) - - # try and find matching mmproj - target = [] - root = os.path.dirname(path) - for fname in os.listdir(root): - name, ext = os.path.splitext(fname) - if ext.lower() != ".gguf": - continue - if "mmproj" not in name.lower(): - continue - if tenc in name.lower(): - target.append(fname) - - if len(target) == 0: - logging.error(f"Error: Can't find mmproj file for '{tenc_fname}' (matching:'{tenc}')! Qwen-Image-Edit will be broken!") - return {} - if len(target) > 1: - logging.error(f"Ambiguous mmproj for text encoder '{tenc_fname}', will use first match.") - - logging.info(f"Using mmproj '{target[0]}' for text encoder '{tenc_fname}'.") - target = os.path.join(root, target[0]) - vsd, _ = gguf_sd_loader(target, is_text_model=True) - - # concat 4D to 5D - if "v.patch_embd.weight.1" in vsd: - w1 = dequantize_tensor(vsd.pop("v.patch_embd.weight"), dtype=torch.float32) - w2 = dequantize_tensor(vsd.pop("v.patch_embd.weight.1"), dtype=torch.float32) - vsd["v.patch_embd.weight"] = torch.stack([w1, w2], dim=2) - - # run main replacement - vsd = sd_map_replace(vsd, CLIP_VISION_SD_MAP) - - # handle split Q/K/V - if "visual.blocks.0.attn_q.weight" in vsd: - attns = {} - # filter out attentions + group - for k,v in vsd.items(): - if any(x in k for x in ["attn_q", "attn_k", "attn_v"]): - k_attn, k_name = k.rsplit(".attn_", 1) - k_attn += ".attn.qkv." + k_name.split(".")[-1] - if k_attn not in attns: - attns[k_attn] = {} - attns[k_attn][k_name] = dequantize_tensor( - v, dtype=(torch.bfloat16 if is_quantized(v) else torch.float16) - ) - - # recombine - for k,v in attns.items(): - suffix = k.split(".")[-1] - vsd[k] = torch.cat([ - v[f"q.{suffix}"], - v[f"k.{suffix}"], - v[f"v.{suffix}"], - ], dim=0) - del attns - - return vsd - -def gguf_tokenizer_loader(path, temb_shape): - # convert gguf tokenizer to spiece - logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") - try: - from sentencepiece import sentencepiece_model_pb2 as model - except ImportError: - raise ImportError("Please make sure sentencepiece and protobuf are installed.\npip install sentencepiece protobuf") - spm = model.ModelProto() - - reader = gguf.GGUFReader(path) - - if get_field(reader, "tokenizer.ggml.model", str) == "t5": - if temb_shape == (256384, 4096): # probably UMT5 - spm.trainer_spec.model_type == 1 # Unigram (do we have a T5 w/ BPE?) - else: - raise NotImplementedError("Unknown model, can't set tokenizer!") - else: - raise NotImplementedError("Unknown model, can't set tokenizer!") - - spm.normalizer_spec.add_dummy_prefix = get_field(reader, "tokenizer.ggml.add_space_prefix", bool) - spm.normalizer_spec.remove_extra_whitespaces = get_field(reader, "tokenizer.ggml.remove_extra_whitespaces", bool) - - tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) - scores = get_list_field(reader, "tokenizer.ggml.scores", float) - toktypes = get_list_field(reader, "tokenizer.ggml.token_type", int) - - for idx, (token, score, toktype) in enumerate(zip(tokens, scores, toktypes)): - # # These aren't present in the original? - # if toktype == 5 and idx >= temb_shape[0]%1000): - # continue - - piece = spm.SentencePiece() - piece.piece = token - piece.score = score - piece.type = toktype - spm.pieces.append(piece) - - # unsure if any of these are correct - spm.trainer_spec.byte_fallback = True - spm.trainer_spec.vocab_size = len(tokens) # split off unused? - spm.trainer_spec.max_sentence_length = 4096 - spm.trainer_spec.eos_id = get_field(reader, "tokenizer.ggml.eos_token_id", int) - spm.trainer_spec.pad_id = get_field(reader, "tokenizer.ggml.padding_token_id", int) - - logging.info(f"Created tokenizer with vocab size of {len(spm.pieces)}") - del reader - return torch.ByteTensor(list(spm.SerializeToString())) - -def gguf_tekken_tokenizer_loader(path, temb_shape): - # convert ggml (hf) tokenizer metadata to tekken/comfy data - logging.info("Attempting to recreate tekken tokenizer from GGUF file metadata...") - import json - import base64 - from transformers.convert_slow_tokenizer import bytes_to_unicode - - reader = gguf.GGUFReader(path) - - model_str = get_field(reader, "tokenizer.ggml.model", str) - if model_str == "gpt2": - if temb_shape == (131072, 5120): # probably Mistral - data = { - "config": {"num_vocab_tokens": 150000, "default_vocab_size": 131072}, - "vocab": [], - "special_tokens": [], - } - else: - raise NotImplementedError("Unknown model, can't set tokenizer!") - else: - raise NotImplementedError("Unknown model, can't set tokenizer!") - - tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) - toktypes = get_list_field(reader, "tokenizer.ggml.token_type", int) - - decoder = {v: k for k, v in bytes_to_unicode().items()} - for idx, (token, toktype) in enumerate(zip(tokens, toktypes)): - if toktype == 3: - data["special_tokens"].append( - {'rank': idx, 'token_str': token, 'is_control': True} - ) - else: - tok = bytes([decoder[char] for char in token]) - data["vocab"].append({ - "rank": len(data["vocab"]), - "token_bytes": base64.b64encode(tok).decode("ascii"), - "token_str": tok.decode("utf-8", errors="replace") # ? - }) - - logging.info(f"Created tekken tokenizer with vocab size of {len(data['vocab'])} (+{len(data['special_tokens'])})") - del reader - return torch.ByteTensor(list(json.dumps(data).encode('utf-8'))) - -def gguf_gemma3_tokenizer_loader(path): - #TODO: merge into gguf_tokenizer_loader - logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") - try: - from sentencepiece import sentencepiece_model_pb2 as model - except ImportError: - raise ImportError("Please install sentencepiece and protobuf.\npip install sentencepiece protobuf") - spm = model.ModelProto() - reader = gguf.GGUFReader(path) - - spm.normalizer_spec.name = "identity" - spm.normalizer_spec.add_dummy_prefix = False - spm.trainer_spec.model_type = 2 - spm.trainer_spec.input_format = "tsv" - spm.trainer_spec.byte_fallback = True - spm.trainer_spec.max_sentence_length = 4192 - spm.trainer_spec.bos_piece = "" - - tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) - scores = get_list_field(reader, "tokenizer.ggml.scores", float) - toktype = get_list_field(reader, "tokenizer.ggml.token_type", int) - - if not tokens or not scores or not toktype: - raise ValueError("Missing tokenizer metadata") - - for idx in range(len(tokens)): - piece = spm.SentencePiece() - piece.piece = tokens[idx] - if idx == 3: # UNK position - piece.type = 2 # UNK Token - piece.score = 0.0 # UNK Score - else: - piece.type = toktype[idx] - piece.score = scores[idx] - spm.pieces.append(piece) - - spm.trainer_spec.vocab_size = len(spm.pieces) - logging.info(f"Created tokenizer with vocab size of {len(spm.pieces)}") - - del reader - return torch.ByteTensor(list(spm.SerializeToString())) - -def gguf_clip_loader(path): - sd, extra = gguf_sd_loader(path, is_text_model=True) - arch = extra.get("arch_str", None) - if arch in {"t5", "t5encoder"}: - temb_key = "token_embd.weight" - if temb_key in sd and sd[temb_key].shape == (256384, 4096): - # non-standard Comfy-Org tokenizer - sd["spiece_model"] = gguf_tokenizer_loader(path, sd[temb_key].shape) - # TODO: dequantizing token embed here is janky but otherwise we OOM due to tensor being massive. - logging.warning(f"Dequantizing {temb_key} to prevent runtime OOM.") - sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16) - sd = sd_map_replace(sd, T5_SD_MAP) - elif arch in {"llama", "qwen2vl", "qwen3", "qwen3vl", "gemma3"}: - # TODO: pass model_options["vocab_size"] to loader somehow - temb_key = "token_embd.weight" - if temb_key in sd and sd[temb_key].shape[0] >= (64 * 1024): - if arch == "llama" and sd[temb_key].shape == (131072, 5120): - # non-standard Comfy-Org tokenizer - sd["tekken_model"] = gguf_tekken_tokenizer_loader(path, sd[temb_key].shape) - elif arch == "gemma3": - sd["spiece_model"] = gguf_gemma3_tokenizer_loader(path) - # See note above for T5. - logging.warning(f"Dequantizing {temb_key} to prevent runtime OOM.") - sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16) - if arch == "gemma3": - sd = sd_map_replace(sd, GEMMA3_SD_MAP) - sd = gemma3_norm_corrections(sd) - else: - sd = sd_map_replace(sd, LLAMA_SD_MAP) - if arch == "llama": - sd = llama_permute(sd, 32, 8) # L3 / Mistral - if arch == "qwen2vl": - vsd = gguf_mmproj_loader(path) - sd.update(vsd) - else: - pass - return sd +# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import warnings +import logging +import torch +import gguf +import json +import re +import os +import threading +import comfy.memory_management +from .ops import GGMLTensor +from .dequant import is_quantized, dequantize_tensor +from .quant_ops import make_quantized + +IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "hidream", "cosmos", "ltxv", "ltxv_upscaler", "hyvid", "wan", "lumina2", "qwen_image", "ideogram", "krea2", "minimax_h3", "minimax_h3_vae", "minimax_music3"} +TXT_ARCH_LIST = {"t5", "t5encoder", "llama", "qwen2vl", "qwen3", "qwen3vl", "qwen35", "gemma3", "gemma4", "minimax_music3"} +VIS_TYPE_LIST = {"clip-vision", "mmproj"} +RAW_BYTE_TENSOR_KEYS = frozenset(("tokenizer_json", "spiece_model", "tekken_model")) + +def device_supports_bf16(): + """ + Return True if the active torch device can run bf16 natively. On devices + without native bf16 support, computation silently falls back to fp32 which + is very slow, so callers should load tensors as fp16 instead. + """ + try: + import comfy.model_management + return comfy.model_management.should_use_bf16(comfy.model_management.get_torch_device()) + except Exception: + # If support can't be determined, keep the previous bf16 behavior. + return True + + +def dynamic_gguf_file_slice(path): + """ + Create the file handle metadata used by DynamicVRAM to transfer a GGUF + tensor directly from its mmap-backed file to GPU memory. + """ + if not comfy.memory_management.aimdo_enabled: + return None + + import comfy_aimdo.model_mmap + + model_mmap = comfy_aimdo.model_mmap.ModelMMAP(path) + return model_mmap, threading.Lock() + + +def attach_dynamic_file_slice(torch_tensor, model_mmap, file_lock, offset, size): + storage = torch_tensor.untyped_storage() + storage._comfy_tensor_file_slice = comfy.memory_management.TensorFileSlice( + model_mmap.get_file_handle(), + file_lock, + offset, + size, + ) + # Keep the native mmap alive for the storage lifetime. + storage._comfy_tensor_mmap_refs = (model_mmap,) + +def get_orig_shape(reader, tensor_name): + field_key = f"comfy.gguf.orig_shape.{tensor_name}" + field = reader.get_field(field_key) + if field is None: + return None + # Has original shape metadata, so we try to decode it. + if len(field.types) != 2 or field.types[0] != gguf.GGUFValueType.ARRAY or field.types[1] != gguf.GGUFValueType.INT32: + raise TypeError(f"Bad original shape metadata for {field_key}: Expected ARRAY of INT32, got {field.types}") + return torch.Size(tuple(int(field.parts[part_idx][0]) for part_idx in field.data)) + +def get_field(reader, field_name, field_type): + field = reader.get_field(field_name) + if field is None: + return None + elif field_type == str: + # extra check here as this is used for checking arch string + if len(field.types) != 1 or field.types[0] != gguf.GGUFValueType.STRING: + raise TypeError(f"Bad type for GGUF {field_name} key: expected string, got {field.types!r}") + return str(field.parts[field.data[-1]], encoding="utf-8") + elif field_type in [int, float, bool]: + return field_type(field.parts[field.data[-1]].item()) + else: + raise TypeError(f"Unknown field type {field_type}") + +def get_list_field(reader, field_name, field_type): + field = reader.get_field(field_name) + if field is None: + return None + elif field_type == str: + return tuple(str(field.parts[part_idx], encoding="utf-8") for part_idx in field.data) + elif field_type in [int, float, bool]: + return tuple(field_type(field.parts[part_idx][0]) for part_idx in field.data) + else: + raise TypeError(f"Unknown field type {field_type}") + +def get_gguf_metadata(reader): + """Extract all simple metadata fields like safetensors""" + metadata = {} + for field_name in reader.fields: + try: + field = reader.get_field(field_name) + if len(field.types) == 1: # Simple scalar fields only + if field.types[0] == gguf.GGUFValueType.STRING: + metadata[field_name] = str(field.parts[field.data[-1]], "utf-8") + elif field.types[0] == gguf.GGUFValueType.INT32: + metadata[field_name] = int(field.parts[field.data[-1]]) + elif field.types[0] == gguf.GGUFValueType.F32: + metadata[field_name] = float(field.parts[field.data[-1]]) + elif field.types[0] == gguf.GGUFValueType.BOOL: + metadata[field_name] = bool(field.parts[field.data[-1]]) + except: + continue + return metadata + +def gguf_tensor_count(path): + return len(gguf.GGUFReader(path).tensors) + + +def normalize_raw_byte_tensor(value): + """Restore tokenizer payloads to the uint8 contract expected by ComfyUI.""" + if isinstance(value, GGMLTensor): + qtype = getattr(value, "tensor_type", None) + if qtype == gguf.GGMLQuantizationType.I8: + return value.data.view(torch.uint8).reshape(value.tensor_shape).contiguous() + if is_quantized(value): + value = dequantize_tensor(value, dtype=torch.float32) + else: + value = torch.Tensor(value).reshape(value.tensor_shape) + elif hasattr(value, "dequantize"): + value = value.dequantize() + + if not torch.is_tensor(value): + raise TypeError(f"Expected a tensor for tokenizer payload, got {type(value).__name__}") + return value.to(torch.uint8).contiguous() + + +def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", is_text_model=False, dynamic=False, progress_callback=None): + """ + Read state dict as fake tensors + """ + reader = gguf.GGUFReader(path) + dynamic_file_slice = dynamic_gguf_file_slice(path) if dynamic else None + + # filter and strip prefix + has_prefix = False + if handle_prefix is not None: + prefix_len = len(handle_prefix) + tensor_names = set(tensor.name for tensor in reader.tensors) + has_prefix = any(s.startswith(handle_prefix) for s in tensor_names) + + tensors = [] + for tensor in reader.tensors: + sd_key = tensor_name = tensor.name + if has_prefix: + if not tensor_name.startswith(handle_prefix): + continue + sd_key = tensor_name[prefix_len:] + tensors.append((sd_key, tensor)) + + # detect and verify architecture + compat = None + arch_str = get_field(reader, "general.architecture", str) + type_str = get_field(reader, "general.type", str) + if arch_str in [None, "pig", "cow"]: + if is_text_model: + raise ValueError(f"This gguf file is incompatible with llama.cpp!\nConsider using safetensors or a compatible gguf file\n({path})") + compat = "sd.cpp" if arch_str is None else arch_str + # import here to avoid changes to convert.py breaking regular models + from .tools.convert import detect_arch + try: + arch_str = detect_arch(set(val[0] for val in tensors)).arch + except Exception as e: + raise ValueError(f"This model is not currently supported - ({e})") + elif arch_str not in TXT_ARCH_LIST and is_text_model: + if type_str not in VIS_TYPE_LIST: + raise ValueError(f"Unexpected text model architecture type in GGUF file: {arch_str!r}") + elif arch_str not in IMG_ARCH_LIST and not is_text_model: + raise ValueError(f"Unexpected architecture type in GGUF file: {arch_str!r}") + + if compat: + logging.warning(f"Warning: This gguf model file is loaded in compatibility mode '{compat}' [arch:{arch_str}]") + + # Q8_CR weights must use ComfyUI's native TensorWiseINT8Layout rather than + # the generic GGML layout so DynamicVRAM retains native INT8 ConvRot kernels. + custom_quant_configs = {} + for field_name in reader.fields: + if field_name.startswith("comfy.gguf.quant."): + key = field_name[len("comfy.gguf.quant."):] + field = reader.get_field(field_name) + custom_quant_configs[key] = json.loads(str(field.parts[field.data[-1]], "utf-8")) + custom_quant_tensor_names = { + tensor_name + for key, quant_conf in custom_quant_configs.items() + if quant_conf.get("format") == "int8_tensorwise" + for tensor_name in (key, f"{key}_scale") + } | { + tensor_name + for key, quant_conf in custom_quant_configs.items() + if quant_conf.get("format") == "int4_cr" + for tensor_name in (key, f"{key}_scale") + } + + # main loading loop + # Devices without native bf16 fall back to slow fp32 compute, so load the + # full-precision BF16 storage tensors as fp16 there instead. + bf16_storage_dtype = torch.bfloat16 if device_supports_bf16() else torch.float16 + state_dict = {} + qtype_dict = {} + for tensor_index, (sd_key, tensor) in enumerate(tensors, start=1): + tensor_name = tensor.name + # torch_tensor = torch.from_numpy(tensor.data) # mmap + + # NOTE: line above replaced with this block to avoid persistent numpy warning about mmap + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="The given NumPy array is not writable") + torch_tensor = torch.from_numpy(tensor.data) # mmap + if dynamic_file_slice is not None: + model_mmap, file_lock = dynamic_file_slice + attach_dynamic_file_slice( + torch_tensor, + model_mmap, + file_lock, + tensor.data_offset, + tensor.n_bytes, + ) + + shape = get_orig_shape(reader, tensor_name) + if shape is None: + shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) + # Workaround for stable-diffusion.cpp SDXL detection. + if compat == "sd.cpp" and arch_str == "sdxl": + if any([tensor_name.endswith(x) for x in (".proj_in.weight", ".proj_out.weight")]): + while len(shape) > 2 and shape[-1] == 1: + shape = shape[:-1] + + # add to state dict + raw_byte_tensor = sd_key in RAW_BYTE_TENSOR_KEYS and len(shape) == 1 + if raw_byte_tensor and tensor.tensor_type == gguf.GGMLQuantizationType.I8: + # GGUF has no U8 type. I8 is used only as a byte-preserving + # container for tokenizer payloads. + state_dict[sd_key] = torch_tensor.view(torch.uint8).reshape(shape) + elif dynamic and sd_key not in custom_quant_tensor_names: + if tensor.tensor_type in { + gguf.GGMLQuantizationType.F32, + gguf.GGMLQuantizationType.F16, + }: + state_dict[sd_key] = torch_tensor.view(*shape) + elif tensor.tensor_type == gguf.GGMLQuantizationType.BF16: + state_dict[sd_key] = torch_tensor.view(torch.bfloat16).reshape(shape).to( + dtype=torch.float32 if len(shape) <= 1 else bf16_storage_dtype, + ) + else: + state_dict[sd_key] = make_quantized(torch_tensor, tensor.tensor_type, shape) + elif tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: + torch_tensor = torch_tensor.view(*shape) + state_dict[sd_key] = GGMLTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape) + else: + state_dict[sd_key] = GGMLTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape) + + # BF16 GGUF tensors are full-precision storage, not compressed quants. + if not dynamic and tensor.tensor_type == gguf.GGMLQuantizationType.BF16: + dtype = torch.float32 if len(shape) <= 1 else bf16_storage_dtype + state_dict[sd_key] = dequantize_tensor(state_dict[sd_key], dtype=dtype) + + # keep track of loaded tensor types + tensor_type_str = getattr(tensor.tensor_type, "name", repr(tensor.tensor_type)) + qtype_dict[tensor_type_str] = qtype_dict.get(tensor_type_str, 0) + 1 + if progress_callback is not None: + progress_callback(tensor_index, len(tensors)) + + # print loaded tensor type counts + logging.info("gguf qtypes: " + ", ".join(f"{k} ({v})" for k, v in qtype_dict.items())) + + # mark largest tensor for vram estimation + qsd = {k:v for k,v in state_dict.items() if is_quantized(v)} + if len(qsd) > 0: + max_key = max(qsd.keys(), key=lambda k: qsd[k].numel()) + state_dict[max_key].is_largest_weight = True + + # extra info to return + extra = { + "arch_str": arch_str, + "metadata": get_gguf_metadata(reader) + } + + # Detect custom ComfyUI native quantization metadata + warned_unrotated_convrot = False + for field_name in reader.fields: + if not field_name.startswith("comfy.gguf.quant."): + continue + key = field_name[len("comfy.gguf.quant."):] + field = reader.get_field(field_name) + quant_conf = custom_quant_configs[key] + fmt = quant_conf.get("format") + + if fmt in {"int4_compact_gemm", "int4_pytorch"}: + raise ValueError( + "Q4_PT GGUF files are retired because PyTorch's Ampere INT4 " + "kernel is not performance-competitive. Reconvert as Q8_CR." + ) + + weight_key = key + scale_key = f"{key}_scale" + if weight_key not in state_dict or scale_key not in state_dict: + logging.warning(f"Missing custom quant tensors for {weight_key}") + continue + + weight_ggml = state_dict[weight_key] + scale_ggml = state_dict[scale_key] + + if fmt == "int8_tensorwise": + if quant_conf.get("convrot") and not quant_conf.get("weight_rotated", False): + if not warned_unrotated_convrot: + logging.warning( + "Disabling ConvRot because this GGUF does not mark its weights " + "as pre-rotated. Reconvert with the current converter to enable ConvRot." + ) + warned_unrotated_convrot = True + quant_conf["convrot"] = False + quant_conf.pop("convrot_groupsize", None) + elif quant_conf.get("convrot"): + groupsize = quant_conf.get("convrot_groupsize", 256) + weight_shape = weight_ggml.shape if dynamic else weight_ggml.tensor_shape + if weight_shape[-1] % groupsize != 0: + logging.warning( + "Disabling ConvRot for %s because %d input features are not " + "divisible by group size %d.", + weight_key, + weight_shape[-1], + groupsize, + ) + quant_conf["convrot"] = False + quant_conf.pop("convrot_groupsize", None) + + # Convert to ComfyUI native state-dict layout + if dynamic: + weight = weight_ggml.view(torch.int8).reshape(weight_ggml.shape) + scale = scale_ggml.view(torch.float32).reshape(scale_ggml.shape) + else: + weight = weight_ggml.data.view(torch.int8).reshape(weight_ggml.tensor_shape) + scale = scale_ggml.data.view(torch.float32).reshape(scale_ggml.tensor_shape) + + state_dict[weight_key] = torch.nn.Parameter(weight, requires_grad=False) + state_dict[scale_key] = torch.nn.Parameter(scale, requires_grad=False) + + layer_prefix = weight_key[:weight_key.rfind("weight")] + quant_json = json.dumps(quant_conf) + state_dict[f"{layer_prefix}comfy_quant"] = torch.nn.Parameter( + torch.tensor(list(quant_json.encode("utf-8")), dtype=torch.uint8), + requires_grad=False, + ) + extra["gguf_quant_mode"] = "int8_convrot" + + # Q4_CR_W4A4: custom W4A4 INT4 backed by comfy_kitchen's fast ConvRot + # int4 tensor-core MMA. On-disk is kitchen-native packed int4 (N, K//2) + # int8 + a per-output-row fp16 scale; no per-group scales/zeros. + # The weight is pre-rotated by a block-diagonal Hadamard along K, matching + # the activation rotation the kernel applies at runtime. + elif fmt == "int4_cr" and quant_conf.get("backing") == "w4a4": + orig_shape = torch.Size(tuple(quant_conf["orig_shape"])) + convrot_groupsize = quant_conf.get("convrot_groupsize", 256) + quant_group_size = quant_conf.get("quant_group_size", 64) + packed_shape = (orig_shape[0], orig_shape[1] // 2) + + if dynamic: + packed = weight_ggml.view(torch.int8).reshape(packed_shape) + scale = scale_ggml.view(torch.float16).to(torch.bfloat16).reshape(orig_shape[0]) + else: + packed = weight_ggml.data.view(torch.int8).reshape(packed_shape) + scale = scale_ggml.data.view(torch.float16).to(torch.bfloat16).reshape(orig_shape[0]) + + layer_prefix = weight_key[:weight_key.rfind("weight")] + + state_dict[weight_key] = torch.nn.Parameter(packed, requires_grad=False) + state_dict[scale_key] = torch.nn.Parameter(scale, requires_grad=False) + state_dict[f"{layer_prefix}comfy_quant"] = torch.nn.Parameter( + torch.tensor(list(json.dumps(quant_conf).encode("utf-8")), dtype=torch.uint8), + requires_grad=False, + ) + extra["gguf_quant_mode"] = "int4_cr_w4a4" + + elif fmt == "int4_cr": + # The retired W4A16 (AWQ GEMV) Q4_CR format is no longer supported. + raise ValueError( + "This Q4_CR GGUF uses the retired W4A16 backing (no 'backing': 'w4a4'). " + "Reconvert with --quant-type Q4_CR_W4A4." + ) + + return (state_dict, extra) + +# for remapping llama.cpp -> original key names +T5_SD_MAP = { + "enc.": "encoder.", + ".blk.": ".block.", + "token_embd": "shared", + "output_norm": "final_layer_norm", + "attn_q": "layer.0.SelfAttention.q", + "attn_k": "layer.0.SelfAttention.k", + "attn_v": "layer.0.SelfAttention.v", + "attn_o": "layer.0.SelfAttention.o", + "attn_norm": "layer.0.layer_norm", + "attn_rel_b": "layer.0.SelfAttention.relative_attention_bias", + "ffn_up": "layer.1.DenseReluDense.wi_1", + "ffn_down": "layer.1.DenseReluDense.wo", + "ffn_gate": "layer.1.DenseReluDense.wi_0", + "ffn_norm": "layer.1.layer_norm", +} + +LLAMA_SD_MAP = { + "blk.": "model.layers.", + "attn_norm": "input_layernorm", + "attn_q_norm.": "self_attn.q_norm.", + "attn_k_norm.": "self_attn.k_norm.", + "attn_v_norm.": "self_attn.v_norm.", + "attn_q": "self_attn.q_proj", + "attn_k": "self_attn.k_proj", + "attn_v": "self_attn.v_proj", + "attn_output": "self_attn.o_proj", + "ffn_up": "mlp.up_proj", + "ffn_down": "mlp.down_proj", + "ffn_gate": "mlp.gate_proj", + "ffn_norm": "post_attention_layernorm", + "token_embd": "model.embed_tokens", + "output_norm": "model.norm", + "output.weight": "lm_head.weight", +} + +GEMMA3_SD_MAP = LLAMA_SD_MAP.copy() +GEMMA3_SD_MAP.update({ + "ffn_norm": "pre_feedforward_layernorm", + "post_ffw_norm": "post_feedforward_layernorm", + "post_attention_norm": "post_attention_layernorm", +}) + +# These specific keys must precede the generic token_embd mapping below. +GEMMA4_SD_MAP = { + "per_layer_token_embd": "model.embed_tokens_per_layer", + "per_layer_model_proj": "model.per_layer_model_projection", + "per_layer_proj_norm": "model.per_layer_projection_norm", + "inp_gate": "per_layer_input_gate", + "layer_output_scale.weight": "layer_scalar", + "post_norm": "post_per_layer_input_norm", + ".proj.": ".per_layer_projection.", + **GEMMA3_SD_MAP, +} + +# Qwen3.5 (llama.cpp ``qwen35`` arch). ComfyUI's detect_te_model() identifies +# Qwen3.5 by the unprefixed ``model.language_model.*`` layout before applying +# its own prefix rename, so the LM must map with that prefix. Hybrid layers +# carry linear attention (``attn_qkv``/``attn_gate``/``ssm_*`` tensors); the +# second layernorm is exported as ``post_attention_norm`` (no ``ffn_norm``). +# Entry order matters because sd_map_replace() does substring replacement: +# the fused/SSM variants must precede the generic ``attn_q``/``ssm_a`` keys. +QWEN35_SD_MAP = { + "blk.": "model.language_model.layers.", + "attn_norm": "input_layernorm", + "attn_q_norm.": "self_attn.q_norm.", + "attn_k_norm.": "self_attn.k_norm.", + "attn_qkv": "linear_attn.in_proj_qkv", + "attn_gate": "linear_attn.in_proj_z", + "attn_q": "self_attn.q_proj", + "attn_k": "self_attn.k_proj", + "attn_v": "self_attn.v_proj", + "attn_output": "self_attn.o_proj", + "ssm_alpha": "linear_attn.in_proj_a", + "ssm_beta": "linear_attn.in_proj_b", + "ssm_conv1d": "linear_attn.conv1d", + "ssm_dt.bias": "linear_attn.dt_bias", + "ssm_norm": "linear_attn.norm", + "ssm_out": "linear_attn.out_proj", + "ssm_a": "linear_attn.A_log", + "post_attention_norm": "post_attention_layernorm", + "ffn_up": "mlp.up_proj", + "ffn_down": "mlp.down_proj", + "ffn_gate": "mlp.gate_proj", + "token_embd": "model.language_model.embed_tokens", + "output_norm": "model.language_model.norm", + "output.weight": "lm_head.weight", +} + +CLIP_VISION_SD_MAP = { + "mm.": "visual.merger.mlp.", + "v.post_ln.": "visual.merger.ln_q.", + "v.patch_embd": "visual.patch_embed.proj", + "v.blk.": "visual.blocks.", + "ffn_up": "mlp.up_proj", + "ffn_down": "mlp.down_proj", + "ffn_gate": "mlp.gate_proj", + "attn_out.": "attn.proj.", + "ln1.": "norm1.", + "ln2.": "norm2.", +} + +CLIP_VISION_QWEN3_MAP = { + "v.blk": "model.visual.blocks", + ".fc": ".linear_fc", + "ck.8.": "st.0.", + "ck.16.": "st.1.", + "ck.24.": "st.2.", + "ck.5.": "st.0.", + "ck.11.": "st.1.", + "ck.17.": "st.2.", + "attn_out": "attn.proj", + "ln1": "norm1", + "ln2": "norm2", + "attn_qkv": "attn.qkv", + "ffn_up": "mlp.linear_fc1", + "ffn_down": "mlp.linear_fc2", + "mm.0": "model.visual.merger.linear_fc1", + "mm.2": "model.visual.merger.linear_fc2", + "v.post_ln": "model.visual.merger.norm", + "v.patch_embd": "model.visual.patch_embed.proj", + "v.position_embd.weight": "visual.pos_embed.weight", + "v.deepstack.": "model.visual.deepstack_merger_list.", + # Older llama.cpp Qwen3-VL exporters misspelled this tensor prefix. + "v.deepstast.": "model.visual.deepstack_merger_list.", +} + +def sd_map_replace(raw_sd, key_map): + sd = {} + for k,v in raw_sd.items(): + for s,d in key_map.items(): + k = k.replace(s,d) + sd[k] = v + return sd + +def llama_permute(raw_sd, n_head, n_head_kv): + # Reverse version of LlamaModel.permute in llama.cpp convert script + sd = {} + permute = lambda x,h: x.reshape(h, x.shape[0] // h // 2, 2, *x.shape[1:]).swapaxes(1, 2).reshape(x.shape) + for k,v in raw_sd.items(): + if k.endswith(("q_proj.weight", "q_proj.bias")): + v.data = permute(v.data, n_head) + if k.endswith(("k_proj.weight", "k_proj.bias")): + v.data = permute(v.data, n_head_kv) + sd[k] = v + return sd + +def gemma3_norm_corrections(sd): + # Reverse change from Gemma3Model modify_tensors in llama.cpp convert script + norm_patterns = [ + "input_layernorm.weight", + "post_attention_layernorm.weight", + "pre_feedforward_layernorm.weight", + "post_feedforward_layernorm.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + "model.norm.weight" + ] + corrected = 0 + for key in list(sd.keys()): + if any(p in key for p in norm_patterns): + if is_quantized(sd[key]): + sd[key] = dequantize_tensor(sd[key], dtype=torch.float32) - 1.0 + else: + sd[key] = sd[key].float() - 1.0 + corrected += 1 + #logging.info(f"Gemma3: Applied -1 norm correction to {corrected} tensors") + return sd + +def _qwen35_v_reorder(value, num_v_heads, num_k_heads, head_dim, qk_rows=0, axis=0): + """Reverse llama.cpp's tiled V-head order for linear attention. + + llama.cpp (``_LinearAttentionVReorderBase``) stores V heads in tiled + order ``[K0_v0, K1_v0, ..., K0_v1, K1_v1, ...]`` so its kernels can + broadcast K with ``ggml_repeat``. ComfyUI expects the grouped-by-K-head + order ``[K0_v0..v{r-1}, K1_v0..v{r-1}, ...]`` that the HF checkpoints + use. Quantized tensors are dequantized first because the head blocks + do not align with every GGML quant block size. ``qk_rows`` skips the + untouched Q/K rows of fused tensors; ``axis=1`` reorders columns. + """ + if num_k_heads <= 0 or num_v_heads <= num_k_heads: + return value + if is_quantized(value): + dtype = torch.bfloat16 if device_supports_bf16() else torch.float16 + value = dequantize_tensor(value, dtype=dtype) + r = num_v_heads // num_k_heads + # grouped head g = (k_idx, v_idx); its tiled position is v_idx * k + k_idx + src = [(g % r) * num_k_heads + (g // r) for g in range(num_v_heads)] + idx = torch.tensor(src, dtype=torch.long) + if value.ndim == 1: + return value.index_select(0, idx) + if qk_rows > 0: + # fused qkv / conv1d: only the V block is reordered + head_rows = value.shape[0] - qk_rows + v = value[qk_rows:].view(num_v_heads, -1).index_select(0, idx) + v = v.reshape(head_rows, value.shape[1]) + return torch.cat([value[:qk_rows], v], dim=0) + if axis == 1: + return ( + value.view(-1, num_v_heads, head_dim) + .index_select(1, idx) + .reshape(value.shape) + ) + return value.view(num_v_heads, -1).index_select(0, idx).reshape(value.shape) + +def qwen35_corrections(sd): + # Reverse llama.cpp's Qwen3.5 conversion tweaks (see Qwen3NextModel and + # _LinearAttentionVReorderBase modify_tensors in llama.cpp): it stores + # RMS norm weights as (w + 1), stores ``ssm_a`` as ``-exp(A_log)``, + # squeezes the depthwise conv1d kernels to 2D, and stores linear-attention + # V heads in tiled order. ComfyUI's Qwen35 TE expects the raw weights and + # applies ``w + 1`` and ``-exp(A_log)`` itself. + norm_patterns = [ + "input_layernorm.weight", + "post_attention_layernorm.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + "model.language_model.norm.weight", + ] + corrected = 0 + for key in list(sd.keys()): + if any(p in key for p in norm_patterns): + if is_quantized(sd[key]): + sd[key] = dequantize_tensor(sd[key], dtype=torch.float32) - 1.0 + else: + sd[key] = sd[key].float() - 1.0 + corrected += 1 + + # derive the linear-attention head layout from the checkpoint shapes + qkv_key = next((k for k in sd if k.endswith(".linear_attn.in_proj_qkv.weight")), None) + z_key = next((k for k in sd if k.endswith(".linear_attn.in_proj_z.weight")), None) + alog_key = next((k for k in sd if k.endswith(".linear_attn.A_log")), None) + num_k_heads = num_v_heads = head_dim = 0 + if qkv_key and z_key and alog_key: + value_dim = sd[z_key].shape[0] + conv_dim = sd[qkv_key].shape[0] + key_dim = (conv_dim - value_dim) // 2 + num_v_heads = sd[alog_key].shape[0] + head_dim = value_dim // num_v_heads + num_k_heads = key_dim // head_dim + + for key in list(sd.keys()): + if key.endswith(".linear_attn.A_log"): + # llama.cpp stores ``ssm_a = -exp(A_log)`` (always negative); + # ComfyUI computes ``-A_log.exp()``, so invert back to A_log. + value = sd[key] + if is_quantized(value): + value = dequantize_tensor(value, dtype=torch.float32) + else: + value = value.float() + sd[key] = torch.log(-value) + corrected += 1 + elif key.endswith(".linear_attn.dt_bias"): + corrected += 1 + elif key.endswith(".linear_attn.conv1d.weight"): + corrected += 1 + elif key.endswith(".linear_attn.in_proj_qkv.weight"): + corrected += 1 + elif key.endswith(( + ".linear_attn.in_proj_z.weight", + ".linear_attn.in_proj_a.weight", + ".linear_attn.in_proj_b.weight", + )): + corrected += 1 + elif key.endswith(".linear_attn.out_proj.weight"): + corrected += 1 + else: + continue + + # reverse the tiled V-head order (identity when heads are balanced) + if key.endswith((".linear_attn.A_log", ".linear_attn.dt_bias")): + sd[key] = _qwen35_v_reorder(sd[key], num_v_heads, num_k_heads, 1) + elif key.endswith(".linear_attn.conv1d.weight"): + value = sd[key] + if value.ndim == 2: + conv_dim = value.shape[0] + qk_channels = conv_dim - num_v_heads * head_dim + # llama.cpp squeezes the depthwise conv1d kernel to 2D + # (out_channels, kernel_size); ComfyUI's Conv1d expects + # (out_channels, 1, kernel_size). + value = _qwen35_v_reorder( + value, num_v_heads, num_k_heads, head_dim, qk_rows=qk_channels + ) + sd[key] = value.unsqueeze(-2) + elif key.endswith(".linear_attn.in_proj_qkv.weight"): + value = sd[key] + qk_rows = value.shape[0] - num_v_heads * head_dim + sd[key] = _qwen35_v_reorder( + value, num_v_heads, num_k_heads, head_dim, qk_rows=qk_rows + ) + elif key.endswith(".linear_attn.out_proj.weight"): + sd[key] = _qwen35_v_reorder( + sd[key], num_v_heads, num_k_heads, head_dim, axis=1 + ) + else: + sd[key] = _qwen35_v_reorder(sd[key], num_v_heads, num_k_heads, 1) + + logging.info(f"qwen35 GGUF: corrected {corrected} norm/A_log/conv1d/V-head tensors") + return sd + +def strip_quant_suffix(name): + pattern = r"[-_]?(?:ud-)?i?q[0-9]_[a-z0-9_\-]{1,8}$" + match = re.search(pattern, name, re.IGNORECASE) + if match: + name = name[:match.start()] + return name + +def gguf_mmproj_loader(path): + # Reverse version of Qwen2VLVisionModel.modify_tensors + logging.info("Attenpting to find mmproj file for text encoder...") + + # get name to match w/o quant suffix + tenc_fname = os.path.basename(path) + tenc = os.path.splitext(tenc_fname)[0].lower() + tenc = strip_quant_suffix(tenc) + + # try and find matching mmproj + target = [] + root = os.path.dirname(path) + for fname in os.listdir(root): + name, ext = os.path.splitext(fname) + if ext.lower() != ".gguf": + continue + if "mmproj" not in name.lower(): + continue + if tenc in name.lower(): + target.append(fname) + + if len(target) == 0: + logging.error(f"Error: Can't find mmproj file for '{tenc_fname}' (matching:'{tenc}')! Qwen-Image-Edit will be broken!") + return {} + if len(target) > 1: + logging.error(f"Ambiguous mmproj for text encoder '{tenc_fname}', will use first match.") + + logging.info(f"Using mmproj '{target[0]}' for text encoder '{tenc_fname}'.") + target = os.path.join(root, target[0]) + vsd, _ = gguf_sd_loader(target, is_text_model=True) + + # concat 4D to 5D + if "v.patch_embd.weight.1" in vsd: + w1 = dequantize_tensor(vsd.pop("v.patch_embd.weight"), dtype=torch.float32) + w2 = dequantize_tensor(vsd.pop("v.patch_embd.weight.1"), dtype=torch.float32) + vsd["v.patch_embd.weight"] = torch.stack([w1, w2], dim=2) + + if any("deepstack" in key or "deepstast" in key or "attn_qkv" in key for key in vsd): + # Qwen3-VL / Qwen3.5 style vision towers use fused ``v.blk.N.attn_qkv`` + # projections and a ``linear_fc`` merger. Qwen3.5 mmprojs omit the + # deepstack tensors, so detect them by the fused projection instead. + return sd_map_replace(vsd, CLIP_VISION_QWEN3_MAP) + + # run main replacement + vsd = sd_map_replace(vsd, CLIP_VISION_SD_MAP) + + # handle split Q/K/V + if "visual.blocks.0.attn_q.weight" in vsd: + attns = {} + # filter out attentions + group + for k,v in vsd.items(): + if any(x in k for x in ["attn_q", "attn_k", "attn_v"]): + k_attn, k_name = k.rsplit(".attn_", 1) + k_attn += ".attn.qkv." + k_name.split(".")[-1] + if k_attn not in attns: + attns[k_attn] = {} + attns[k_attn][k_name] = dequantize_tensor( + v, dtype=(torch.bfloat16 if is_quantized(v) else torch.float16) + ) + + # recombine + for k,v in attns.items(): + suffix = k.split(".")[-1] + vsd[k] = torch.cat([ + v[f"q.{suffix}"], + v[f"k.{suffix}"], + v[f"v.{suffix}"], + ], dim=0) + del attns + + return vsd + +def gguf_tokenizer_loader(path, temb_shape): + # convert gguf tokenizer to spiece + logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") + try: + from sentencepiece import sentencepiece_model_pb2 as model + except ImportError: + raise ImportError("Please make sure sentencepiece and protobuf are installed.\npip install sentencepiece protobuf") + spm = model.ModelProto() + + reader = gguf.GGUFReader(path) + + if get_field(reader, "tokenizer.ggml.model", str) == "t5": + if temb_shape == (256384, 4096): # probably UMT5 + spm.trainer_spec.model_type == 1 # Unigram (do we have a T5 w/ BPE?) + else: + raise NotImplementedError("Unknown model, can't set tokenizer!") + else: + raise NotImplementedError("Unknown model, can't set tokenizer!") + + spm.normalizer_spec.add_dummy_prefix = get_field(reader, "tokenizer.ggml.add_space_prefix", bool) + spm.normalizer_spec.remove_extra_whitespaces = get_field(reader, "tokenizer.ggml.remove_extra_whitespaces", bool) + + tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) + scores = get_list_field(reader, "tokenizer.ggml.scores", float) + toktypes = get_list_field(reader, "tokenizer.ggml.token_type", int) + + for idx, (token, score, toktype) in enumerate(zip(tokens, scores, toktypes)): + # # These aren't present in the original? + # if toktype == 5 and idx >= temb_shape[0]%1000): + # continue + + piece = spm.SentencePiece() + piece.piece = token + piece.score = score + piece.type = toktype + spm.pieces.append(piece) + + # unsure if any of these are correct + spm.trainer_spec.byte_fallback = True + spm.trainer_spec.vocab_size = len(tokens) # split off unused? + spm.trainer_spec.max_sentence_length = 4096 + spm.trainer_spec.eos_id = get_field(reader, "tokenizer.ggml.eos_token_id", int) + spm.trainer_spec.pad_id = get_field(reader, "tokenizer.ggml.padding_token_id", int) + + logging.info(f"Created tokenizer with vocab size of {len(spm.pieces)}") + del reader + return torch.ByteTensor(list(spm.SerializeToString())) + +def gguf_tekken_tokenizer_loader(path, temb_shape): + # convert ggml (hf) tokenizer metadata to tekken/comfy data + logging.info("Attempting to recreate tekken tokenizer from GGUF file metadata...") + import json + import base64 + from transformers.convert_slow_tokenizer import bytes_to_unicode + + reader = gguf.GGUFReader(path) + + model_str = get_field(reader, "tokenizer.ggml.model", str) + if model_str == "gpt2": + if temb_shape == (131072, 5120): # probably Mistral + data = { + "config": {"num_vocab_tokens": 150000, "default_vocab_size": 131072}, + "vocab": [], + "special_tokens": [], + } + else: + raise NotImplementedError("Unknown model, can't set tokenizer!") + else: + raise NotImplementedError("Unknown model, can't set tokenizer!") + + tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) + toktypes = get_list_field(reader, "tokenizer.ggml.token_type", int) + + decoder = {v: k for k, v in bytes_to_unicode().items()} + for idx, (token, toktype) in enumerate(zip(tokens, toktypes)): + if toktype == 3: + data["special_tokens"].append( + {'rank': idx, 'token_str': token, 'is_control': True} + ) + else: + tok = bytes([decoder[char] for char in token]) + data["vocab"].append({ + "rank": len(data["vocab"]), + "token_bytes": base64.b64encode(tok).decode("ascii"), + "token_str": tok.decode("utf-8", errors="replace") # ? + }) + + logging.info(f"Created tekken tokenizer with vocab size of {len(data['vocab'])} (+{len(data['special_tokens'])})") + del reader + return torch.ByteTensor(list(json.dumps(data).encode('utf-8'))) + +def gguf_gemma3_tokenizer_loader(path): + #TODO: merge into gguf_tokenizer_loader + logging.info("Attempting to recreate sentencepiece tokenizer from GGUF file metadata...") + try: + from sentencepiece import sentencepiece_model_pb2 as model + except ImportError: + raise ImportError("Please install sentencepiece and protobuf.\npip install sentencepiece protobuf") + spm = model.ModelProto() + reader = gguf.GGUFReader(path) + + spm.normalizer_spec.name = "identity" + spm.normalizer_spec.add_dummy_prefix = False + spm.trainer_spec.model_type = 2 + spm.trainer_spec.input_format = "tsv" + spm.trainer_spec.byte_fallback = True + spm.trainer_spec.max_sentence_length = 4192 + spm.trainer_spec.bos_piece = "" + + tokens = get_list_field(reader, "tokenizer.ggml.tokens", str) + scores = get_list_field(reader, "tokenizer.ggml.scores", float) + toktype = get_list_field(reader, "tokenizer.ggml.token_type", int) + + if not tokens or not scores or not toktype: + raise ValueError("Missing tokenizer metadata") + + for idx in range(len(tokens)): + piece = spm.SentencePiece() + piece.piece = tokens[idx] + if idx == 3: # UNK position + piece.type = 2 # UNK Token + piece.score = 0.0 # UNK Score + else: + piece.type = toktype[idx] + piece.score = scores[idx] + spm.pieces.append(piece) + + spm.trainer_spec.vocab_size = len(spm.pieces) + logging.info(f"Created tokenizer with vocab size of {len(spm.pieces)}") + + del reader + return torch.ByteTensor(list(spm.SerializeToString())) + +def gemma4_tokenizer_json(tokens, merges, token_types): + if not tokens or not merges or not token_types: + raise ValueError("Missing Gemma 4 tokenizer metadata") + if len(tokens) != len(token_types): + raise ValueError("Gemma 4 tokenizer token and token-type counts differ") + + # Gemma 4 stores its BPE vocabulary in standard GGUF token and merge fields. + data = { + "version": "1.0", + "truncation": None, + "padding": None, + "added_tokens": [ + { + "id": index, + "content": token, + "single_word": False, + "lstrip": False, + "rstrip": False, + "normalized": False, + "special": True, + } + for index, (token, token_type) in enumerate(zip(tokens, token_types)) + if token_type == 3 + ], + "normalizer": None, + "pre_tokenizer": { + "type": "Metaspace", + "replacement": "\u2581", + "prepend_scheme": "never", + "split": True, + }, + "post_processor": None, + "decoder": { + "type": "Metaspace", + "replacement": "\u2581", + "prepend_scheme": "never", + "split": True, + }, + "model": { + "type": "BPE", + "dropout": None, + "unk_token": "", + "continuing_subword_prefix": None, + "end_of_word_suffix": None, + "fuse_unk": False, + "byte_fallback": False, + "ignore_merges": False, + "vocab": {token: index for index, token in enumerate(tokens)}, + "merges": list(merges), + }, + } + return torch.ByteTensor(list(json.dumps(data, ensure_ascii=False).encode("utf-8"))) + +def gguf_gemma4_tokenizer_loader(path): + logging.info("Recreating Gemma 4 BPE tokenizer from GGUF metadata...") + reader = gguf.GGUFReader(path) + if get_field(reader, "tokenizer.ggml.model", str) != "gemma4": + raise ValueError("Expected a Gemma 4 tokenizer in the GGUF metadata") + + tokenizer_json = gemma4_tokenizer_json( + get_list_field(reader, "tokenizer.ggml.tokens", str), + get_list_field(reader, "tokenizer.ggml.merges", str), + get_list_field(reader, "tokenizer.ggml.token_type", int), + ) + del reader + return tokenizer_json + +def inject_qwen3vl_detection_markers(sd): + """Add visual sentinels when a llama.cpp Qwen3-VL GGUF excludes its vision tower.""" + ln_key = "model.layers.0.input_layernorm.weight" + lm_hidden = int(sd[ln_key].shape[0]) if ln_key in sd else 2560 + vis_hidden = 1024 if lm_hidden == 2560 else 1152 + merge_dim = vis_hidden * 4 # spatial_merge_size=2 + + if lm_hidden == 5120: + # MiniMax H3 uses the truncated Qwen3-VL-32B encoder. Its detector + # deliberately checks this unprefixed visual key plus layer 49. + marker_key = "visual.deepstack_merger_list.0.norm.weight" + else: + marker_key = "model.visual.deepstack_merger_list.0.norm.weight" + + sd[marker_key] = torch.zeros(merge_dim) + if lm_hidden != 5120: + sd["model.visual.merger.linear_fc2.weight"] = torch.zeros(lm_hidden, merge_dim) + logging.info( + "qwen3vl GGUF: injected visual marker tensor " + "(lm_hidden=%d, merge_dim=%d)", + lm_hidden, + merge_dim, + ) + +def gguf_clip_loader(path, dynamic=False, progress_callback=None): + sd, extra = gguf_sd_loader( + path, + is_text_model=True, + dynamic=dynamic, + progress_callback=progress_callback, + ) + arch = extra.get("arch_str", None) + if arch == "minimax_music3" and "tokenizer_json" in sd: + sd["tokenizer_json"] = normalize_raw_byte_tensor(sd["tokenizer_json"]) + if arch in {"t5", "t5encoder"}: + temb_key = "token_embd.weight" + if temb_key in sd and sd[temb_key].shape == (256384, 4096): + # non-standard Comfy-Org tokenizer + sd["spiece_model"] = gguf_tokenizer_loader(path, sd[temb_key].shape) + # TODO: dequantizing token embed here is janky but otherwise we OOM due to tensor being massive. + logging.warning(f"Dequantizing {temb_key} to prevent runtime OOM.") + sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16) + sd = sd_map_replace(sd, T5_SD_MAP) + elif arch in {"llama", "qwen2vl", "qwen3", "qwen3vl", "qwen35", "gemma3", "gemma4"}: + # TODO: pass model_options["vocab_size"] to loader somehow + temb_key = "token_embd.weight" + if temb_key in sd and sd[temb_key].shape[0] >= (64 * 1024): + if arch == "llama" and sd[temb_key].shape == (131072, 5120): + # non-standard Comfy-Org tokenizer + sd["tekken_model"] = gguf_tekken_tokenizer_loader(path, sd[temb_key].shape) + elif arch == "gemma3": + sd["spiece_model"] = gguf_gemma3_tokenizer_loader(path) + elif arch == "gemma4": + sd["tokenizer_json"] = gguf_gemma4_tokenizer_loader(path) + # See note above for T5. + logging.warning(f"Dequantizing {temb_key} to prevent runtime OOM.") + sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16) + if arch == "gemma3": + sd = sd_map_replace(sd, GEMMA3_SD_MAP) + sd = gemma3_norm_corrections(sd) + elif arch == "gemma4": + # ComfyUI calculates Gemma 4 RoPE frequencies itself. + sd.pop("rope_freqs.weight", None) + sd = sd_map_replace(sd, GEMMA4_SD_MAP) + elif arch == "qwen35": + sd = sd_map_replace(sd, QWEN35_SD_MAP) + sd = qwen35_corrections(sd) + vsd = gguf_mmproj_loader(path) + sd.update(vsd) + else: + sd = sd_map_replace(sd, LLAMA_SD_MAP) + if arch == "llama": + sd = llama_permute(sd, 32, 8) # L3 / Mistral + if arch == "qwen2vl": + vsd = gguf_mmproj_loader(path) + sd.update(vsd) + if arch == "qwen3vl": + vsd = gguf_mmproj_loader(path) + if vsd and "model.layers.49.self_attn.q_proj.weight" in sd: + # MiniMax H3 receives the Qwen3-VL vision tower under + # ``visual.*``, unlike the standalone Qwen3-VL variants. + vsd = { + key.replace("model.visual.", "visual.", 1) + if key.startswith("model.visual.") + else key: value + for key, value in vsd.items() + } + sd.update(vsd) + if arch == "qwen3vl" and not ( + "model.visual.deepstack_merger_list.0.norm.weight" in sd + or "visual.deepstack_merger_list.0.norm.weight" in sd + ): + # Standard llama.cpp Qwen3-VL GGUFs omit the visual tower. Without it, + # detect_te_model() mis-classifies the state dict as a Qwen3 LM instead + # of Qwen3-VL. MiniMax H3 additionally uses Qwen3-VL-32B truncated to + # 50 layers, whose detector relies on an unprefixed visual marker. + # Inject zero sentinel tensors with shapes that exactly match the model + # parameters so that load_state_dict(strict=False) doesn't raise a size + # mismatch error while still satisfying detect_te_model()'s key checks. + inject_qwen3vl_detection_markers(sd) + if "lm_head.weight" in sd and is_quantized(sd["lm_head.weight"]): + lm_head = sd["lm_head.weight"] + if getattr(lm_head, "tensor_type", None) != gguf.GGMLQuantizationType.BF16: + # BaseGenerate.logits() feeds lm_head straight to F.linear, + # bypassing the GGML ops path that dequantizes on the fly. + # Block-quantized GGUF data has a byte-expanded shape (e.g. + # Q8_0 stores scales interleaved), so raw data must never + # reach the matmul. BF16 storage is already full precision + # and keeps its logical shape, so it can stay quantized. + logging.warning(f"Dequantizing lm_head.weight to prevent raw-block matmul.") + sd["lm_head.weight"] = dequantize_tensor(lm_head, dtype=torch.float16) + elif arch == "ideogram": + # Dequantize Ideogram model for inference + logging.info("Dequantizing Ideogram model for inference...") + # Use BF16 to save VRAM while maintaining quality, but fall back to FP16 + # on devices that don't support bf16 (avoids slow fp32 compute fallback). + target_dtype = torch.bfloat16 if device_supports_bf16() else torch.float16 + dequantized_count = 0 + for key in list(sd.keys()): + if is_quantized(sd[key]): + sd[key] = dequantize_tensor(sd[key], dtype=target_dtype) + dequantized_count += 1 + logging.info(f"Dequantized {dequantized_count} tensors for Ideogram model ({target_dtype})") + else: + pass + return sd diff --git a/lora.py b/lora.py new file mode 100644 index 00000000..638e2d91 --- /dev/null +++ b/lora.py @@ -0,0 +1,487 @@ +import json +import logging +import os + +import gguf +import torch +from safetensors import safe_open + + +_SUPPORTED_FACTOR_TYPES = { + gguf.GGMLQuantizationType.F32, + gguf.GGMLQuantizationType.F16, + gguf.GGMLQuantizationType.BF16, + gguf.GGMLQuantizationType.Q8_0, +} +_FP8_SOURCE_TYPES = { + dtype + for dtype in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e5m2", None), + ) + if dtype is not None +} + + +def _read_int8_quant_config(state_dict, weight_key): + quant_key = f"{weight_key[:-len('weight')]}comfy_quant" + quant_data = state_dict.get(quant_key) + if quant_data is None: + return None, quant_key + if not isinstance(quant_data, torch.Tensor) or quant_data.dtype != torch.uint8: + raise ValueError( + f"INT8 source quantization metadata {quant_key!r} must be a uint8 JSON tensor." + ) + try: + return json.loads(bytes(quant_data.detach().cpu().tolist()).decode("utf-8")), quant_key + except (UnicodeDecodeError, json.JSONDecodeError, ValueError) as error: + raise ValueError( + f"INT8 source quantization metadata {quant_key!r} is invalid JSON." + ) from error + + +def _unrotate_convrot_weight(weight, group_size): + from comfy_kitchen.tensor.int8_utils import _build_hadamard + + if group_size <= 0 or weight.shape[1] % group_size: + raise ValueError( + f"ConvRot group size {group_size} does not divide weight shape {tuple(weight.shape)}." + ) + hadamard = _build_hadamard(group_size, device=weight.device, dtype=weight.dtype) + return torch.matmul( + weight.reshape(weight.shape[0], -1, group_size), + hadamard, + ).reshape_as(weight) + + +def materialize_int8_source_weights(state_dict): + """Restore scaled INT8 weights, reversing ConvRot before LoRA fusion or export.""" + int8_keys = [ + key for key, value in state_dict.items() + if isinstance(value, torch.Tensor) and value.dtype == torch.int8 + ] + for weight_key in int8_keys: + if not weight_key.endswith(".weight"): + raise ValueError( + f"INT8 source tensor {weight_key!r} is not a Linear weight. " + "Only scaled INT8 Linear checkpoint weights are supported." + ) + scale_key = f"{weight_key}_scale" + scale = state_dict.get(scale_key) + if scale is None: + raise ValueError( + f"INT8 source weight {weight_key!r} is missing its scale tensor " + f"{scale_key!r}." + ) + if not isinstance(scale, torch.Tensor) or not scale.is_floating_point(): + raise ValueError( + f"INT8 source weight scale {scale_key!r} must be a floating-point tensor." + ) + + weight = state_dict[weight_key] + scale = scale.to(dtype=torch.float32) + if scale.ndim == 1 and scale.shape[0] == weight.shape[0]: + scale = scale.unsqueeze(1) + try: + restored = weight.to(dtype=torch.float32) * scale + except RuntimeError as error: + raise ValueError( + f"INT8 source weight scale {scale_key!r} with shape {tuple(scale.shape)} " + f"cannot be broadcast over {weight_key!r} with shape {tuple(weight.shape)}." + ) from error + + quant_conf, quant_key = _read_int8_quant_config(state_dict, weight_key) + if quant_conf is not None: + if quant_conf.get("format") != "int8_tensorwise": + raise ValueError( + f"INT8 source quantization metadata {quant_key!r} uses unsupported " + f"format {quant_conf.get('format')!r}." + ) + if quant_conf.get("convrot"): + restored = _unrotate_convrot_weight( + restored, + int(quant_conf.get("convrot_groupsize", 256)), + ) + del state_dict[quant_key] + + state_dict[weight_key] = restored.to(dtype=torch.float16) + del state_dict[scale_key] + return len(int8_keys) + + +def _read_string_field(reader, name, required=True): + field = reader.get_field(name) + if field is None: + if required: + raise ValueError(f"GGUF LoRA is missing required metadata {name!r}.") + return None + if len(field.types) != 1 or field.types[0] != gguf.GGUFValueType.STRING: + raise ValueError(f"GGUF LoRA metadata {name!r} must be a string.") + return str(field.parts[field.data[-1]], encoding="utf-8") + + +def _read_float_field(reader, name): + field = reader.get_field(name) + if field is None: + return None + if len(field.types) != 1 or field.types[0] != gguf.GGUFValueType.FLOAT32: + raise ValueError(f"GGUF LoRA metadata {name!r} must be a float32.") + return float(field.parts[field.data[-1]].item()) + + +def _tensor_to_float(tensor): + if tensor.tensor_type not in _SUPPORTED_FACTOR_TYPES: + raise ValueError( + f"GGUF LoRA tensor {tensor.name!r} uses unsupported {tensor.tensor_type.name}. " + "Only F32, F16, BF16, and Q8_0 factors are supported." + ) + shape = tuple(int(value) for value in reversed(tensor.shape)) + data = torch.from_numpy(tensor.data.copy()) + if tensor.tensor_type == gguf.GGMLQuantizationType.F32: + return data.view(torch.float32).reshape(shape) + if tensor.tensor_type == gguf.GGMLQuantizationType.F16: + return data.view(torch.float16).reshape(shape) + if tensor.tensor_type == gguf.GGMLQuantizationType.BF16: + return data.view(torch.bfloat16).reshape(shape) + return torch.from_numpy( + gguf.quants.dequantize(tensor.data, tensor.tensor_type) + ).reshape(shape).to(torch.float16) + + +def _build_targets(pairs, path, architecture=None, default_alpha=None): + if not pairs: + raise ValueError("LoRA contains no factor tensors.") + + lora = {} + targets = {} + for base_name, pair in pairs.items(): + if not {"a", "b"}.issubset(pair): + raise ValueError(f"LoRA target {base_name!r} is missing one factor.") + down, up = pair["a"], pair["b"] + if down.ndim != 2 or up.ndim != 2: + raise ValueError( + f"LoRA target {base_name!r} is not 2-D; convolutional and other " + "adapter layouts are not supported." + ) + if down.shape[0] != up.shape[1]: + raise ValueError( + f"LoRA target {base_name!r} has incompatible factor shapes " + f"{tuple(down.shape)} and {tuple(up.shape)}." + ) + target_name = base_name.removesuffix(".weight") + source_name = base_name if base_name.endswith(".weight") else f"{base_name}.weight" + alpha = pair.get("alpha", default_alpha) + lora[f"{target_name}.lora_A.weight"] = down + lora[f"{target_name}.lora_B.weight"] = up + if alpha is not None: + lora[f"{target_name}.alpha"] = torch.tensor(alpha, dtype=torch.float32) + targets[target_name] = { + "base_name": source_name, + "down": down, + "up": up, + "alpha": alpha, + } + + metadata = { + "path": os.path.abspath(path), + "architecture": architecture, + "alpha": default_alpha, + "target_count": len(targets), + } + return lora, targets, metadata + + +def load_gguf_lora(path): + """Read a standard GGUF LoRA into ComfyUI's lora_A/lora_B dictionary form.""" + reader = gguf.GGUFReader(path) + try: + if _read_string_field(reader, "general.type") != "adapter": + raise ValueError("GGUF file is not an adapter (general.type must be 'adapter').") + if _read_string_field(reader, "adapter.type") != "lora": + raise ValueError("GGUF adapter is not a LoRA (adapter.type must be 'lora').") + + pairs = {} + for tensor in reader.tensors: + if tensor.name.endswith(".lora_a"): + base_name = tensor.name[:-len(".lora_a")] + pairs.setdefault(base_name, {})["a"] = _tensor_to_float(tensor) + elif tensor.name.endswith(".lora_b"): + base_name = tensor.name[:-len(".lora_b")] + pairs.setdefault(base_name, {})["b"] = _tensor_to_float(tensor) + else: + raise ValueError( + f"Unsupported GGUF LoRA tensor {tensor.name!r}; only paired " + "'.lora_a' and '.lora_b' factors are supported." + ) + + return _build_targets( + pairs, + path, + architecture=_read_string_field(reader, "general.architecture", required=False), + default_alpha=_read_float_field(reader, "adapter.lora.alpha"), + ) + finally: + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + + +_SAFETENSORS_FACTOR_SUFFIXES = ( + (".lora_A.weight", "a"), + (".lora_B.weight", "b"), + (".lora_down.weight", "a"), + (".lora_up.weight", "b"), +) +_SAFETENSORS_LOKR_SUFFIXES = ( + (".lokr_w1", "w1"), + (".lokr_w2", "w2"), +) + + +def load_safetensors_lora(path): + """Load Linear LoRA and direct-factor LyCORIS LoKr safetensors adapters.""" + pairs = {} + lokr_pairs = {} + with safe_open(path, framework="pt", device="cpu") as checkpoint: + keys = set(checkpoint.keys()) + for key in keys: + match = next( + ((suffix, factor) for suffix, factor in _SAFETENSORS_FACTOR_SUFFIXES if key.endswith(suffix)), + None, + ) + if match is None: + match = next( + ( + (suffix, factor) + for suffix, factor in _SAFETENSORS_LOKR_SUFFIXES + if key.endswith(suffix) + ), + None, + ) + if match is None: + continue + suffix, factor = match + base_name = key[:-len(suffix)] + lokr_pairs.setdefault(base_name, {})[factor] = checkpoint.get_tensor(key) + else: + suffix, factor = match + base_name = key[:-len(suffix)] + pairs.setdefault(base_name, {})[factor] = checkpoint.get_tensor(key) + + for base_name, pair in pairs.items(): + for alpha_key in (f"{base_name}.alpha", f"{base_name}.lora_alpha"): + if alpha_key not in keys: + continue + alpha = checkpoint.get_tensor(alpha_key) + if alpha.numel() != 1: + raise ValueError(f"LoRA alpha {alpha_key!r} must be a scalar.") + pair["alpha"] = float(alpha.item()) + break + + for base_name, pair in lokr_pairs.items(): + for alpha_key in (f"{base_name}.alpha", f"{base_name}.lora_alpha"): + if alpha_key not in keys: + continue + alpha = checkpoint.get_tensor(alpha_key) + if alpha.numel() != 1: + raise ValueError(f"LoKr alpha {alpha_key!r} must be a scalar.") + pair["alpha"] = float(alpha.item()) + break + + if pairs: + lora, targets, metadata = _build_targets(pairs, path) + else: + lora = {} + targets = {} + metadata = { + "path": os.path.abspath(path), + "architecture": None, + "alpha": None, + "target_count": 0, + } + + for base_name, pair in lokr_pairs.items(): + if not {"w1", "w2"}.issubset(pair): + raise ValueError( + f"LoKr target {base_name!r} is missing a direct lokr_w1 or lokr_w2 factor. " + "Factorized and Tucker LoKr layouts are not supported for offline fusion." + ) + w1, w2 = pair["w1"], pair["w2"] + if w1.ndim != 2 or w2.ndim != 2: + raise ValueError( + f"LoKr target {base_name!r} is not a 2-D Linear adapter." + ) + target_name = base_name.removesuffix(".weight") + source_name = base_name if base_name.endswith(".weight") else f"{base_name}.weight" + lora[f"{target_name}.lokr_w1"] = w1 + lora[f"{target_name}.lokr_w2"] = w2 + targets[target_name] = { + "base_name": source_name, + "kind": "lokr", + "w1": w1, + "w2": w2, + "alpha": pair.get("alpha"), + } + + metadata["target_count"] = len(targets) + if not targets: + raise ValueError("LoRA contains no supported factor tensors.") + return lora, targets, metadata + + +def load_lora(path): + """Load a GGUF or safetensors LoRA suitable for offline fusion.""" + extension = os.path.splitext(path)[1].lower() + if extension == ".gguf": + return load_gguf_lora(path) + if extension == ".safetensors": + return load_safetensors_lora(path) + raise ValueError(f"LoRA fusion accepts .gguf or .safetensors adapters, got {path!r}.") + + +def _find_state_key(state_dict, candidates): + return next((candidate for candidate in candidates if candidate in state_dict), None) + + +def _comfy_model_lora_target_map(state_dict): + """Build ComfyUI's current diffusion-model LoRA map without loading weights.""" + import comfy.lora + import comfy.model_detection + + model_config = comfy.model_detection.model_config_from_unet(state_dict, "") + if model_config is None: + return {} + + model = model_config.get_model({}, device=torch.device("meta")) + source_state = { + f"diffusion_model.{key}": value + for key, value in state_dict.items() + } + object.__setattr__( + model, + "state_dict", + lambda *args, **kwargs: source_state, + ) + return comfy.lora.model_lora_keys_unet(model, {}) + + +def _target_shape(target): + if target.get("kind") == "lokr": + return ( + target["w1"].shape[0] * target["w2"].shape[0], + target["w1"].shape[1] * target["w2"].shape[1], + ) + return (target["up"].shape[0], target["down"].shape[1]) + + +def resolve_fusion_targets(state_dict, targets): + resolved = {} + missing = [] + comfy_map = None + for target_name, target in targets.items(): + candidates = ( + target["base_name"], + f"{target_name}.weight", + target_name, + target["base_name"].removeprefix("diffusion_model."), + f"{target_name.removeprefix('diffusion_model.')}.weight", + ) + state_key = _find_state_key(state_dict, candidates) + target_slice = None + if state_key is None: + if comfy_map is None: + comfy_map = _comfy_model_lora_target_map(state_dict) + mapped = comfy_map.get(target_name) + if isinstance(mapped, tuple): + mapped, target_slice = mapped[:2] + if mapped is not None: + bare_mapped = mapped.removeprefix("model.diffusion_model.").removeprefix( + "diffusion_model." + ) + state_key = _find_state_key( + state_dict, + ( + mapped, + bare_mapped, + f"diffusion_model.{bare_mapped}", + f"model.diffusion_model.{bare_mapped}", + ), + ) + if state_key is None: + missing.append(target["base_name"]) + continue + weight = state_dict[state_key] + target_weight = weight + if target_slice is not None: + target_weight = weight.narrow(*target_slice) + if target_weight.ndim != 2: + raise ValueError( + f"GGUF LoRA target {target['base_name']!r} resolves to {state_key!r}, " + f"which is {target_weight.ndim}-D. Only 2-D Linear weights can be fused." + ) + if tuple(target_weight.shape) != _target_shape(target): + raise ValueError( + f"GGUF LoRA target {target['base_name']!r} does not match {state_key!r}: " + f"base shape {tuple(target_weight.shape)}, adapter delta shape " + f"{_target_shape(target)}." + ) + resolved.setdefault(state_key, []).append((target, target_slice)) + if missing: + logging.warning( + "Skipping %d LoRA target(s) not present in the source checkpoint: %s", + len(missing), + ", ".join(sorted(missing)), + ) + return resolved + + +def fuse_targets_into_state_dict(state_dict, targets, strength, device): + """Apply one adapter's linear deltas in GPU/CPU FP32, returning the target count.""" + resolved = resolve_fusion_targets(state_dict, targets) + if not resolved: + logging.warning( + "No compatible LoRA targets were found in the source checkpoint; " + "this adapter does not change the fused cache." + ) + fused_count = 0 + for state_key, target_entries in resolved.items(): + source = state_dict[state_key] + state_dict[state_key], count = fuse_target_entries_into_tensor( + source, + target_entries, + strength, + device, + state_dict.get(f"{state_key}_scale"), + ) + fused_count += count + return fused_count + + +def fuse_target_entries_into_tensor(source, target_entries, strength, device, source_scale=None): + """Fuse resolved LoRA targets into one tensor without retaining unrelated weights.""" + source_dtype = source.dtype + if source_dtype not in {torch.float16, torch.bfloat16, torch.float32, *_FP8_SOURCE_TYPES}: + raise ValueError( + f"Fusion target has unsupported dtype {source.dtype}. " + "Fuse from an FP16, BF16, FP32, or scaled FP8 checkpoint." + ) + fused = source.to(device=device, dtype=torch.float32) + if source_dtype in _FP8_SOURCE_TYPES and source_scale is not None: + if source_scale.numel() != 1: + raise ValueError("FP8 scale tensor must be a scalar.") + fused.mul_(source_scale.to(device=device, dtype=torch.float32)) + for target, target_slice in target_entries: + target_fused = fused if target_slice is None else fused.narrow(*target_slice) + if target.get("kind") == "lokr": + w1 = target["w1"].to(device=device, dtype=torch.float32) + w2 = target["w2"].to(device=device, dtype=torch.float32) + target_fused.add_(torch.kron(w1, w2), alpha=strength) + del w1, w2 + else: + down = target["down"].to(device=device, dtype=torch.float32) + up = target["up"].to(device=device, dtype=torch.float32) + alpha = target["alpha"] if target["alpha"] is not None else down.shape[0] + target_fused.add_(up.matmul(down), alpha=strength * alpha / down.shape[0]) + del down, up + output_dtype = torch.float16 if source_dtype in _FP8_SOURCE_TYPES else source_dtype + return fused.to(device="cpu", dtype=output_dtype), len(target_entries) diff --git a/nodes.py b/nodes.py index 46835146..426efb73 100644 --- a/nodes.py +++ b/nodes.py @@ -1,329 +1,1272 @@ -# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) -import torch -import logging -import inspect -import collections - -import nodes -import comfy.sd -import comfy.lora -import comfy.float -import comfy.utils -import comfy.model_patcher -import comfy.model_management -import folder_paths - -from .ops import GGMLOps, move_patch_to_device -from .loader import gguf_sd_loader, gguf_clip_loader -from .dequant import is_quantized, is_torch_compatible - -def update_folder_names_and_paths(key, targets=[]): - # check for existing key - base = folder_paths.folder_names_and_paths.get(key, ([], {})) - base = base[0] if isinstance(base[0], (list, set, tuple)) else [] - # find base key & add w/ fallback, sanity check + warning - target = next((x for x in targets if x in folder_paths.folder_names_and_paths), targets[0]) - orig, _ = folder_paths.folder_names_and_paths.get(target, ([], {})) - folder_paths.folder_names_and_paths[key] = (orig or base, {".gguf"}) - if base and base != orig: - logging.warning(f"Unknown file list already present on key {key}: {base}") - -# Add a custom keys for files ending in .gguf -update_folder_names_and_paths("unet_gguf", ["diffusion_models", "unet"]) -update_folder_names_and_paths("clip_gguf", ["text_encoders", "clip"]) - -class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): - patch_on_device = False - - def patch_weight_to_device(self, key, device_to=None, inplace_update=False): - if key not in self.patches: - return - weight = comfy.utils.get_attr(self.model, key) - - patches = self.patches[key] - if is_quantized(weight): - out_weight = weight.to(device_to) - patches = move_patch_to_device(patches, self.load_device if self.patch_on_device else self.offload_device) - # TODO: do we ever have legitimate duplicate patches? (i.e. patch on top of patched weight) - out_weight.patches = [(patches, key)] - else: - inplace_update = self.weight_inplace_update or inplace_update - if key not in self.backup: - self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])( - weight.to(device=self.offload_device, copy=inplace_update), inplace_update - ) - - if device_to is not None: - temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True) - else: - temp_weight = weight.to(torch.float32, copy=True) - - out_weight = comfy.lora.calculate_weight(patches, temp_weight, key) - out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype) - - if inplace_update: - comfy.utils.copy_to_param(self.model, key, out_weight) - else: - comfy.utils.set_attr_param(self.model, key, out_weight) - - def unpatch_model(self, device_to=None, unpatch_weights=True): - if unpatch_weights: - for p in self.model.parameters(): - if is_torch_compatible(p): - continue - patches = getattr(p, "patches", []) - if len(patches) > 0: - p.patches = [] - # TODO: Find another way to not unload after patches - return super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights) - - - def pin_weight_to_device(self, key): - op_key = key.rsplit('.', 1)[0] - if not self.mmap_released and op_key in self.named_modules_to_munmap: - # TODO: possible to OOM, find better way to detach - self.named_modules_to_munmap[op_key].to(self.load_device).to(self.offload_device) - del self.named_modules_to_munmap[op_key] - super().pin_weight_to_device(key) - - mmap_released = False - named_modules_to_munmap = {} - - def load(self, *args, force_patch_weights=False, **kwargs): - if not self.mmap_released: - self.named_modules_to_munmap = dict(self.model.named_modules()) - - # always call `patch_weight_to_device` even for lowvram - super().load(*args, force_patch_weights=True, **kwargs) - - # make sure nothing stays linked to mmap after first load - if not self.mmap_released: - linked = [] - if kwargs.get("lowvram_model_memory", 0) > 0: - for n, m in self.named_modules_to_munmap.items(): - if hasattr(m, "weight"): - device = getattr(m.weight, "device", None) - if device == self.offload_device: - linked.append((n, m)) - continue - if hasattr(m, "bias"): - device = getattr(m.bias, "device", None) - if device == self.offload_device: - linked.append((n, m)) - continue - if linked and self.load_device != self.offload_device: - logging.info(f"Attempting to release mmap ({len(linked)})") - for n, m in linked: - # TODO: possible to OOM, find better way to detach - m.to(self.load_device).to(self.offload_device) - self.mmap_released = True - self.named_modules_to_munmap = {} - - def clone(self, *args, **kwargs): - src_cls = self.__class__ - self.__class__ = GGUFModelPatcher - n = super().clone(*args, **kwargs) - n.__class__ = GGUFModelPatcher - self.__class__ = src_cls - # GGUF specific clone values below - n.patch_on_device = getattr(self, "patch_on_device", False) - n.mmap_released = getattr(self, "mmap_released", False) - if src_cls != GGUFModelPatcher: - n.size = 0 # force recalc - return n - -class UnetLoaderGGUF: - @classmethod - def INPUT_TYPES(s): - unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")] - return { - "required": { - "unet_name": (unet_names,), - } - } - - RETURN_TYPES = ("MODEL",) - FUNCTION = "load_unet" - CATEGORY = "bootleg" - TITLE = "Unet Loader (GGUF)" - - def load_unet(self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None): - ops = GGMLOps() - - if dequant_dtype in ("default", None): - ops.Linear.dequant_dtype = None - elif dequant_dtype in ["target"]: - ops.Linear.dequant_dtype = dequant_dtype - else: - ops.Linear.dequant_dtype = getattr(torch, dequant_dtype) - - if patch_dtype in ("default", None): - ops.Linear.patch_dtype = None - elif patch_dtype in ["target"]: - ops.Linear.patch_dtype = patch_dtype - else: - ops.Linear.patch_dtype = getattr(torch, patch_dtype) - - # init model - unet_path = folder_paths.get_full_path("unet", unet_name) - sd, extra = gguf_sd_loader(unet_path) - - kwargs = {} - valid_params = inspect.signature(comfy.sd.load_diffusion_model_state_dict).parameters - if "metadata" in valid_params: - kwargs["metadata"] = extra.get("metadata", {}) - - model = comfy.sd.load_diffusion_model_state_dict( - sd, model_options={"custom_operations": ops}, **kwargs, - ) - if model is None: - logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path)) - raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path)) - model = GGUFModelPatcher.clone(model) - model.patch_on_device = patch_on_device - return (model,) - -class UnetLoaderGGUFAdvanced(UnetLoaderGGUF): - @classmethod - def INPUT_TYPES(s): - unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")] - return { - "required": { - "unet_name": (unet_names,), - "dequant_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}), - "patch_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}), - "patch_on_device": ("BOOLEAN", {"default": False}), - } - } - TITLE = "Unet Loader (GGUF/Advanced)" - -class CLIPLoaderGGUF: - @classmethod - def INPUT_TYPES(s): - base = nodes.CLIPLoader.INPUT_TYPES() - return { - "required": { - "clip_name": (s.get_filename_list(),), - "type": base["required"]["type"], - } - } - - RETURN_TYPES = ("CLIP",) - FUNCTION = "load_clip" - CATEGORY = "bootleg" - TITLE = "CLIPLoader (GGUF)" - - @classmethod - def get_filename_list(s): - files = [] - files += folder_paths.get_filename_list("clip") - files += folder_paths.get_filename_list("clip_gguf") - return sorted(files) - - def load_data(self, ckpt_paths): - clip_data = [] - for p in ckpt_paths: - if p.endswith(".gguf"): - sd = gguf_clip_loader(p) - else: - sd = comfy.utils.load_torch_file(p, safe_load=True) - if "scaled_fp8" in sd: # NOTE: Scaled FP8 would require different custom ops, but only one can be active - raise NotImplementedError(f"Mixing scaled FP8 with GGUF is not supported! Use regular CLIP loader or switch model(s)\n({p})") - clip_data.append(sd) - return clip_data - - def load_patcher(self, clip_paths, clip_type, clip_data): - clip = comfy.sd.load_text_encoder_state_dicts( - clip_type = clip_type, - state_dicts = clip_data, - model_options = { - "custom_operations": GGMLOps, - "initial_device": comfy.model_management.text_encoder_offload_device() - }, - embedding_directory = folder_paths.get_folder_paths("embeddings"), - ) - clip.patcher = GGUFModelPatcher.clone(clip.patcher) - return clip - - def load_clip(self, clip_name, type="stable_diffusion"): - clip_path = folder_paths.get_full_path("clip", clip_name) - clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher([clip_path], clip_type, self.load_data([clip_path])),) - -class DualCLIPLoaderGGUF(CLIPLoaderGGUF): - @classmethod - def INPUT_TYPES(s): - base = nodes.DualCLIPLoader.INPUT_TYPES() - file_options = (s.get_filename_list(), ) - return { - "required": { - "clip_name1": file_options, - "clip_name2": file_options, - "type": base["required"]["type"], - } - } - - TITLE = "DualCLIPLoader (GGUF)" - - def load_clip(self, clip_name1, clip_name2, type): - clip_path1 = folder_paths.get_full_path("clip", clip_name1) - clip_path2 = folder_paths.get_full_path("clip", clip_name2) - clip_paths = (clip_path1, clip_path2) - clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) - -class TripleCLIPLoaderGGUF(CLIPLoaderGGUF): - @classmethod - def INPUT_TYPES(s): - file_options = (s.get_filename_list(), ) - return { - "required": { - "clip_name1": file_options, - "clip_name2": file_options, - "clip_name3": file_options, - } - } - - TITLE = "TripleCLIPLoader (GGUF)" - - def load_clip(self, clip_name1, clip_name2, clip_name3, type="sd3"): - clip_path1 = folder_paths.get_full_path("clip", clip_name1) - clip_path2 = folder_paths.get_full_path("clip", clip_name2) - clip_path3 = folder_paths.get_full_path("clip", clip_name3) - clip_paths = (clip_path1, clip_path2, clip_path3) - clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) - -class QuadrupleCLIPLoaderGGUF(CLIPLoaderGGUF): - @classmethod - def INPUT_TYPES(s): - file_options = (s.get_filename_list(), ) - return { - "required": { - "clip_name1": file_options, - "clip_name2": file_options, - "clip_name3": file_options, - "clip_name4": file_options, - } - } - - TITLE = "QuadrupleCLIPLoader (GGUF)" - - def load_clip(self, clip_name1, clip_name2, clip_name3, clip_name4, type="stable_diffusion"): - clip_path1 = folder_paths.get_full_path("clip", clip_name1) - clip_path2 = folder_paths.get_full_path("clip", clip_name2) - clip_path3 = folder_paths.get_full_path("clip", clip_name3) - clip_path4 = folder_paths.get_full_path("clip", clip_name4) - clip_paths = (clip_path1, clip_path2, clip_path3, clip_path4) - clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) - return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) - -NODE_CLASS_MAPPINGS = { - "UnetLoaderGGUF": UnetLoaderGGUF, - "CLIPLoaderGGUF": CLIPLoaderGGUF, - "DualCLIPLoaderGGUF": DualCLIPLoaderGGUF, - "TripleCLIPLoaderGGUF": TripleCLIPLoaderGGUF, - "QuadrupleCLIPLoaderGGUF": QuadrupleCLIPLoaderGGUF, - "UnetLoaderGGUFAdvanced": UnetLoaderGGUFAdvanced, -} - +# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import torch +import logging +import inspect +import collections +import json +import os +from contextlib import nullcontext + +import nodes +import comfy.sd +import comfy.lora +import comfy.float +import comfy.ops +import comfy.utils +import comfy.model_patcher +import comfy.model_management +import comfy.memory_management +import folder_paths + +from .ops import ( + GGMLTensor, + GGMLOps, + get_gguf_q8_ops, + get_gguf_q4_w4a4_ops, + int4_lora_offload_enabled, + log_cuda_oom_loaded_models, + move_patch_to_device, +) +from .loader import gguf_sd_loader, gguf_clip_loader, gguf_tensor_count +from .dequant import dequantize_tensor, is_quantized, is_torch_compatible +from .tools.convert import ( + DEFAULT_TARGET_SIZE_Q8_TYPE, + QUANT_TYPE_MAP, + QUANTIZATION_DEVICE_OPTIONS, + TARGET_SIZE_Q8_TYPES, + TARGET_SIZE_QUANT_TYPE, + convert_file, +) +from .lora import load_gguf_lora + + +_DYNAMIC_VRAM_LORA_WARNING_MIN_BYTES = 64 * 1024 * 1024 + + +def update_folder_names_and_paths(key, targets=[]): + # check for existing key + base = folder_paths.folder_names_and_paths.get(key, ([], {})) + base = base[0] if isinstance(base[0], (list, set, tuple)) else [] + # find base key & add w/ fallback, sanity check + warning + target = next((x for x in targets if x in folder_paths.folder_names_and_paths), targets[0]) + orig, _ = folder_paths.folder_names_and_paths.get(target, ([], {})) + folder_paths.folder_names_and_paths[key] = (orig or base, {".gguf"}) + if base and base != orig: + logging.warning(f"Unknown file list already present on key {key}: {base}") + +# Add a custom keys for files ending in .gguf +update_folder_names_and_paths("unet_gguf", ["diffusion_models", "unet"]) +update_folder_names_and_paths("clip_gguf", ["text_encoders", "clip"]) +update_folder_names_and_paths("lora_gguf", ["loras"]) +update_folder_names_and_paths("vae_gguf", ["vae"]) +update_folder_names_and_paths("latent_upscale_models_gguf", ["latent_upscale_models"]) + + +class GGUFLoadProgress: + """Coordinate one ComfyUI progress bar across one or more model files.""" + + def __init__(self, paths): + self.path_totals = { + path: gguf_tensor_count(path) if path.endswith(".gguf") else 1 + for path in paths + } + self.total = sum(self.path_totals.values()) + self.pbar = comfy.utils.ProgressBar(self.total) + self.completed = 0 + + def callback_for(self, path): + offset = self.completed + total = self.path_totals[path] + + def update(current, _loader_total): + self.pbar.update_absolute(offset + current, self.total) + + return update + + def complete_file(self, path): + self.completed += self.path_totals[path] + self.pbar.update_absolute(self.completed, self.total) + +class GGUFModelPatcher(comfy.model_patcher.ModelPatcher): + patch_on_device = False + + def _evict_gguf_quantized_caches(self): + for module in self.model.modules(): + evict = getattr(module, "evict_quantized_caches", None) + if evict is not None: + evict() + + def _move_gguf_quantized_caches(self, device): + for module in self.model.modules(): + move = getattr(module, "move_fused_caches", None) + if move is not None: + move(device) + + def _prepare_gguf_quantized_weights(self, device=None): + """Eagerly prepare fallback caches for active, non-native INT4 patches.""" + prepared = [] + layout = [] + active_modules = [] + for name, module in self.model.named_modules(): + signature = getattr(module, "_fused_patch_signature", None) + prepare = getattr(module, "prepare_fused_weight", None) + if signature is None or prepare is None: + continue + module_signature = signature() + layout.append((name, module_signature)) + if getattr(module, "weight_function", ()) or getattr(module, "bias_function", ()): + active_modules.append((name, module)) + + layout = tuple(layout) + previous_layout = getattr(self, "_gguf_patch_layout_signature", None) + if previous_layout is not None and previous_layout != layout: + self._evict_gguf_quantized_caches() + self._gguf_patch_layout_signature = layout + + if device is None: + # Fallback cache preparation is compute-heavy. Use the execution + # device rather than Dynamic VRAM's commonly-CPU offload device. + device = getattr(self, "load_device", None) + if device is None: + device = getattr(self, "offload_device", None) + device = torch.device(device) if device is not None else torch.device("cpu") + if not active_modules: + return 0 + + interrupt = getattr(comfy.model_management, "throw_exception_if_processing_interrupted", None) + cuda_context = getattr(comfy.model_management, "cuda_device_context", None) + device_context = cuda_context(device) if callable(cuda_context) else nullcontext() + try: + with device_context: + for _, module in active_modules: + if interrupt is not None: + interrupt() + if module.prepare_fused_weight(device): + prepared.append(module) + # Fallback-patched INT4 layers retain a full-precision matrix + # for correctness. Stage it on the offload device so all + # layers do not accumulate on the execution GPU at load time. + move = getattr(module, "move_fused_caches", None) + offload_device = getattr(self, "offload_device", None) + if move is not None and offload_device is not None and device != offload_device: + move(offload_device) + except BaseException: + # Do not leave a mixture of old and newly fused representations after + # an interrupt or a failed adapter preparation. + self._evict_gguf_quantized_caches() + raise + return len(prepared) + + def patch_weight_to_device(self, key, device_to=None, inplace_update=False): + if key not in self.patches: + return + weight = comfy.utils.get_attr(self.model, key) + + patches = self.patches[key] + # Legacy Q4_CR (INT4) path: the base weight stays packed (plain + # Parameter) and adapters run through the module's weight_function + # (native low-rank bypass or fused fallback). Fusing directly with + # calculate_weight here would corrupt the packed INT4 bytes. + # Standard GGML quants and all other tensors keep their paths below + # unchanged. + if key.rsplit('.', 1)[-1] == "weight" and '.' in key: + module = None + try: + module = comfy.utils.get_attr(self.model, key.rsplit('.', 1)[0]) + except Exception: + module = None + install_patch = getattr(module, "install_patch_entries", None) + if install_patch is not None and bool(getattr(module, "_quantized", False)): + out_weight = weight.to(device_to) + patches = move_patch_to_device(patches, self.load_device if self.patch_on_device else self.offload_device) + install_patch(patches, key) + logging.warning(f"Q4_CR type GGUF is best used with the f{UnetLoaderGGUFDynamicVRAM.TITLE} loader node, because using it with Legacy VRAM is much slower.") + if inplace_update: + comfy.utils.copy_to_param(self.model, key, out_weight) + else: + comfy.utils.set_attr_param(self.model, key, out_weight) + return + if is_quantized(weight): + out_weight = weight.to(device_to) + patches = move_patch_to_device(patches, self.load_device if self.patch_on_device else self.offload_device) + module = comfy.utils.get_attr(self.model, key.rsplit('.', 1)[0]) + install_patch = getattr(module, "install_patch_entries", None) + if install_patch is not None: + install_patch(patches, key) + # TODO: do we ever have legitimate duplicate patches? (i.e. patch on top of patched weight) + out_weight.patches = [(patches, key)] + else: + inplace_update = self.weight_inplace_update or inplace_update + if key not in self.backup: + self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])( + weight.to(device=self.offload_device, copy=inplace_update), inplace_update + ) + + if device_to is not None: + temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True) + else: + temp_weight = weight.to(torch.float32, copy=True) + + out_weight = comfy.lora.calculate_weight(patches, temp_weight, key) + out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype) + + if inplace_update: + comfy.utils.copy_to_param(self.model, key, out_weight) + else: + comfy.utils.set_attr_param(self.model, key, out_weight) + + def unpatch_model(self, device_to=None, unpatch_weights=True): + if unpatch_weights: + self._evict_gguf_quantized_caches() + self._gguf_patch_layout_signature = None + for module in self.model.modules(): + if getattr(module, "_gguf_static_patch", False): + module.weight_function = [] + module._gguf_static_patch = False + if unpatch_weights: + for p in self.model.parameters(): + if is_torch_compatible(p): + continue + patches = getattr(p, "patches", []) + if len(patches) > 0: + p.patches = [] + # TODO: Find another way to not unload after patches + return super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights) + + + def pin_weight_to_device(self, key): + op_key = key.rsplit('.', 1)[0] + if not self.mmap_released and op_key in self.named_modules_to_munmap: + # TODO: possible to OOM, find better way to detach + self.named_modules_to_munmap[op_key].to(self.load_device).to(self.offload_device) + del self.named_modules_to_munmap[op_key] + super().pin_weight_to_device(key) + + mmap_released = False + named_modules_to_munmap = {} + + def partially_unload(self, *args, force_patch_weights=False, **kwargs): + result = super().partially_unload(*args, force_patch_weights=True, **kwargs) + if args: + self._move_gguf_quantized_caches(args[0]) + elif "device_to" in kwargs: + self._move_gguf_quantized_caches(kwargs["device_to"]) + return result + + def partially_load(self, *args, force_patch_weights=False, **kwargs): + return super().partially_load(*args, force_patch_weights=True, **kwargs) + + def load(self, *args, force_patch_weights=False, **kwargs): + if not self.mmap_released: + self.named_modules_to_munmap = dict(self.model.named_modules()) + + # always call `patch_weight_to_device` even for lowvram + super().load(*args, force_patch_weights=True, **kwargs) + self._prepare_gguf_quantized_weights() + + # make sure nothing stays linked to mmap after first load + if not self.mmap_released: + linked = [] + if kwargs.get("lowvram_model_memory", 0) > 0: + for n, m in self.named_modules_to_munmap.items(): + if hasattr(m, "weight"): + device = getattr(m.weight, "device", None) + if device == self.offload_device: + linked.append((n, m)) + continue + if hasattr(m, "bias"): + device = getattr(m.bias, "device", None) + if device == self.offload_device: + linked.append((n, m)) + continue + if linked and self.load_device != self.offload_device: + logging.info(f"Attempting to release mmap ({len(linked)})") + for n, m in linked: + # TODO: possible to OOM, find better way to detach + m.to(self.load_device).to(self.offload_device) + self.mmap_released = True + self.named_modules_to_munmap = {} + + def clone(self, *args, **kwargs): + src_cls = self.__class__ + self.__class__ = GGUFModelPatcher + n = super().clone(*args, **kwargs) + n.__class__ = GGUFModelPatcher + self.__class__ = src_cls + # GGUF specific clone values below + n.patch_on_device = getattr(self, "patch_on_device", False) + n.mmap_released = getattr(self, "mmap_released", False) + n._gguf_patch_layout_signature = getattr(self, "_gguf_patch_layout_signature", None) + if src_cls != GGUFModelPatcher: + n.size = 0 # force recalc + return n + +class UnetLoaderGGUF: + @classmethod + def INPUT_TYPES(s): + unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")] + return { + "required": { + "unet_name": (unet_names,), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "load_unet" + CATEGORY = "bootleg" + TITLE = "Unet Loader (GGUF, Legacy node)" + + def load_unet(self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None): + unet_path = folder_paths.get_full_path("unet", unet_name) + progress = GGUFLoadProgress([unet_path]) + sd, extra = gguf_sd_loader(unet_path, progress_callback=progress.callback_for(unet_path)) + progress.complete_file(unet_path) + + mode = extra.get("gguf_quant_mode") + if mode == "int8_convrot": + # Use ComfyUI native INT8 path (weights stay INT8) + ops = get_gguf_q8_ops(compute_dtype=torch.bfloat16)() + elif mode == "int4_pytorch": + raise RuntimeError( + "Q4_PT is retired because PyTorch's Ampere INT4 kernel is not " + "performance-competitive. Reconvert the model as Q8_CR." + ) + elif mode == "int4_cr_w4a4": + # Q4_CR_W4A4: custom W4A4 INT4 backed by comfy_kitchen's fast + # ConvRot int4 tensor-core MMA. + ops = get_gguf_q4_w4a4_ops(compute_dtype=torch.bfloat16)() + else: + ops = GGMLOps() + + if dequant_dtype in ("default", None): + ops.Linear.dequant_dtype = None + elif dequant_dtype in ["target"]: + ops.Linear.dequant_dtype = dequant_dtype + else: + ops.Linear.dequant_dtype = getattr(torch, dequant_dtype) + + if patch_dtype in ("default", None): + ops.Linear.patch_dtype = None + elif patch_dtype in ["target"]: + ops.Linear.patch_dtype = patch_dtype + else: + ops.Linear.patch_dtype = getattr(torch, patch_dtype) + + # init model + + kwargs = {} + valid_params = inspect.signature(comfy.sd.load_diffusion_model_state_dict).parameters + if "metadata" in valid_params: + kwargs["metadata"] = extra.get("metadata", {}) + + model = comfy.sd.load_diffusion_model_state_dict( + sd, model_options={"custom_operations": ops}, **kwargs, + ) + if model is None: + logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path)) + raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path)) + model = GGUFModelPatcher.clone(model) + model.patch_on_device = patch_on_device + return (model,) + + +class VAELoaderGGUF: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "vae_name": (folder_paths.get_filename_list("vae_gguf"),), + } + } + + RETURN_TYPES = ("VAE",) + FUNCTION = "load_vae" + CATEGORY = "bootleg" + TITLE = "VAE Loader (GGUF, Legacy node)" + + def load_vae(self, vae_name): + vae_path = folder_paths.get_full_path("vae", vae_name) + progress = GGUFLoadProgress([vae_path]) + sd, extra = gguf_sd_loader( + vae_path, + handle_prefix=None, + progress_callback=progress.callback_for(vae_path), + ) + progress.complete_file(vae_path) + + if extra["arch_str"] != "minimax_h3_vae": + raise ValueError( + "VAE Loader (GGUF) currently supports only MiniMax H3 video VAE GGUF files." + ) + if extra.get("gguf_quant_mode") != "int8_convrot": + raise ValueError( + "MiniMax H3 VAE GGUF must use Q8_CR so decoder Linear weights stay on the native INT8 path." + ) + + for key, value in tuple(sd.items()): + if isinstance(value, GGMLTensor): + sd[key] = dequantize_tensor(value, dtype=torch.Tensor(value).dtype) + + operations = get_gguf_q8_ops(compute_dtype=torch.float16)() + vae_kwargs = {"sd": sd, "metadata": extra.get("metadata", {})} + vae_init_params = inspect.signature(comfy.sd.VAE.__init__).parameters + if "operations" in vae_init_params: + vae_kwargs["operations"] = operations + if "disable_dynamic" in vae_init_params: + vae_kwargs["disable_dynamic"] = True + vae = comfy.sd.VAE(**vae_kwargs) + else: + minimax_vae = comfy.ldm.minimax.vae + if not hasattr(minimax_vae, "ops"): + raise RuntimeError( + "This ComfyUI version cannot inject Q8_CR operations into MiniMax H3 VAE." + ) + original_operations = minimax_vae.ops + minimax_vae.ops = operations + try: + vae = comfy.sd.VAE(**vae_kwargs) + finally: + minimax_vae.ops = original_operations + vae.throw_exception_if_invalid() + return (vae,) + + +class LTXVLatentUpscaleModelLoaderGGUF: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model_name": ( + folder_paths.get_filename_list("latent_upscale_models_gguf"), + ), + } + } + + RETURN_TYPES = ("LATENT_UPSCALE_MODEL",) + FUNCTION = "load_model" + CATEGORY = "bootleg" + TITLE = "LTXV Latent Upscale Model Loader (GGUF)" + + def load_model(self, model_name): + model_path = folder_paths.get_full_path( + "latent_upscale_models_gguf", model_name + ) + progress = GGUFLoadProgress([model_path]) + state_dict, extra = gguf_sd_loader( + model_path, + handle_prefix=None, + progress_callback=progress.callback_for(model_path), + ) + progress.complete_file(model_path) + + if extra["arch_str"] != "ltxv_upscaler": + raise ValueError( + "LTXV Latent Upscale Model Loader (GGUF) requires an LTX 2.5 " + "latent upscaler GGUF file." + ) + config_json = extra["metadata"].get("config") + if config_json is None: + raise ValueError("LTX 2.5 latent upscaler GGUF is missing its config metadata.") + + from comfy.ldm.lightricks.latent_upsampler import LatentUpsampler + + config = json.loads(config_json) + model = LatentUpsampler.from_config(config, operations=GGMLOps) + model_dtype = comfy.model_management.vae_dtype( + allowed_dtypes=[torch.bfloat16, torch.float32] + ) + model = model.to(dtype=model_dtype) + comfy.model_management.archive_model_dtypes(model) + model_patcher = comfy.model_patcher.CoreModelPatcher( + model, + load_device=comfy.model_management.get_torch_device(), + offload_device=comfy.model_management.unet_offload_device(), + ) + model.load_state_dict(state_dict, assign=model_patcher.is_dynamic()) + return (model_patcher,) + + +class TargetedQuantizationGGUF: + """Convert a source checkpoint to GGUF from a ComfyUI workflow.""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "source_path": ( + "STRING", + { + "default": "", + "tooltip": "Absolute path to a .safetensors, .ckpt, .pt, .bin, or .pth source model.", + }, + ), + "destination_path": ( + "STRING", + { + "default": "", + "tooltip": "Output .gguf path. Leave empty to derive it from the source path.", + }, + ), + "quantization": ( + [TARGET_SIZE_QUANT_TYPE, *QUANT_TYPE_MAP.keys()], + { + "default": TARGET_SIZE_QUANT_TYPE, + "tooltip": ( + "TARGET_SIZE starts at the selected Q8 type, reduces central core matrices " + "to Q5_0 then Q4_0, then ordinary 1-D tensors to BF16 only when necessary." + ), + }, + ), + "max_size_mb": ( + "FLOAT", + { + "default": 0.0, + "min": 0.0, + "max": 1000000.0, + "step": 1.0, + "tooltip": "Maximum output size in MiB. Required only for TARGET_SIZE.", + }, + ), + "target_size_q8_type": ( + list(TARGET_SIZE_Q8_TYPES), + { + "default": DEFAULT_TARGET_SIZE_Q8_TYPE, + "tooltip": ( + "TARGET_SIZE baseline: Q8_CR uses native INT8 ConvRot; " + "Q8_0 uses standard GGUF Q8 before layers are reduced to Q4_0." + ), + }, + ), + "quantization_device": ( + list(QUANTIZATION_DEVICE_OPTIONS), + { + "default": "auto", + "tooltip": ( + "Q8_CR conversion device. auto uses CUDA when available and " + "falls back to CPU per matrix when VRAM is insufficient." + ), + }, + ), + "overwrite": ("BOOLEAN", {"default": False}), + "streamed": ( + "BOOLEAN", + { + "default": False, + "tooltip": ( + "For safetensors sources, process one tensor at a time and stage " + "GGUF data on disk to reduce RAM and VRAM usage." + ), + }, + ), + }, + "optional": { + "lora_paths": ( + "STRING", + { + "multiline": True, + "default": "", + "tooltip": "Absolute .safetensors or .gguf LoRA paths, one per line or comma-separated.", + }, + ), + "lora_strengths": ( + "STRING", + { + "default": "", + "tooltip": "Comma-separated merge strengths matching lora_paths; blank uses 1.0.", + }, + ), + }, + } + + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("gguf_path", "quantization_info") + FUNCTION = "quantize" + CATEGORY = "bootleg/quantization" + TITLE = "Targeted Quantization (GGUF)" + + def quantize( + self, + source_path, + destination_path, + quantization, + max_size_mb, + target_size_q8_type, + quantization_device, + overwrite, + streamed=False, + lora_paths="", + lora_strengths="", + ): + source_path = os.path.abspath(os.path.expanduser(source_path)) + if not os.path.isfile(source_path): + raise FileNotFoundError(f"Source model does not exist: {source_path}") + + if quantization == TARGET_SIZE_QUANT_TYPE and max_size_mb <= 0: + raise ValueError("TARGET_SIZE requires max_size_mb greater than zero.") + lora_paths = _parse_lora_paths(lora_paths) if lora_paths.strip() else [] + lora_strengths = _parse_strengths(lora_strengths, len(lora_paths)) if lora_paths else [] + + progress = {"bar": None, "read_total": None, "total": None} + + def report_progress(stage, current, total): + if stage == "read": + if progress["bar"] is None: + progress["read_total"] = total + progress["total"] = total * 2 + 1 + progress["bar"] = comfy.utils.ProgressBar(progress["total"]) + progress["bar"].update_absolute(current, progress["total"]) + elif stage == "quantize": + if progress["bar"] is None: + progress["read_total"] = 0 + progress["total"] = total + 1 + progress["bar"] = comfy.utils.ProgressBar(progress["total"]) + progress["bar"].update_absolute(progress["read_total"] + current, progress["total"]) + + output_path, _ = convert_file( + source_path, + dst_path=destination_path or None, + interact=False, + overwrite=overwrite, + quant_type_name=None if quantization == TARGET_SIZE_QUANT_TYPE else quantization, + max_size_mb=max_size_mb if quantization == TARGET_SIZE_QUANT_TYPE else None, + target_size_q8_type=target_size_q8_type, + quantization_device=quantization_device, + progress_callback=report_progress, + lora_paths=lora_paths, + lora_strengths=lora_strengths, + streamed=streamed, + ) + if progress["bar"] is not None: + progress["bar"].update_absolute(progress["total"], progress["total"]) + + output_size_mb = os.path.getsize(output_path) / (1024 * 1024) + info = f"{quantization}: {output_size_mb:.2f} MiB written to {output_path}" + return (output_path, info) + + +def _gguf_lora_path(lora_name): + path = folder_paths.get_full_path("loras", lora_name) + if path is None: + raise FileNotFoundError(f"GGUF LoRA does not exist: {lora_name}") + if not path.lower().endswith(".gguf"): + raise ValueError(f"GGUF LoRA Import requires a .gguf file, got {lora_name!r}.") + return path + + +def _gguf_lora_key_map(model, clip): + key_map = {} + if model is not None: + key_map = comfy.lora.model_lora_keys_unet(model.model, key_map) + if clip is not None: + key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map) + return key_map + + +def _remap_gguf_lora_for_comfy(targets, key_map): + lora = {} + missing = [] + for target_name, target in targets.items(): + candidates = ( + target_name, + f"diffusion_model.{target_name}", + f"text_encoders.{target_name}", + target_name.removeprefix("diffusion_model."), + target_name.removeprefix("text_encoders."), + ) + mapped_name = next((candidate for candidate in candidates if candidate in key_map), None) + if mapped_name is None: + missing.append(target["base_name"]) + continue + lora[f"{mapped_name}.lora_A.weight"] = target["down"] + lora[f"{mapped_name}.lora_B.weight"] = target["up"] + if target["alpha"] is not None: + lora[f"{mapped_name}.alpha"] = torch.tensor( + target["alpha"], dtype=torch.float32 + ) + if missing: + raise ValueError( + "GGUF LoRA targets do not match the connected MODEL/CLIP: " + + ", ".join(sorted(missing)) + ) + return lora + + +class GGUFLoraImport: + """Load standard GGUF LoRA factors through ComfyUI's normal patch mechanism.""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "lora_name": (folder_paths.get_filename_list("lora_gguf"),), + "strength_model": ( + "FLOAT", + {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}, + ), + }, + "optional": { + "clip": ("CLIP",), + "strength_clip": ( + "FLOAT", + {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}, + ), + }, + } + + RETURN_TYPES = ("MODEL", "CLIP") + FUNCTION = "load_lora" + CATEGORY = "bootleg/LoRA" + TITLE = "Load LoRA (GGUF)" + + def load_lora(self, model, lora_name, strength_model, clip=None, strength_clip=1.0): + path = _gguf_lora_path(lora_name) + _, targets, metadata = load_gguf_lora(path) + key_map = _gguf_lora_key_map(model, clip) + lora = _remap_gguf_lora_for_comfy(targets, key_map) + return comfy.sd.load_lora_for_models( + model, + clip, + lora, + strength_model, + strength_clip, + lora_metadata={"gguf_lora": metadata}, + ) + + +def _parse_lora_paths(lora_paths): + paths = [ + os.path.abspath(os.path.expanduser(path.strip())) + for path in lora_paths.replace(",", "\n").splitlines() + if path.strip() + ] + if not paths: + raise ValueError("Provide at least one LoRA path.") + for path in paths: + if not path.lower().endswith((".gguf", ".safetensors")): + raise ValueError(f"LoRA fusion accepts only .gguf or .safetensors adapters, got {path!r}.") + if not os.path.isfile(path): + raise FileNotFoundError(f"LoRA does not exist: {path}") + return paths + + +def _parse_strengths(strengths, count): + values = [value.strip() for value in strengths.replace("\n", ",").split(",") if value.strip()] + if not values: + return [1.0] * count + if len(values) != count: + raise ValueError( + f"Expected {count} LoRA strength value(s), received {len(values)}." + ) + try: + return [float(value) for value in values] + except ValueError as error: + raise ValueError("LoRA strengths must be comma-separated numbers.") from error + + +def _require_dynamic_vram(): + if not comfy.memory_management.aimdo_enabled: + raise RuntimeError( + "Dynamic VRAM is not enabled in this ComfyUI installation. " + "Start ComfyUI without --disable-dynamic-vram and use a build that supports DynamicVRAM." + ) + + +def _clone_as_dynamic_gguf_patcher(model_patcher): + source_class = model_patcher.__class__ + model_patcher.__class__ = GGUFModelPatcherDynamic + cloned = model_patcher.clone() + model_patcher.__class__ = source_class + return cloned + + +def _legacy_gguf_ops(extra): + mode = extra.get("gguf_quant_mode") + if mode == "int8_convrot": + return get_gguf_q8_ops(compute_dtype=torch.bfloat16)() + if mode == "int4_pytorch": + raise RuntimeError( + "Q4_PT is retired because PyTorch's Ampere INT4 kernel is not " + "performance-competitive. Reconvert the model as Q8_CR." + ) + if mode == "int4_cr_w4a4": + return get_gguf_q4_w4a4_ops(compute_dtype=torch.bfloat16)() + return GGMLOps() + + +def _load_dynamic_gguf_unet(unet_path, disable_dynamic=False, progress=None): + if not disable_dynamic: + _require_dynamic_vram() + progress = progress or GGUFLoadProgress([unet_path]) + sd, extra = gguf_sd_loader( + unet_path, + dynamic=not disable_dynamic, + progress_callback=progress.callback_for(unet_path), + ) + progress.complete_file(unet_path) + + kwargs = {} + valid_params = inspect.signature(comfy.sd.load_diffusion_model_state_dict).parameters + if "metadata" in valid_params: + kwargs["metadata"] = extra.get("metadata", {}) + + # Target-size models combine native Q8_CR layers with standard GGML Q4_0 + # layers. Dynamic VRAM's default mixed-precision Linear cannot serialize a + # bare GGML Q4_0 weight, so use the GGUF Q8 ops for both paths. Those ops + # retain Q8_CR metadata and materialize only standard GGML layers as needed. + model_options = { + "custom_operations": _legacy_gguf_ops(extra), + } + model = comfy.sd.load_diffusion_model_state_dict( + sd, + model_options=model_options, + disable_dynamic=disable_dynamic, + **kwargs, + ) + if model is None: + logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path)) + raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path)) + model = GGUFModelPatcher.clone(model) if disable_dynamic else _clone_as_dynamic_gguf_patcher(model) + model.cached_patcher_init = (_load_dynamic_gguf_unet, (unet_path,)) + return model + + +class GGUFModelPatcherDynamic(comfy.model_patcher.ModelPatcherDynamic): + _evict_gguf_quantized_caches = GGUFModelPatcher._evict_gguf_quantized_caches + _move_gguf_quantized_caches = GGUFModelPatcher._move_gguf_quantized_caches + _prepare_gguf_quantized_weights = GGUFModelPatcher._prepare_gguf_quantized_weights + + def _dynamic_vram_lora_cpu_factors(self): + """Return CPU LoRA/LoKr factors and adapters from active low-VRAM patches.""" + execution_device = torch.device(getattr(self, "load_device", "cpu")) + offload_device = torch.device(getattr(self, "offload_device", "cpu")) + if execution_device.type != "cuda" or offload_device.type != "cpu": + return execution_device, offload_device, {}, () + + factors = {} + adapters = {} + for _, module in self.model.named_modules(): + functions = ( + *getattr(module, "weight_function", ()), + *getattr(module, "bias_function", ()), + ) + for patch_function in functions: + if not getattr(patch_function, "is_lowvram_patch", False): + continue + patches = getattr(patch_function, "patches", None) + key = getattr(patch_function, "key", None) + patch_entries = getattr(patch_function, "prepared_patches", None) + if patch_entries is None and patches is not None and key is not None: + patch_entries = patches.get(key) + if patch_entries is None: + continue + for patch in patch_entries: + if len(patch) < 2: + continue + adapter = patch[1] + if getattr(adapter, "name", None) not in {"lora", "lokr"}: + continue + has_cpu_factor = False + for factor in getattr(adapter, "weights", ()): + if isinstance(factor, torch.Tensor) and factor.device.type == "cpu": + factors[id(factor)] = factor + has_cpu_factor = True + if has_cpu_factor: + adapters[id(adapter)] = adapter + return execution_device, offload_device, factors, tuple(adapters.values()) + + def _move_dynamic_vram_lora_factor(self, factor, device): + return factor.to(device) + + def _preload_dynamic_vram_lora_factors(self): + """Atomically place active low-VRAM LoRA/LoKr factors on the execution GPU.""" + execution_device, _, factors, adapters = ( + GGUFModelPatcherDynamic._dynamic_vram_lora_cpu_factors(self) + ) + if not factors or not torch.cuda.is_available(): + return False + + try: + free_bytes, _ = torch.cuda.mem_get_info(execution_device) + except RuntimeError: + return False + + factor_bytes = sum( + factor.numel() * factor.element_size() for factor in factors.values() + ) + if ( + int4_lora_offload_enabled() + and free_bytes - factor_bytes < comfy.model_management.extra_reserved_memory() + ): + return False + + transferred = {} + try: + for factor in factors.values(): + transferred[id(factor)] = self._move_dynamic_vram_lora_factor( + factor, execution_device + ) + except torch.OutOfMemoryError: + transferred.clear() + log_cuda_oom_loaded_models("preloading Dynamic VRAM INT4 LoRA/LoKr factors") + raise + + for adapter in adapters: + adapter.weights = tuple( + transferred.get(id(weight), weight) + if isinstance(weight, torch.Tensor) + else weight + for weight in adapter.weights + ) + logging.info( + "Dynamic VRAM preloaded %d unique LoRA/LoKr factor tensors (%.1f MiB) on %s.", + len(factors), + factor_bytes / 1024**2, + execution_device, + ) + return True + + def _warn_dynamic_vram_lora_streaming(self): + """Warn when active low-VRAM LoRA/LoKr factors must stream to CUDA repeatedly.""" + if not int4_lora_offload_enabled(): + return False + execution_device, offload_device, factors, _ = ( + GGUFModelPatcherDynamic._dynamic_vram_lora_cpu_factors(self) + ) + if execution_device.type != "cuda" or offload_device.type != "cpu": + return False + + factor_bytes = sum(factor.numel() * factor.element_size() for factor in factors.values()) + if factor_bytes < _DYNAMIC_VRAM_LORA_WARNING_MIN_BYTES: + return False + + signature = ( + tuple(sorted((id(factor), factor.numel(), factor.element_size()) for factor in factors.values())), + str(execution_device), + str(offload_device), + ) + if getattr(self, "_gguf_lora_stream_warning_signature", None) == signature: + return False + + memory_report = "CUDA memory telemetry unavailable" + if torch.cuda.is_available(): + try: + free_bytes, total_bytes = torch.cuda.mem_get_info(execution_device) + except RuntimeError: + pass + else: + memory_report = ( + f"CUDA free/total {free_bytes / 1024**3:.1f}/{total_bytes / 1024**3:.1f} GiB" + ) + logging.warning( + "Dynamic VRAM runtime LoRA/LoKr warning: %.1f MiB across %d unique offloaded " + "LoRA/LoKr factor tensors may be repeatedly streamed from %s to %s (%s). " + "This can cause severe slowdown and increase CUDA OOM risk, but does not " + "guarantee an OOM.", + factor_bytes / 1024**2, + len(factors), + offload_device, + execution_device, + memory_report, + ) + self._gguf_lora_stream_warning_signature = signature + return True + + def unpatch_model(self, device_to=None, unpatch_weights=True): + if unpatch_weights: + self._evict_gguf_quantized_caches() + self._gguf_patch_layout_signature = None + return super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights) + + def partially_unload(self, device_to, memory_to_free=0, force_patch_weights=False): + result = super().partially_unload( + device_to, memory_to_free=memory_to_free, force_patch_weights=force_patch_weights + ) + self._move_gguf_quantized_caches(device_to) + return result + + def load(self, *args, **kwargs): + super().load(*args, **kwargs) + # GGML weights cannot be requantized after applying a LoRA patch. + for _, module in self.model.named_modules(): + for param_key in ("weight", "bias"): + attr = f"{param_key}_lowvram_function" + lowvram_function = getattr(module, attr, None) + if lowvram_function is not None: + setattr(module, attr, None) + functions = getattr(module, f"{param_key}_function", []) + functions.append(lowvram_function) + setattr(module, f"{param_key}_function", functions) + self._preload_dynamic_vram_lora_factors() + self._warn_dynamic_vram_lora_streaming() + self._prepare_gguf_quantized_weights() + + def clone(self, disable_dynamic=False, model_override=None): + if disable_dynamic: + if model_override is None: + fallback = self.cached_patcher_init[0]( + *self.cached_patcher_init[1], + disable_dynamic=True, + ) + model_override = fallback.get_clone_model_override() + return GGUFModelPatcher.clone(self, model_override=model_override) + return super().clone(disable_dynamic=disable_dynamic, model_override=model_override) + + +class UnetLoaderGGUFDynamicVRAM(UnetLoaderGGUF): + TITLE = "Unet Loader (GGUF, Dynamic VRAM)" + + def load_unet(self, unet_name, **kwargs): + unet_path = folder_paths.get_full_path("unet", unet_name) + return (_load_dynamic_gguf_unet(unet_path, progress=GGUFLoadProgress([unet_path])),) + +class UnetLoaderGGUFAdvanced(UnetLoaderGGUF): + @classmethod + def INPUT_TYPES(s): + unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")] + return { + "required": { + "unet_name": (unet_names,), + "dequant_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}), + "patch_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}), + "patch_on_device": ("BOOLEAN", {"default": False}), + } + } + TITLE = "Unet Loader (GGUF/Advanced)" + +class CLIPLoaderGGUF: + @classmethod + def INPUT_TYPES(s): + base = nodes.CLIPLoader.INPUT_TYPES() + return { + "required": { + "clip_name": (s.get_filename_list(),), + "type": base["required"]["type"], + } + } + + RETURN_TYPES = ("CLIP",) + FUNCTION = "load_clip" + CATEGORY = "bootleg" + TITLE = "CLIPLoader (GGUF, Legacy node)" + + @classmethod + def get_filename_list(s): + files = [] + files += folder_paths.get_filename_list("clip") + files += folder_paths.get_filename_list("clip_gguf") + return sorted(files) + + def load_data(self, ckpt_paths): + clip_data = [] + progress = GGUFLoadProgress(ckpt_paths) + for p in ckpt_paths: + if p.endswith(".gguf"): + sd = gguf_clip_loader(p, progress_callback=progress.callback_for(p)) + else: + sd = comfy.utils.load_torch_file(p, safe_load=True) + if "scaled_fp8" in sd: # NOTE: Scaled FP8 would require different custom ops, but only one can be active + raise NotImplementedError(f"Mixing scaled FP8 with GGUF is not supported! Use regular CLIP loader or switch model(s)\n({p})") + clip_data.append(sd) + progress.complete_file(p) + return clip_data + + def load_patcher(self, clip_paths, clip_type, clip_data): + clip = comfy.sd.load_text_encoder_state_dicts( + clip_type = clip_type, + state_dicts = clip_data, + model_options = { + "custom_operations": GGMLOps, + "initial_device": comfy.model_management.text_encoder_offload_device() + }, + embedding_directory = folder_paths.get_folder_paths("embeddings"), + ) + clip.patcher = GGUFModelPatcher.clone(clip.patcher) + return clip + + def load_clip(self, clip_name, type="stable_diffusion"): + clip_path = folder_paths.get_full_path("clip", clip_name) + clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) + return (self.load_patcher([clip_path], clip_type, self.load_data([clip_path])),) + + +def _load_dynamic_gguf_clip(clip_paths, clip_type, disable_dynamic=False, progress=None): + if not disable_dynamic: + _require_dynamic_vram() + progress = progress or GGUFLoadProgress(clip_paths) + clip_data = [] + for path in clip_paths: + if path.endswith(".gguf"): + clip_data.append( + gguf_clip_loader( + path, + dynamic=not disable_dynamic, + progress_callback=progress.callback_for(path), + ) + ) + else: + clip_data.append(comfy.utils.load_torch_file(path, safe_load=True)) + progress.complete_file(path) + + model_options = { + "initial_device": comfy.model_management.text_encoder_offload_device(), + } + if disable_dynamic: + model_options["custom_operations"] = GGMLOps + + clip = comfy.sd.load_text_encoder_state_dicts( + clip_type=clip_type, + state_dicts=clip_data, + model_options=model_options, + embedding_directory=folder_paths.get_folder_paths("embeddings"), + disable_dynamic=disable_dynamic, + ) + clip.patcher = ( + GGUFModelPatcher.clone(clip.patcher) + if disable_dynamic + else _clone_as_dynamic_gguf_patcher(clip.patcher) + ) + clip.patcher.cached_patcher_init = (_load_dynamic_gguf_clip_patcher, (clip_paths, clip_type)) + return clip + + +def _load_dynamic_gguf_clip_patcher(clip_paths, clip_type, disable_dynamic=False): + return _load_dynamic_gguf_clip( + clip_paths, + clip_type, + disable_dynamic=disable_dynamic, + ).patcher + + +class CLIPLoaderGGUFDynamicVRAM(CLIPLoaderGGUF): + TITLE = "CLIPLoader (GGUF, Dynamic VRAM)" + + def load_clip(self, clip_name, type="stable_diffusion"): + clip_path = folder_paths.get_full_path("clip", clip_name) + clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) + return (_load_dynamic_gguf_clip([clip_path], clip_type, progress=GGUFLoadProgress([clip_path])),) + +class DualCLIPLoaderGGUF(CLIPLoaderGGUF): + @classmethod + def INPUT_TYPES(s): + base = nodes.DualCLIPLoader.INPUT_TYPES() + file_options = (s.get_filename_list(), ) + return { + "required": { + "clip_name1": file_options, + "clip_name2": file_options, + "type": base["required"]["type"], + } + } + + TITLE = "DualCLIPLoader (GGUF, Legacy node)" + + def load_clip(self, clip_name1, clip_name2, type): + clip_path1 = folder_paths.get_full_path("clip", clip_name1) + clip_path2 = folder_paths.get_full_path("clip", clip_name2) + clip_paths = (clip_path1, clip_path2) + clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) + return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) + + +class DualCLIPLoaderGGUFDynamicVRAM(DualCLIPLoaderGGUF): + TITLE = "DualCLIPLoader (GGUF, Dynamic VRAM)" + + def load_clip(self, clip_name1, clip_name2, type): + clip_paths = ( + folder_paths.get_full_path("clip", clip_name1), + folder_paths.get_full_path("clip", clip_name2), + ) + clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) + return (_load_dynamic_gguf_clip(clip_paths, clip_type),) + +class TripleCLIPLoaderGGUF(CLIPLoaderGGUF): + @classmethod + def INPUT_TYPES(s): + file_options = (s.get_filename_list(), ) + return { + "required": { + "clip_name1": file_options, + "clip_name2": file_options, + "clip_name3": file_options, + } + } + + TITLE = "TripleCLIPLoader (GGUF, Legacy node)" + + def load_clip(self, clip_name1, clip_name2, clip_name3, type="sd3"): + clip_path1 = folder_paths.get_full_path("clip", clip_name1) + clip_path2 = folder_paths.get_full_path("clip", clip_name2) + clip_path3 = folder_paths.get_full_path("clip", clip_name3) + clip_paths = (clip_path1, clip_path2, clip_path3) + clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) + return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) + + +class TripleCLIPLoaderGGUFDynamicVRAM(TripleCLIPLoaderGGUF): + TITLE = "TripleCLIPLoader (GGUF, Dynamic VRAM)" + + def load_clip(self, clip_name1, clip_name2, clip_name3, type="sd3"): + clip_paths = ( + folder_paths.get_full_path("clip", clip_name1), + folder_paths.get_full_path("clip", clip_name2), + folder_paths.get_full_path("clip", clip_name3), + ) + clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) + return (_load_dynamic_gguf_clip(clip_paths, clip_type),) + +class QuadrupleCLIPLoaderGGUF(CLIPLoaderGGUF): + @classmethod + def INPUT_TYPES(s): + file_options = (s.get_filename_list(), ) + return { + "required": { + "clip_name1": file_options, + "clip_name2": file_options, + "clip_name3": file_options, + "clip_name4": file_options, + } + } + + TITLE = "QuadrupleCLIPLoader (GGUF, Legacy node)" + + def load_clip(self, clip_name1, clip_name2, clip_name3, clip_name4, type="stable_diffusion"): + clip_path1 = folder_paths.get_full_path("clip", clip_name1) + clip_path2 = folder_paths.get_full_path("clip", clip_name2) + clip_path3 = folder_paths.get_full_path("clip", clip_name3) + clip_path4 = folder_paths.get_full_path("clip", clip_name4) + clip_paths = (clip_path1, clip_path2, clip_path3, clip_path4) + clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) + return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),) + + +class QuadrupleCLIPLoaderGGUFDynamicVRAM(QuadrupleCLIPLoaderGGUF): + TITLE = "QuadrupleCLIPLoader (GGUF, Dynamic VRAM)" + + def load_clip(self, clip_name1, clip_name2, clip_name3, clip_name4, type="stable_diffusion"): + clip_paths = ( + folder_paths.get_full_path("clip", clip_name1), + folder_paths.get_full_path("clip", clip_name2), + folder_paths.get_full_path("clip", clip_name3), + folder_paths.get_full_path("clip", clip_name4), + ) + clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) + return (_load_dynamic_gguf_clip(clip_paths, clip_type),) + +NODE_CLASS_MAPPINGS = { + "TargetedQuantizationGGUF": TargetedQuantizationGGUF, + "GGUFLoraImport": GGUFLoraImport, + "UnetLoaderGGUF": UnetLoaderGGUF, + "VAELoaderGGUF": VAELoaderGGUF, + "LTXVLatentUpscaleModelLoaderGGUF": LTXVLatentUpscaleModelLoaderGGUF, + "CLIPLoaderGGUF": CLIPLoaderGGUF, + "DualCLIPLoaderGGUF": DualCLIPLoaderGGUF, + "TripleCLIPLoaderGGUF": TripleCLIPLoaderGGUF, + "QuadrupleCLIPLoaderGGUF": QuadrupleCLIPLoaderGGUF, + "UnetLoaderGGUFAdvanced": UnetLoaderGGUFAdvanced, + "UnetLoaderGGUFDynamicVRAM": UnetLoaderGGUFDynamicVRAM, + "CLIPLoaderGGUFDynamicVRAM": CLIPLoaderGGUFDynamicVRAM, + "DualCLIPLoaderGGUFDynamicVRAM": DualCLIPLoaderGGUFDynamicVRAM, + "TripleCLIPLoaderGGUFDynamicVRAM": TripleCLIPLoaderGGUFDynamicVRAM, + "QuadrupleCLIPLoaderGGUFDynamicVRAM": QuadrupleCLIPLoaderGGUFDynamicVRAM, +} diff --git a/ops.py b/ops.py index 88a352e2..e4c659a6 100644 --- a/ops.py +++ b/ops.py @@ -1,281 +1,1325 @@ -# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) -import gguf -import torch -import logging - -import comfy.ops -import comfy.lora -import comfy.model_management -from .dequant import dequantize_tensor, is_quantized - -def chained_hasattr(obj, chained_attr): - probe = obj - for attr in chained_attr.split('.'): - if hasattr(probe, attr): - probe = getattr(probe, attr) - else: - return False - return True - -# A bakcward and forward compatible way to get `torch.compiler.disable`. -def get_torch_compiler_disable_decorator(): - def dummy_decorator(*args, **kwargs): - def noop(x): - return x - return noop - - from packaging import version - - if not chained_hasattr(torch, "compiler.disable"): - logging.info("ComfyUI-GGUF: Torch too old for torch.compile - bypassing") - return dummy_decorator # torch too old - elif version.parse(torch.__version__) >= version.parse("2.8"): - logging.info("ComfyUI-GGUF: Allowing full torch compile") - return dummy_decorator # torch compile works - if chained_hasattr(torch, "_dynamo.config.nontraceable_tensor_subclasses"): - logging.info("ComfyUI-GGUF: Allowing full torch compile (nightly)") - return dummy_decorator # torch compile works, nightly before 2.8 release - else: - logging.info("ComfyUI-GGUF: Partial torch compile only, consider updating pytorch") - return torch.compiler.disable - -torch_compiler_disable = get_torch_compiler_disable_decorator() - -class GGMLTensor(torch.Tensor): - """ - Main tensor-like class for storing quantized weights - """ - def __init__(self, *args, tensor_type, tensor_shape, patches=[], **kwargs): - super().__init__() - self.tensor_type = tensor_type - self.tensor_shape = tensor_shape - self.patches = patches - - def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs): - return super().__new__(cls, *args, **kwargs) - - def to(self, *args, **kwargs): - new = super().to(*args, **kwargs) - new.tensor_type = getattr(self, "tensor_type", None) - new.tensor_shape = getattr(self, "tensor_shape", new.data.shape) - new.patches = getattr(self, "patches", []).copy() - return new - - def clone(self, *args, **kwargs): - return self - - def detach(self, *args, **kwargs): - return self - - def copy_(self, *args, **kwargs): - # fixes .weight.copy_ in comfy/clip_model/CLIPTextModel - try: - return super().copy_(*args, **kwargs) - except Exception as e: - logging.warning(f"ignoring 'copy_' on tensor: {e}") - - def new_empty(self, size, *args, **kwargs): - # Intel Arc fix, ref#50 - new_tensor = super().new_empty(size, *args, **kwargs) - return GGMLTensor( - new_tensor, - tensor_type = getattr(self, "tensor_type", None), - tensor_shape = size, - patches = getattr(self, "patches", []).copy() - ) - - @property - def shape(self): - if not hasattr(self, "tensor_shape"): - self.tensor_shape = self.size() - return self.tensor_shape - -class GGMLLayer(torch.nn.Module): - """ - This (should) be responsible for de-quantizing on the fly - """ - comfy_cast_weights = True - dequant_dtype = None - patch_dtype = None - largest_layer = False - torch_compatible_tensor_types = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16} - - def is_ggml_quantized(self, *, weight=None, bias=None): - if weight is None: - weight = self.weight - if bias is None: - bias = self.bias - return is_quantized(weight) or is_quantized(bias) - - def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): - weight, bias = state_dict.get(f"{prefix}weight"), state_dict.get(f"{prefix}bias") - # NOTE: using modified load for linear due to not initializing on creation, see GGMLOps todo - if self.is_ggml_quantized(weight=weight, bias=bias) or isinstance(self, torch.nn.Linear): - return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs) - # Not strictly required, but fixes embedding shape mismatch. Threshold set in loader.py - if isinstance(self, torch.nn.Embedding) and self.weight.shape[0] >= (64 * 1024): - return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs) - return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) - - def ggml_load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): - prefix_len = len(prefix) - for k,v in state_dict.items(): - if k[prefix_len:] == "weight": - self.weight = torch.nn.Parameter(v, requires_grad=False) - elif k[prefix_len:] == "bias" and v is not None: - self.bias = torch.nn.Parameter(v, requires_grad=False) - else: - unexpected_keys.append(k) - - # For Linear layer with missing weight - if self.weight is None and isinstance(self, torch.nn.Linear): - v = torch.zeros(self.in_features, self.out_features) - self.weight = torch.nn.Parameter(v, requires_grad=False) - missing_keys.append(prefix+"weight") - - # for vram estimation (TODO: less fragile logic?) - if getattr(self.weight, "is_largest_weight", False): - self.largest_layer = True - - def _save_to_state_dict(self, *args, **kwargs): - if self.is_ggml_quantized(): - return self.ggml_save_to_state_dict(*args, **kwargs) - return super()._save_to_state_dict(*args, **kwargs) - - def ggml_save_to_state_dict(self, destination, prefix, keep_vars): - # This is a fake state dict for vram estimation - weight = torch.zeros_like(self.weight, device=torch.device("meta")) - destination[prefix + "weight"] = weight - if self.bias is not None: - bias = torch.zeros_like(self.bias, device=torch.device("meta")) - destination[prefix + "bias"] = bias - - # Take into account space required for dequantizing the largest tensor - if self.largest_layer: - shape = getattr(self.weight, "tensor_shape", self.weight.shape) - dtype = self.dequant_dtype if self.dequant_dtype and self.dequant_dtype != "target" else torch.float16 - temp = torch.empty(*shape, device=torch.device("meta"), dtype=dtype) - destination[prefix + "temp.weight"] = temp - - return - # This would return the dequantized state dict - destination[prefix + "weight"] = self.get_weight(self.weight) - if bias is not None: - destination[prefix + "bias"] = self.get_weight(self.bias) - - def get_weight(self, tensor, dtype): - if tensor is None: - return - - # consolidate and load patches to GPU in async - patch_list = [] - device = tensor.device - for patches, key in getattr(tensor, "patches", []): - patch_list += move_patch_to_device(patches, device) - - # dequantize tensor while patches load - weight = dequantize_tensor(tensor, dtype, self.dequant_dtype) - - # prevent propagating custom tensor class - if isinstance(weight, GGMLTensor): - weight = torch.Tensor(weight) - - # apply patches - if len(patch_list) > 0: - if self.patch_dtype is None: - weight = comfy.lora.calculate_weight(patch_list, weight, key) - else: - # for testing, may degrade image quality - patch_dtype = dtype if self.patch_dtype == "target" else self.patch_dtype - weight = comfy.lora.calculate_weight(patch_list, weight, key, patch_dtype) - return weight - - @torch_compiler_disable() - def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None): - if input is not None: - if dtype is None: - dtype = getattr(input, "dtype", torch.float32) - if bias_dtype is None: - bias_dtype = dtype - if device is None: - device = input.device - - bias = None - non_blocking = comfy.model_management.device_supports_non_blocking(device) - if s.bias is not None: - bias = s.get_weight(s.bias.to(device), dtype) - bias = comfy.ops.cast_to(bias, bias_dtype, device, non_blocking=non_blocking, copy=False) - - weight = s.get_weight(s.weight.to(device), dtype) - weight = comfy.ops.cast_to(weight, dtype, device, non_blocking=non_blocking, copy=False) - return weight, bias - - def forward_comfy_cast_weights(self, input, *args, **kwargs): - if self.is_ggml_quantized(): - out = self.forward_ggml_cast_weights(input, *args, **kwargs) - else: - out = super().forward_comfy_cast_weights(input, *args, **kwargs) - - # non-ggml forward might still propagate custom tensor class - if isinstance(out, GGMLTensor): - out = torch.Tensor(out) - return out - - def forward_ggml_cast_weights(self, input): - raise NotImplementedError - -class GGMLOps(comfy.ops.manual_cast): - """ - Dequantize weights on the fly before doing the compute - """ - class Linear(GGMLLayer, comfy.ops.manual_cast.Linear): - def __init__(self, in_features, out_features, bias=True, device=None, dtype=None): - torch.nn.Module.__init__(self) - # TODO: better workaround for reserved memory spike on windows - # Issue is with `torch.empty` still reserving the full memory for the layer - # Windows doesn't over-commit memory so without this 24GB+ of pagefile is used - self.in_features = in_features - self.out_features = out_features - self.weight = None - self.bias = None - - def forward_ggml_cast_weights(self, input): - weight, bias = self.cast_bias_weight(input) - return torch.nn.functional.linear(input, weight, bias) - - class Conv2d(GGMLLayer, comfy.ops.manual_cast.Conv2d): - def forward_ggml_cast_weights(self, input): - weight, bias = self.cast_bias_weight(input) - return self._conv_forward(input, weight, bias) - - class Embedding(GGMLLayer, comfy.ops.manual_cast.Embedding): - def forward_ggml_cast_weights(self, input, out_dtype=None): - output_dtype = out_dtype - if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16: - out_dtype = None - weight, _bias = self.cast_bias_weight(self, device=input.device, dtype=out_dtype) - return torch.nn.functional.embedding( - input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse - ).to(dtype=output_dtype) - - class LayerNorm(GGMLLayer, comfy.ops.manual_cast.LayerNorm): - def forward_ggml_cast_weights(self, input): - if self.weight is None: - return super().forward_comfy_cast_weights(input) - weight, bias = self.cast_bias_weight(input) - return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps) - - class GroupNorm(GGMLLayer, comfy.ops.manual_cast.GroupNorm): - def forward_ggml_cast_weights(self, input): - weight, bias = self.cast_bias_weight(input) - return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps) - -def move_patch_to_device(item, device): - if isinstance(item, torch.Tensor): - return item.to(device, non_blocking=True) - elif isinstance(item, tuple): - return tuple(move_patch_to_device(x, device) for x in item) - elif isinstance(item, list): - return [move_patch_to_device(x, device) for x in item] - else: - return item +# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import copy +import gguf +import json +import torch +import logging +import os +import time +from numbers import Real + +import comfy.ops +import comfy.lora +import comfy.model_management +from .dequant import dequantize_tensor, is_quantized + + +_PERF_LOG_ENV = "COMFYUI_GGUF_PERF_LOG" +_INT4_LORA_OFFLOAD_ENV = "COMFYUI_GGUF_INT4_LORA_OFFLOAD" + + +def int4_lora_offload_enabled(): + return os.environ.get(_INT4_LORA_OFFLOAD_ENV, "").strip().lower() in { + "true", + "1", + "yes", + "on", + } + + +def log_cuda_oom_loaded_models(context): + """Log ComfyUI's managed model inventory before propagating a CUDA OOM.""" + loaded_models = getattr(comfy.model_management, "current_loaded_models", ()) + logging.error("ComfyUI-GGUF CUDA OOM while %s. Loaded models:", context) + if not loaded_models: + logging.error(" (none)") + return + + for loaded_model in loaded_models: + patcher = getattr(loaded_model, "model", None) + model = getattr(patcher, "model", None) + model_class = type(model).__name__ if model is not None else "" + model_name = getattr(model, "name", None) or getattr(patcher, "model_name", None) + identifier = ( + f"{model_class} ({model_name})" + if isinstance(model_name, str) and model_name + else model_class + ) + + size_bytes = None + model_memory = getattr(loaded_model, "model_memory", None) + if callable(model_memory): + try: + size_bytes = model_memory() + except (AttributeError, TypeError): + pass + if not isinstance(size_bytes, Real): + size_bytes = getattr(patcher, "size", None) + + if isinstance(size_bytes, Real) and size_bytes >= 0: + logging.error(" %s: %.1f MiB", identifier, size_bytes / 1024**2) + else: + logging.error(" %s: size unavailable", identifier) + + +def _configure_perf_logger(): + log_path = os.environ.get(_PERF_LOG_ENV, "").strip() + if not log_path: + return None + if log_path.strip().lower() in {"1", "true", "yes", "on"}: + log_path = "comfyui-gguf-performance.log" + + logger = logging.getLogger("comfyui_gguf.performance") + logger.setLevel(logging.INFO) + logger.propagate = False + try: + handler = logging.FileHandler(log_path, encoding="utf-8") + except (OSError, ValueError) as error: + logging.getLogger(__name__).warning( + "ComfyUI-GGUF: unable to open performance log %r: %s", log_path, error + ) + return None + handler.setFormatter(logging.Formatter("%(asctime)s %(message)s")) + handler._comfyui_gguf_perf = True + logger.addHandler(handler) + logger.info("performance logging enabled path=%s", log_path) + return logger + + +_PERF_LOGGER = _configure_perf_logger() + + +def _perf_sync(device): + if device is not None and getattr(device, "type", None) == "cuda" and torch.cuda.is_available(): + torch.cuda.synchronize(device) + + +def _perf_forward(mode, layer, input_tensor, operation): + """Run and optionally time one quantized Linear forward. + + CUDA synchronization is deliberately limited to the opt-in diagnostic path; + normal inference does not pay for timing or synchronization overhead. + """ + if _PERF_LOGGER is None: + return operation() + + device = getattr(input_tensor, "device", None) + _perf_sync(device) + started = time.perf_counter() + call_number = getattr(layer, "_gguf_perf_calls", 0) + 1 + layer._gguf_perf_calls = call_number + cache_before = getattr(layer, "_quantized_weight", None) is not None + fused_cache_before = getattr(layer, "_fused_weight", None) is not None + try: + result = operation() + _perf_sync(device) + except BaseException: + _perf_sync(device) + elapsed_ms = (time.perf_counter() - started) * 1000 + _PERF_LOGGER.exception( + "forward mode=%s call=%d in_features=%s out_features=%s " + "input_shape=%s device=%s elapsed_ms=%.3f failed=true", + mode, + call_number, + getattr(layer, "in_features", None), + getattr(layer, "out_features", None), + tuple(getattr(input_tensor, "shape", ())), + device, + elapsed_ms, + ) + raise + + elapsed_ms = (time.perf_counter() - started) * 1000 + _PERF_LOGGER.info( + "forward mode=%s call=%d in_features=%s out_features=%s " + "input_shape=%s device=%s elapsed_ms=%.3f cache_before=%s cache_after=%s " + "fused_cache_before=%s fused_cache_after=%s patched=%s", + mode, + call_number, + getattr(layer, "in_features", None), + getattr(layer, "out_features", None), + tuple(getattr(input_tensor, "shape", ())), + device, + elapsed_ms, + cache_before, + getattr(layer, "_quantized_weight", None) is not None, + fused_cache_before, + getattr(layer, "_fused_weight", None) is not None, + bool(getattr(layer, "weight_function", ())) or bool(getattr(layer, "bias_function", ())), + ) + return result + +def _build_regular_hadamard(size, dtype=torch.float32, device="cpu"): + """Build a normalized regular (Sylvester) Hadamard of a power-of-4 size.""" + if size < 4 or (size & (size - 1)) != 0: + import math + if not math.log(size, 4).is_integer(): + raise ValueError(f"Regular Hadamard size must be a power of 4, got {size}") + h4 = torch.tensor( + [[1, 1, 1, -1], [1, 1, -1, 1], [1, -1, 1, 1], [-1, 1, 1, 1]], + dtype=dtype, + device=device, + ) + h = h4 + current_size = 4 + while current_size < size: + h = torch.kron(h, h4) + current_size *= 4 + return h / (size ** 0.5) + +def _valid_compute_dtype(dtype): + return dtype in {torch.float16, torch.bfloat16, torch.float32, torch.float64} + +def _infer_compute_dtype(tensor_type, fallback=None): + if _valid_compute_dtype(fallback): + return fallback + if tensor_type == gguf.GGMLQuantizationType.BF16: + return torch.bfloat16 + if tensor_type == gguf.GGMLQuantizationType.F32: + return torch.float32 + return torch.float16 + +def chained_hasattr(obj, chained_attr): + probe = obj + for attr in chained_attr.split('.'): + if hasattr(probe, attr): + probe = getattr(probe, attr) + else: + return False + return True + +# A bakcward and forward compatible way to get `torch.compiler.disable`. +def get_torch_compiler_disable_decorator(): + def dummy_decorator(*args, **kwargs): + def noop(x): + return x + return noop + + from packaging import version + + if not chained_hasattr(torch, "compiler.disable"): + logging.info("ComfyUI-GGUF: Torch too old for torch.compile - bypassing") + return dummy_decorator # torch too old + elif version.parse(torch.__version__) >= version.parse("2.8"): + logging.info("ComfyUI-GGUF: Allowing full torch compile") + return dummy_decorator # torch compile works + if chained_hasattr(torch, "_dynamo.config.nontraceable_tensor_subclasses"): + logging.info("ComfyUI-GGUF: Allowing full torch compile (nightly)") + return dummy_decorator # torch compile works, nightly before 2.8 release + else: + logging.info("ComfyUI-GGUF: Partial torch compile only, consider updating pytorch") + return torch.compiler.disable + +torch_compiler_disable = get_torch_compiler_disable_decorator() + +class GGMLTensor(torch.Tensor): + """ + Main tensor-like class for storing quantized weights + """ + def __init__(self, *args, tensor_type, tensor_shape, patches=[], compute_dtype=None, **kwargs): + super().__init__() + self.tensor_type = tensor_type + self.tensor_shape = tensor_shape + self.patches = patches + self.compute_dtype = compute_dtype + + def __new__(cls, *args, tensor_type, tensor_shape, patches=[], compute_dtype=None, **kwargs): + return super().__new__(cls, *args, **kwargs) + + def to(self, *args, **kwargs): + new = super().to(*args, **kwargs) + new.tensor_type = getattr(self, "tensor_type", None) + new.tensor_shape = getattr(self, "tensor_shape", new.data.shape) + new.patches = getattr(self, "patches", []).copy() + new.compute_dtype = getattr(self, "compute_dtype", None) + return new + + def clone(self, *args, **kwargs): + return self + + def detach(self, *args, **kwargs): + return self + + def copy_(self, *args, **kwargs): + # fixes .weight.copy_ in comfy/clip_model/CLIPTextModel + try: + return super().copy_(*args, **kwargs) + except Exception as e: + logging.warning(f"ignoring 'copy_' on tensor: {e}") + + def new_empty(self, size, *args, **kwargs): + # Intel Arc fix, ref#50 + new_tensor = super().new_empty(size, *args, **kwargs) + return GGMLTensor( + new_tensor, + tensor_type = getattr(self, "tensor_type", None), + tensor_shape = size, + patches = getattr(self, "patches", []).copy(), + compute_dtype = getattr(self, "compute_dtype", None), + ) + + @property + def dtype(self): + qtype = getattr(self, "tensor_type", None) + if qtype in GGMLLayer.torch_compatible_tensor_types: + # NOTE: use the base-class descriptor instead of torch.Tensor(self): + # constructing a Tensor from an inference-mode tensor raises + # "Inference tensors do not track version counter" (hit via + # low_vram_patch_estimate_vram on bias keys when LoRAs patch biases) + return torch.Tensor.dtype.__get__(self) + return _infer_compute_dtype(qtype, getattr(self, "compute_dtype", None)) + + @property + def shape(self): + if not hasattr(self, "tensor_shape"): + self.tensor_shape = self.size() + return self.tensor_shape + +class GGMLLayer(torch.nn.Module): + """ + This (should) be responsible for de-quantizing on the fly + """ + comfy_cast_weights = True + dequant_dtype = None + patch_dtype = None + largest_layer = False + torch_compatible_tensor_types = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16} + + def is_ggml_quantized(self, *, weight=None, bias=None): + if weight is None: + weight = self.weight + if bias is None: + bias = self.bias + return is_quantized(weight) or is_quantized(bias) + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + weight, bias = state_dict.get(f"{prefix}weight"), state_dict.get(f"{prefix}bias") + # NOTE: using modified load for linear due to not initializing on creation, see GGMLOps todo + if self.is_ggml_quantized(weight=weight, bias=bias) or isinstance(self, torch.nn.Linear): + return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs) + # Not strictly required, but fixes embedding shape mismatch. Threshold set in loader.py + if isinstance(self, torch.nn.Embedding) and self.weight.shape[0] >= (64 * 1024): + return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs) + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + def ggml_load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): + prefix_len = len(prefix) + for k,v in state_dict.items(): + if k[prefix_len:] == "weight": + if isinstance(v, GGMLTensor): + v.compute_dtype = self._ggml_compute_dtype(v, "weight") + self.weight = torch.nn.Parameter(v, requires_grad=False) + elif k[prefix_len:] == "bias" and v is not None: + if isinstance(v, GGMLTensor): + v.compute_dtype = self._ggml_compute_dtype(v, "bias") + self.bias = torch.nn.Parameter(v, requires_grad=False) + else: + unexpected_keys.append(k) + + # For Linear layer with missing weight + if self.weight is None and isinstance(self, torch.nn.Linear): + v = torch.zeros(self.in_features, self.out_features) + self.weight = torch.nn.Parameter(v, requires_grad=False) + missing_keys.append(prefix+"weight") + + # for vram estimation (TODO: less fragile logic?) + if getattr(self.weight, "is_largest_weight", False): + self.largest_layer = True + + def _ggml_compute_dtype(self, tensor, param_name): + if self.dequant_dtype is not None and self.dequant_dtype != "target": + return self.dequant_dtype + model_dtype = getattr(self, f"{param_name}_comfy_model_dtype", None) + return _infer_compute_dtype(getattr(tensor, "tensor_type", None), model_dtype) + + def _save_to_state_dict(self, *args, **kwargs): + if self.is_ggml_quantized(): + return self.ggml_save_to_state_dict(*args, **kwargs) + return super()._save_to_state_dict(*args, **kwargs) + + def ggml_save_to_state_dict(self, destination, prefix, keep_vars): + # This is a fake state dict for vram estimation + weight = torch.zeros_like(self.weight, device=torch.device("meta")) + destination[prefix + "weight"] = weight + if self.bias is not None: + bias = torch.zeros_like(self.bias, device=torch.device("meta")) + destination[prefix + "bias"] = bias + + # Take into account space required for dequantizing the largest tensor + if self.largest_layer: + shape = getattr(self.weight, "tensor_shape", self.weight.shape) + dtype = self.dequant_dtype if self.dequant_dtype and self.dequant_dtype != "target" else torch.float16 + temp = torch.empty(*shape, device=torch.device("meta"), dtype=dtype) + destination[prefix + "temp.weight"] = temp + + return + # This would return the dequantized state dict + destination[prefix + "weight"] = self.get_weight(self.weight) + if bias is not None: + destination[prefix + "bias"] = self.get_weight(self.bias) + + def get_weight(self, tensor, dtype): + if tensor is None: + return + + # consolidate and load patches to GPU in async + patch_list = [] + device = tensor.device + for patches, key in getattr(tensor, "patches", []): + patch_list += move_patch_to_device(patches, device) + + # dequantize tensor while patches load + weight = dequantize_tensor(tensor, dtype, self.dequant_dtype) + + # prevent propagating custom tensor class + if isinstance(weight, GGMLTensor): + weight = torch.Tensor(weight) + + # apply patches + if len(patch_list) > 0: + if self.patch_dtype is None: + weight = comfy.lora.calculate_weight(patch_list, weight, key) + else: + # for testing, may degrade image quality + patch_dtype = dtype if self.patch_dtype == "target" else self.patch_dtype + weight = comfy.lora.calculate_weight(patch_list, weight, key, patch_dtype) + return weight + + @torch_compiler_disable() + def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None): + if input is not None: + if dtype is None: + dtype = getattr(input, "dtype", torch.float32) + if bias_dtype is None: + bias_dtype = dtype + if device is None: + device = input.device + + bias = None + non_blocking = comfy.model_management.device_supports_non_blocking(device) + if s.bias is not None: + bias = s.get_weight(s.bias.to(device), dtype) + bias = comfy.ops.cast_to(bias, bias_dtype, device, non_blocking=non_blocking, copy=False) + + weight = s.get_weight(s.weight.to(device), dtype) + weight = comfy.ops.cast_to(weight, dtype, device, non_blocking=non_blocking, copy=False) + return weight, bias + + def forward_comfy_cast_weights(self, input, *args, **kwargs): + if self.is_ggml_quantized(): + out = self.forward_ggml_cast_weights(input, *args, **kwargs) + else: + out = super().forward_comfy_cast_weights(input, *args, **kwargs) + + # non-ggml forward might still propagate custom tensor class + if isinstance(out, GGMLTensor): + out = torch.Tensor(out) + return out + + def forward_ggml_cast_weights(self, input): + raise NotImplementedError + +class GGMLOps(comfy.ops.manual_cast): + """ + Dequantize weights on the fly before doing the compute + """ + class Linear(GGMLLayer, comfy.ops.manual_cast.Linear): + def __init__(self, in_features, out_features, bias=True, device=None, dtype=None): + torch.nn.Module.__init__(self) + # TODO: better workaround for reserved memory spike on windows + # Issue is with `torch.empty` still reserving the full memory for the layer + # Windows doesn't over-commit memory so without this 24GB+ of pagefile is used + self.in_features = in_features + self.out_features = out_features + self.weight = None + self.bias = None + self.weight_comfy_model_dtype = dtype + self.bias_comfy_model_dtype = dtype + + def forward_ggml_cast_weights(self, input): + weight, bias = self.cast_bias_weight(input) + return torch.nn.functional.linear(input, weight, bias) + + class Conv2d(GGMLLayer, comfy.ops.manual_cast.Conv2d): + def forward_ggml_cast_weights(self, input): + weight, bias = self.cast_bias_weight(input) + return self._conv_forward(input, weight, bias) + + class Conv3d(GGMLLayer, comfy.ops.manual_cast.Conv3d): + def forward_ggml_cast_weights(self, input): + weight, bias = self.cast_bias_weight(input) + return self._conv_forward(input, weight, bias) + + class Embedding(GGMLLayer, comfy.ops.manual_cast.Embedding): + def forward_ggml_cast_weights(self, input, out_dtype=None): + output_dtype = out_dtype + if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16: + out_dtype = None + weight, _bias = self.cast_bias_weight(self, device=input.device, dtype=out_dtype) + return torch.nn.functional.embedding( + input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse + ).to(dtype=output_dtype) + + class LayerNorm(GGMLLayer, comfy.ops.manual_cast.LayerNorm): + def forward_ggml_cast_weights(self, input): + if self.weight is None: + return super().forward_comfy_cast_weights(input) + weight, bias = self.cast_bias_weight(input) + return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps) + + class GroupNorm(GGMLLayer, comfy.ops.manual_cast.GroupNorm): + def forward_ggml_cast_weights(self, input): + weight, bias = self.cast_bias_weight(input) + return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps) + +def move_patch_to_device(item, device): + if isinstance(item, torch.Tensor): + return item.to(device, non_blocking=True) + elif isinstance(item, tuple): + return tuple(move_patch_to_device(x, device) for x in item) + elif isinstance(item, list): + return [move_patch_to_device(x, device) for x in item] + else: + return item + +def get_gguf_q8_ops(compute_dtype=torch.bfloat16, full_precision_mm=False): + """ + Factory for an ops class that uses ComfyUI's native mixed_precision_ops INT8 path. + Weights are kept as INT8 and matmul uses comfy_kitchen's TensorWiseINT8Layout. + """ + BaseOps = comfy.ops.mixed_precision_ops( + quant_config={}, + compute_dtype=compute_dtype, + full_precision_mm=full_precision_mm, + ) + + class GGUFQ8Ops(BaseOps): + class Linear(BaseOps.Linear): + def __init__(self, in_features, out_features, bias=True, device=None, dtype=None): + # Lazy init: don't allocate weight here; it will be loaded from state dict + torch.nn.Module.__init__(self) + self.factory_kwargs = {"device": device, "dtype": BaseOps._compute_dtype} + self.in_features = in_features + self.out_features = out_features + self.weight = None + if bias: + self.bias = torch.nn.Parameter(torch.empty(out_features, **self.factory_kwargs)) + else: + self.register_parameter("bias", None) + self._orig_shape = (out_features, in_features) + self.tensor_class = None + self._full_precision_mm = BaseOps._full_precision_mm + self._full_precision_mm_config = False + + def forward(self, *args, **kwargs): + input_tensor = args[0] if args else kwargs.get("input") + parent_forward = super().forward + return _perf_forward( + "int8_tensorwise", + self, + input_tensor, + lambda: parent_forward(*args, **kwargs), + ) + + def _load_from_state_dict(self, *args): + state_dict, prefix = args[:2] + weight_key = f"{prefix}weight" + # Target-size GGUFs combine native Q8_CR weights with standard + # Q4_0 weights. The mixed-precision loader only understands + # the former's comfy_quant metadata, so materialize standard + # GGML weights before handing them to that loader. + if weight_key in state_dict and f"{prefix}comfy_quant" not in state_dict: + weight = state_dict[weight_key] + if is_quantized(weight): + state_dict[weight_key] = dequantize_tensor( + weight, + dtype=self.factory_kwargs["dtype"], + ) + elif hasattr(weight, "dequantize"): + state_dict[weight_key] = weight.dequantize().to( + dtype=self.factory_kwargs["dtype"], + ) + return comfy.ops._load_quantized_module( + self, + torch.nn.Module._load_from_state_dict.__get__(self, type(self)), + *args, + load_extra_params=True, + ) + + return GGUFQ8Ops + + +# Q4_CR_W4A4: custom W4A4 INT4 format backed by comfy_kitchen's fast ConvRot +# int4 tensor-core MMA. On-disk is kitchen-native packed int4 (N, K//2) int8 +# + a per-output-row fp16 scale. The weight is pre-rotated by a block-diagonal +# Hadamard along K; at runtime the kernel rotates the activation into the same +# basis, so output is directly in the original space. +def get_gguf_q4_w4a4_ops(compute_dtype=torch.bfloat16, full_precision_mm=False): + try: + from comfy_kitchen.tensor.convrot_w4a4 import ( + TensorCoreConvRotW4A4Layout, + quantize_convrot_w4a4_weight, + ) + from comfy_kitchen.tensor.base import QuantizedTensor + _HAVE_KITCHEN = True + except Exception: + TensorCoreConvRotW4A4Layout = None + quantize_convrot_w4a4_weight = None + QuantizedTensor = None + _HAVE_KITCHEN = False + + class GGUFQ4W4A4Ops(comfy.ops.disable_weight_init): + class Linear(torch.nn.Module, comfy.ops.CastWeightBiasOp): + comfy_cast_weights = True + + def __init__(self, in_features, out_features, bias=True, device=None, dtype=None): + torch.nn.Module.__init__(self) + self.in_features = in_features + self.out_features = out_features + self.weight = None + self.register_parameter("bias", None) + self._orig_shape = (out_features, in_features) + self._convrot_groupsize = 256 + self._quant_group_size = 64 + self._quantized = False + self.weight_scale = None + self._compute_dtype = torch.bfloat16 + self._quantized_weight = None + self._quantized_weight_device = None + # Cache for a dequantized non-native patch. Re-quantizing a small + # delta to INT4 can round it away; compatible LoRA/LoKr instead use + # the packed base plus their low-rank output correction. + self._fused_weight = None + self._fused_patch_id = None + self._fused_weight_device = None + self._fused_bias = None + self._fused_bias_patch_id = None + self._fused_bias_device = None + self._lora_factor_cache = {} + + def install_patch_entries(self, patches, key): + """Expose a static patch as the same callable used by Dynamic VRAM.""" + self.evict_quantized_caches() + + def apply_patch(weight, *args, **kwargs): + return comfy.lora.calculate_weight(patches, weight, key) + + apply_patch._gguf_static_patch = True + apply_patch._gguf_patch_entries = patches + apply_patch._gguf_patch_key = key + self.weight_function = [apply_patch] + self._gguf_static_patch = True + + def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, + missing_keys, unexpected_keys, error_msgs): + self.evict_quantized_caches() + weight_key = f"{prefix}weight" + scale_key = f"{prefix}weight_scale" + quant_key = f"{prefix}comfy_quant" + + weight = state_dict.pop(weight_key, None) + scale = state_dict.pop(scale_key, None) + quant_raw = state_dict.pop(quant_key, None) + + bias_key = f"{prefix}bias" + bias = state_dict.pop(bias_key, None) + if bias is not None: + self.bias = torch.nn.Parameter(bias, requires_grad=False) + else: + self.bias = None + if bias_key in missing_keys: + missing_keys.remove(bias_key) + + if quant_raw is None or weight is None: + # Plain (non-quantized) Linear: keep the raw weight. + if weight is None: + missing_keys.append(weight_key) + return + self.weight = torch.nn.Parameter( + weight if isinstance(weight, torch.Tensor) else torch.Tensor(weight), + requires_grad=False, + ) + self._quantized = False + return + + if scale is None: + raise RuntimeError(f"Missing Q4_CR_W4A4 scale tensor for {prefix}") + + quant_conf = json.loads(bytes(quant_raw.tolist()).decode("utf-8")) + if quant_conf.get("format") != "int4_cr" or quant_conf.get("backing") != "w4a4": + raise ValueError( + f"Unsupported Q4_CR_W4A4 format for {prefix}: {quant_conf.get('format')!r} " + f"/ {quant_conf.get('backing')!r}" + ) + self._convrot_groupsize = quant_conf.get("convrot_groupsize", 256) + self._quant_group_size = quant_conf.get("quant_group_size", 64) + self._orig_shape = tuple(quant_conf.get("orig_shape", (self.out_features, self.in_features))) + self.out_features = self._orig_shape[0] + self.in_features = self._orig_shape[1] + + self.weight = torch.nn.Parameter(weight, requires_grad=False) + self.weight_scale = torch.nn.Parameter(scale, requires_grad=False) + self._quantized = True + self._compute_dtype = quant_conf.get("orig_dtype", torch.bfloat16) + self._quantized_weight = None + self._quantized_weight_device = None + self._fused_weight = None + self._fused_patch_id = None + self._fused_weight_device = None + self._fused_bias = None + self._fused_bias_patch_id = None + self._fused_bias_device = None + + def evict_quantized_caches(self): + """Release every derived representation of this quantized layer.""" + self._quantized_weight = None + self._quantized_weight_device = None + self._fused_weight = None + self._fused_patch_id = None + self._fused_weight_device = None + self._fused_bias = None + self._fused_bias_patch_id = None + self._fused_bias_device = None + self._fused_patch_keepalive = () + self._lora_factor_cache.clear() + + def move_fused_caches(self, device): + """Move prepared patched caches without rebuilding them.""" + if self._fused_weight is not None and self._fused_weight_device != str(device): + self._fused_weight = self._fused_weight.to(device=device) + self._fused_weight_device = str(device) + if self._fused_bias is not None and self._fused_bias_device != str(device): + self._fused_bias = self._fused_bias.to(device=device) + self._fused_bias_device = str(device) + + def _build_quantized_weight(self, device, dtype, packed=None): + if not _HAVE_KITCHEN: + return None + if packed is None: + packed = self.weight.to(device=device).to(torch.int8) + else: + packed = packed.to(device=device).to(torch.int8) + scale = self.weight_scale.to(device=device, dtype=torch.float32).contiguous() + params = TensorCoreConvRotW4A4Layout.Params( + scale=scale, + orig_dtype=dtype, + orig_shape=self._orig_shape, + convrot_groupsize=self._convrot_groupsize, + quant_group_size=self._quant_group_size, + linear_dtype="int4", + ) + return QuantizedTensor(packed, "TensorCoreConvRotW4A4Layout", params) + + def _get_cached_quantized_weight(self, device, packed=None): + if packed is not None: + # A patched/offloaded packed weight may differ from self.weight + # (e.g. moved to device each forward). Build fresh; don't cache, + # because the packed content can change per forward. + return self._build_quantized_weight(device, self._compute_dtype, packed=packed) + if self._quantized_weight is not None and self._quantized_weight_device == str(device): + return self._quantized_weight + if self._quantized_weight is not None: + self._quantized_weight = None + if hasattr(self, "_quantized_weight_scale") and self._quantized_weight_scale is not None: + # Free the previous device's cached scale/params to avoid a leak. + self._quantized_weight_scale = None + qt = self._build_quantized_weight(device, self._compute_dtype) + self._quantized_weight = qt + self._quantized_weight_device = str(device) + return qt + + def _fused_patch_signature(self): + """A cheap, stable identity for the active weight/bias patch set. + + ComfyUI installs the adapter (LoRA/LoKR) functions once per model + load/patch, so the callable objects are stable across forwards. We key + the fused-weight cache on the callable objects themselves so a re-bound + patch (even if a function object is reused) is detected and re-fused. + + Under Dynamic VRAM, however, ``GGUFModelPatcherDynamic.load`` promotes a + freshly created ``LowVramPatch`` into ``weight_function`` every time the + model is moved to device. That object is recreated on each reload, so + keying the cache on ``id(f)`` would miss on every forward and re-run the + expensive Hadamard rotate + int4 pack (a 12x slowdown on large layers). + For ``LowVramPatch`` we therefore key on its stable *content*: the tensor + key and the identity of the shared patches dict/list (both live on the + model patcher and survive recreation). A genuine patch change (e.g. a new + ``add_patches`` call) rebuilds that list, which is detected as a change. + All other adapter callables are stable objects, so we key on them directly. + """ + sig = [] + keepalive = [] + for f in list(getattr(self, "weight_function", ())) + list(getattr(self, "bias_function", ())): + if getattr(f, "is_lowvram_patch", False) and getattr(f, "key", None): + patches = getattr(f, "patches", None) + plist = patches.get(f.key) if patches is not None else None + if patches is not None: + keepalive.append(patches) + if plist is not None: + keepalive.append(plist) + sig.append(("lowvram", f.key, id(patches), id(plist))) + else: + sig.append(("callable", f)) + # Keep references alive so address reuse never aliases a stale signature. + self._fused_patch_keepalive = tuple(keepalive) + return tuple(sig) + + def _get_cached_fused_bias(self, device, patch_id=None): + bias_functions = list(getattr(self, "bias_function", ())) + if not bias_functions: + return self.bias.to(device=device, dtype=self._compute_dtype) if self.bias is not None else None + patch_id = patch_id if patch_id is not None else self._fused_patch_signature() + if ( + self._fused_bias is not None + and self._fused_bias_patch_id == patch_id + ): + # The cache stays on the lifecycle-managed offload device. The + # caller creates a temporary device/dtype copy for the operation; + # moving this persistent tensor here would retain every patched + # layer on CUDA after it executes. + return self._fused_bias + if self.bias is None: + self._fused_bias = None + self._fused_bias_patch_id = patch_id + self._fused_bias_device = str(device) + return None + bias = self.bias.to(device=device, dtype=self._compute_dtype) + for f in bias_functions: + bias = f(bias) + self._fused_bias = bias + self._fused_bias_patch_id = patch_id + self._fused_bias_device = str(device) + return bias + + def prepare_fused_weight(self, device): + """Prepare fallback caches for active patch sets before inference. + + Compatible standard LoRA/LoKr patches keep the native INT4 base and + bypass with a low-rank output correction. Other patch forms use a + fused compute-dtype cache, kept on one device only. + """ + if not self._quantized or self.weight is None: + return False + weight_functions = list(getattr(self, "weight_function", ())) + bias_functions = list(getattr(self, "bias_function", ())) + if not weight_functions and not bias_functions: + if self._fused_weight is not None or self._fused_bias is not None: + self.evict_quantized_caches() + return False + + patch_id = self._fused_patch_signature() + if self._fused_patch_id is not None and self._fused_patch_id != patch_id: + self.evict_quantized_caches() + + fused = True + if weight_functions: + # Standard LoRA can use the packed INT4 base plus an exact + # low-rank output correction. Do not eagerly expand all such + # layers into full-precision matrices during model loading. + native_entries = self._native_lora_patch_entries() + if native_entries: + fused = False + if self._fused_weight is not None: + self._fused_weight = None + self._fused_patch_id = None + self._fused_weight_device = None + else: + self._quantized_weight = None + self._quantized_weight_device = None + try: + fused = self._get_cached_fused_weight(device) is not None + except torch.OutOfMemoryError: + if torch.device(device).type != "cuda": + raise + self._move_derived_caches_to_cpu() + logging.warning( + "Q4_CR_W4A4 CUDA fallback for INT4 adapter layer: " + "CUDA ran out of memory while preparing the patched weight." + ) + fused = self._get_cached_fused_weight(torch.device("cpu")) is not None + if bias_functions: + self._get_cached_fused_bias(device, patch_id=patch_id) + return fused + + def _get_cached_fused_weight(self, device): + """Apply a non-native patch once and cache the compute-dtype result. + + Re-quantizing a patched INT4 matrix can erase the small LoRA delta, so + patched layers use the dequantized compute-dtype matrix. The pristine + path below still uses the native INT4 kernel. + """ + # A patched layer must never retain the ordinary packed cache alongside + # either a fused representation or the dequantized fallback. + self._quantized_weight = None + self._quantized_weight_device = None + sig = self._fused_patch_signature() + if ( + self._fused_weight is not None + and self._fused_patch_id == sig + ): + # Keep the persistent fused matrix on its offload device. A + # forward may copy it temporarily to its input device, but must + # not promote the cache permanently; otherwise Dynamic VRAM + # accumulates one full-precision matrix for every LoRA layer. + return self._fused_weight + if self._fused_weight is not None: + self.evict_quantized_caches() + else: + # Do not retain the ordinary packed device cache while the + # patched representation is being prepared. + self._quantized_weight = None + self._quantized_weight_device = None + # Apply adapters in the compute (BF16) domain on the un-rotated weight. + fused = self._dequantized_weight(device, self._compute_dtype) + for f in getattr(self, "weight_function", ()): + fused = f(fused) + fused = fused.to(device=device, dtype=self._compute_dtype) + self._fused_weight = fused + self._fused_patch_id = sig + self._fused_weight_device = str(device) + self._quantized_weight = None + self._quantized_weight_device = None + return fused + + def _native_lora_patch_entries(self): + """Return active low-rank patches that can bypass native INT4. + + A native base matmul plus the adapter's additive output is exact for + ordinary LoRA/LoKr patches and avoids expanding the complete weight. + Other patch forms retain the full-precision fusion path because they + can transform the base weight or otherwise change its shape. + """ + entries = [] + for patch_function in getattr(self, "weight_function", ()): + if getattr(patch_function, "is_lowvram_patch", False): + patches = getattr(patch_function, "patches", None) + key = getattr(patch_function, "key", None) + patch_list = patches.get(key) if patches is not None else None + elif getattr(patch_function, "_gguf_static_patch", False): + patch_list = getattr(patch_function, "_gguf_patch_entries", None) + else: + return None + if patch_list is None: + return None + + prepared = getattr(patch_function, "prepared_patches", None) + if prepared is not None: + patch_list = prepared + for patch in patch_list: + if len(patch) != 5: + return None + strength, adapter, strength_model, offset, function = patch + if ( + strength_model != 1.0 + or offset is not None + or function is not None + or isinstance(adapter, list) + or getattr(adapter, "name", None) not in {"lora", "lokr"} + or not callable(getattr(adapter, "h", None)) + ): + return None + weights = getattr(adapter, "weights", ()) + # The current adapter h() implementations do not implement + # DoRA rescaling. Keep those patches on the exact fusion path. + dora_index = {"lora": 4, "lokr": 8}.get(adapter.name) + if dora_index is not None and len(weights) > dora_index and weights[dora_index] is not None: + return None + entries.append((strength, adapter)) + return entries + + def _get_cached_lora_factor(self, factor, device, dtype): + """Move a CPU adapter factor to CUDA, retaining it only with headroom.""" + if not isinstance(factor, torch.Tensor): + return factor + device = torch.device(device) + try: + if factor.device.type != "cpu" or device.type != "cuda": + return comfy.model_management.cast_to_device(factor, device, dtype) + + key = (id(factor), str(device), dtype) + cached = self._lora_factor_cache.get(key) + if cached is not None and cached[0] is factor: + return cached[1] + + free_memory, _total_memory = torch.cuda.mem_get_info(device) + cache_bytes = factor.numel() * torch.empty((), dtype=dtype).element_size() + if ( + free_memory + <= comfy.model_management.extra_reserved_memory() + cache_bytes + ): + return comfy.model_management.cast_to_device(factor, device, dtype) + cached_factor = comfy.model_management.cast_to_device(factor, device, dtype) + except torch.OutOfMemoryError: + log_cuda_oom_loaded_models("caching INT4 LoRA/LoKr factors") + raise + self._lora_factor_cache[key] = (factor, cached_factor) + return cached_factor + + def _native_lora_bypass(self, input, base_out): + entries = self._native_lora_patch_entries() + if not entries: + return None + out = base_out + for strength, adapter in entries: + if adapter.name == "lora": + up, down, alpha, mid, dora_scale, reshape = adapter.weights + if mid is not None or dora_scale is not None or reshape is not None: + return None + up = self._get_cached_lora_factor(up, input.device, input.dtype) + down = self._get_cached_lora_factor(down, input.device, input.dtype) + rank = down.shape[0] + scale = (alpha / rank) if alpha is not None else 1.0 + delta = torch.nn.functional.linear( + torch.nn.functional.linear(input, down), up + ) * scale + else: + # LoKr's h() is already a low-rank/Kronecker matmul. Its + # implementation casts factors to the input dtype but does + # not move them to the input device, so use a shallow, + # device-local adapter view instead of mutating the + # persistent CPU/offload adapter. + local_adapter = copy.copy(adapter) + local_adapter.weights = tuple( + self._get_cached_lora_factor( + weight, input.device, input.dtype + ) + if isinstance(weight, torch.Tensor) else weight + for weight in getattr(adapter, "weights", ()) + ) + delta = local_adapter.h(input, base_out) + # Avoid materializing a scaled copy of delta and a second output + # tensor. This matters for large Flux activations where the + # residual can be hundreds of MiB even though the base is INT4. + out.add_(delta, alpha=strength) + return out + + def _move_derived_caches_to_cpu(self): + """Release CUDA derived caches before a system-memory retry.""" + self._quantized_weight = None + self._quantized_weight_device = None + self.move_fused_caches(torch.device("cpu")) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def _cpu_forward_fallback(self, input, bias, native_entries=None, use_fused_weight=False): + """Run one failed INT4 layer through CPU system memory.""" + original_device = input.device + cpu = torch.device("cpu") + if original_device.type == "cuda": + self._move_derived_caches_to_cpu() + cpu_input = input.to(device=cpu) + cpu_bias = bias.to(device=cpu, dtype=cpu_input.dtype) if bias is not None else None + + if native_entries: + cpu_weight = self._dequantized_weight(cpu, cpu_input.dtype) + cpu_out = torch.nn.functional.linear(cpu_input, cpu_weight) + if cpu_bias is not None: + cpu_out = cpu_out + cpu_bias + cpu_out = self._native_lora_bypass(cpu_input, cpu_out) + elif use_fused_weight: + cpu_weight = self._get_cached_fused_weight(cpu) + cpu_out = torch.nn.functional.linear( + cpu_input, + cpu_weight.to(device=cpu, dtype=cpu_input.dtype), + cpu_bias, + ) + else: + cpu_weight = self._dequantized_weight(cpu, cpu_input.dtype) + cpu_out = torch.nn.functional.linear(cpu_input, cpu_weight, cpu_bias) + + return cpu_out.to(device=original_device, dtype=input.dtype) + + def forward_comfy_cast_weights(self, input, *args, **kwargs): + has_weight_functions = bool(getattr(self, "weight_function", ())) + has_bias_functions = bool(getattr(self, "bias_function", ())) + if not has_weight_functions and not has_bias_functions and ( + self._fused_weight is not None or self._fused_bias is not None + ): + self.evict_quantized_caches() + patch_id = self._fused_patch_signature() if (has_weight_functions or has_bias_functions) else None + if has_bias_functions and self._fused_bias_patch_id == patch_id: + bias = ( + self._fused_bias.to(device=input.device, dtype=input.dtype) + if self._fused_bias is not None else None + ) + else: + bias = self.bias.to(device=input.device, dtype=input.dtype) if self.bias is not None else None + if bias is not None: + for f in getattr(self, "bias_function", ()): + bias = f(bias) + if not self._quantized or self.weight is None: + # Plain Linear path + weight = self.weight.to(device=input.device, dtype=input.dtype) + return torch.nn.functional.linear(input, weight, bias) + if has_weight_functions or has_bias_functions: + # Compatible LoRA/LoKr uses the native base plus a low-rank output + # correction. Other patch forms use a cached un-rotated BF16 + # fallback rather than re-quantizing a delta that INT4 could erase. + # Bias-only patches retain the native INT4 weight path. + if has_weight_functions: + native_entries = self._native_lora_patch_entries() + if ( + native_entries is None + and self._fused_weight is not None + and self._fused_weight.device.type == "cpu" + and input.device.type == "cuda" + ): + return self._cpu_forward_fallback( + input, bias, use_fused_weight=True + ) + try: + weight_qt = self._get_cached_quantized_weight(input.device) + except torch.OutOfMemoryError: + if input.device.type != "cuda": + raise + logging.warning( + "Q4_CR_W4A4 CUDA fallback for INT4 adapter layer: " + "CUDA ran out of memory; retrying this layer in system memory." + ) + return self._cpu_forward_fallback( + input, + bias, + native_entries=native_entries, + use_fused_weight=not bool(native_entries), + ) + if weight_qt is not None: + base_out = None + try: + base_out = torch.nn.functional.linear(input, weight_qt) + if bias is not None: + base_out = base_out + bias + bypass_out = self._native_lora_bypass(input, base_out) + if bypass_out is not None: + return bypass_out + except torch.OutOfMemoryError: + if input.device.type != "cuda": + raise + del base_out + del weight_qt + logging.warning( + "Q4_CR_W4A4 CUDA fallback for INT4 adapter layer: " + "CUDA ran out of memory; retrying this layer in system memory." + ) + return self._cpu_forward_fallback( + input, bias, native_entries=native_entries + ) + try: + fused_weight = self._get_cached_fused_weight(input.device) + except torch.OutOfMemoryError: + if input.device.type != "cuda": + raise + logging.warning( + "Q4_CR_W4A4 CUDA fallback for INT4 adapter layer: " + "CUDA ran out of memory; retrying this layer in system memory." + ) + return self._cpu_forward_fallback(input, bias, use_fused_weight=True) + if fused_weight.device.type == "cpu" and input.device.type == "cuda": + return self._cpu_forward_fallback(input, bias, use_fused_weight=True) + try: + out = torch.nn.functional.linear( + input, + fused_weight.to(device=input.device, dtype=input.dtype), + ) + if bias is not None: + out = out + bias + return out + except torch.OutOfMemoryError: + if input.device.type != "cuda": + raise + del fused_weight + logging.warning( + "Q4_CR_W4A4 CUDA fallback for INT4 adapter layer: " + "CUDA ran out of memory; retrying this layer in system memory." + ) + return self._cpu_forward_fallback(input, bias, use_fused_weight=True) + + try: + weight_qt = self._get_cached_quantized_weight(input.device) + if weight_qt is None: + # No comfy_kitchen / non-CUDA: dequant fallback. + _orig_w = self._dequantized_weight(input.device, input.dtype) + return torch.nn.functional.linear(input, _orig_w, bias) + + out = torch.nn.functional.linear(input, weight_qt) + if bias is not None: + out = out + bias + return out + except torch.OutOfMemoryError: + if input.device.type != "cuda": + raise + logging.warning( + "Q4_CR_W4A4 CUDA fallback for INT4 layer: " + "CUDA ran out of memory; retrying this layer in system memory." + ) + return self._cpu_forward_fallback(input, bias) + + def forward(self, *args, **kwargs): + input_tensor = args[0] if args else kwargs.get("input") + + def run_forward(): + comfy.ops.run_every_op() + if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0: + return self.forward_comfy_cast_weights(*args, **kwargs) + return torch.nn.functional.linear( + input=input_tensor, + weight=self.weight, + bias=self.bias, + ) + + return _perf_forward("int4_convrot_w4a4", self, input_tensor, run_forward) + + def _dequantized_weight(self, device, dtype): + packed = self.weight.to(device=device).to(torch.int8) + scale = self.weight_scale.to(device=device, dtype=dtype) + n, k_half = packed.shape + k = k_half * 2 + x32 = packed.to(torch.int32) + lo = (x32 & 0xF).to(torch.int8) + hi = ((x32 >> 4) & 0xF).to(torch.int8) + nibbles = torch.stack([lo, hi], dim=-1).reshape(n, k).to(torch.float32) + # Signed two's-complement int4 -> [-8, 7], scaled per output row. + nibbles = torch.where(nibbles >= 8, nibbles - 16, nibbles) + w_rot = nibbles * scale.to(torch.float32).reshape(-1, 1) + # The weight was pre-rotated by a block-diagonal Hadamard along K; + # un-rotate so this dequant matches the kernel's effective weight. + cg = self._convrot_groupsize + if k % cg == 0: + h = _build_regular_hadamard(cg, dtype=torch.float32, device=device) + n_groups = k // cg + w_rot = (w_rot.reshape(n, n_groups, cg) @ h).reshape(n, k) + return w_rot.to(dtype) + + def reset_parameters(self): + return None + + return GGUFQ4W4A4Ops + + +# Retired experimental Q4_PT implementation. No loader or node runtime path +# references this class until a performant W4A16 backend is available. +class RetiredGGUFQ4Ops(comfy.ops.manual_cast): + """ + Ops class for PyTorch's native compact INT4 GEMM. + + Packed weights are created transiently at invocation time. This keeps + low-VRAM offloading viable without retaining a second copy of all weights. + """ + class Linear(torch.nn.Module, comfy.ops.CastWeightBiasOp): + comfy_cast_weights = True + + def __init__(self, in_features, out_features, bias=True, device=None, dtype=None): + torch.nn.Module.__init__(self) + self.in_features = in_features + self.out_features = out_features + self.weight = None + self.register_parameter("bias", None) + self._orig_shape = (out_features, in_features) + self._group_size = None + self._pad = 0 + self._orig_in_features = in_features + self._is_int4 = False + + def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): + weight_key = f"{prefix}weight" + scale_key = f"{prefix}weight_scale" + quant_key = f"{prefix}comfy_quant" + + weight = state_dict.pop(weight_key, None) + scale = state_dict.pop(scale_key, None) + quant_raw = state_dict.pop(quant_key, None) + + bias_key = f"{prefix}bias" + bias = state_dict.pop(bias_key, None) + if quant_raw is None: + if weight is None: + missing_keys.append(weight_key) + return + self.weight = torch.nn.Parameter(torch.Tensor(weight), requires_grad=False) + if bias is not None: + self.bias = torch.nn.Parameter(torch.Tensor(bias), requires_grad=False) + self._is_int4 = False + return + + if weight is None or scale is None: + raise RuntimeError(f"Missing INT4 tensors for {prefix}") + + quant_conf = json.loads(bytes(quant_raw.tolist()).decode("utf-8")) + if quant_conf.get("format") not in {"int4_compact_gemm", "int4_pytorch"}: + raise ValueError(f"Unsupported INT4 format for {prefix}") + self._group_size = quant_conf["group_size"] + self._pad = quant_conf.get("pad", 0) + orig_shape = tuple(quant_conf["orig_shape"]) + self._orig_in_features = orig_shape[1] + self._orig_shape = orig_shape + self.out_features = orig_shape[0] + + self.weight = torch.nn.Parameter(weight, requires_grad=False) + self.weight_scale = torch.nn.Parameter(scale, requires_grad=False) + self._is_int4 = True + + if bias is not None: + self.bias = torch.nn.Parameter(bias, requires_grad=False) + for key in (weight_key, scale_key, quant_key, bias_key): + if key in missing_keys: + missing_keys.remove(key) + + def forward(self, input): + if self.weight is None: + raise RuntimeError("Q4_PT weight was not loaded.") + if not self._is_int4: + weight = self.weight.to(device=input.device, dtype=input.dtype) + bias = self.bias.to(device=input.device, dtype=input.dtype) if self.bias is not None else None + return torch.nn.functional.linear(input, weight, bias) + if self.weight_function or self.bias_function: + raise RuntimeError("Q4_PT does not support weight patches or LoRAs without dequantization.") + input_shape = input.shape + input_2d = input.reshape(-1, input_shape[-1]) + if self._pad: + input_2d = torch.nn.functional.pad(input_2d, (0, self._pad)) + + if self._group_size != 64: + raise RuntimeError( + f"Q4_PT requires PyTorch's group-size-64 INT4 operator, got {self._group_size}." + ) + if input_2d.device.type != "cuda": + raise RuntimeError("Q4_PT requires a CUDA device.") + if input_2d.dtype != torch.bfloat16: + raise RuntimeError( + f"Q4_PT requires BF16 activations for PyTorch's native INT4 operator, got {input_2d.dtype}." + ) + + weight = torch.Tensor( + self.weight.to(device=input.device, dtype=torch.uint8, non_blocking=True) + ) + scale_and_offset = self.weight_scale.to( + device=input.device, + dtype=torch.bfloat16, + non_blocking=True, + ) + packed_features = weight.size(-1) * 2 + native_padding = (-packed_features) % 128 + if native_padding: + # PyTorch's INT4 packer needs K to be a multiple of 128. + # A zero-valued group and zero input features preserve the + # original result for K=64 projections such as Krea2's input. + if native_padding % self._group_size: + raise RuntimeError( + f"Cannot pad Q4_PT input width {packed_features} to PyTorch's INT4 tile size." + ) + input_2d = torch.nn.functional.pad(input_2d, (0, native_padding)) + weight = torch.nn.functional.pad(weight, (0, native_padding // 2)) + scale_and_offset = torch.nn.functional.pad( + scale_and_offset, + (0, 0, 0, 0, 0, native_padding // self._group_size), + ) + + packed_weight = torch._convert_weight_to_int4pack(weight, 8) + output = torch._weight_int4pack_mm( + input_2d, + packed_weight, + self._group_size, + scale_and_offset, + ) + if self.bias is not None: + output = output + self.bias.to(device=input.device, dtype=input.dtype) + return output.reshape(*input_shape[:-1], self.out_features) diff --git a/pyproject.toml b/pyproject.toml index 8b8cf53b..46219ae3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,14 +1,30 @@ -[project] -name = "comfyui-gguf" -description = "GGUF Quantization support for native ComfyUI models." -version = "2.0.0" # 2.0.0 = GitHub main, 1.X.X = ComfyUI Registry -license = { file = "LICENSE" } -dependencies = ["gguf>=0.13.0", "sentencepiece", "protobuf"] - -[project.urls] -Repository = "https://github.com/city96/ComfyUI-GGUF" - -[tool.comfy] -PublisherId = "city96" -DisplayName = "ComfyUI-GGUF" -Icon = "" +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "comfyui-gguf-reboot" +version = "26.09.04" +description = "Actively maintained fork of ComfyUI-GGUF node pack with support for new models, INT8/INT4 GGUFs and optionally ComfyUI Dynamic VRAM" +readme = "README.md" +license = { file = "LICENSE" } +authors = [ + { name = "molbal" } +] +dependencies = [ + "gguf>=0.13.0", + "torch", + "numpy", + "sentencepiece", + "protobuf", + "tqdm", +] + +[project.urls] +Repository = "https://github.com/molbal/ComfyUI-GGUF" +Documentation = "https://github.com/molbal/ComfyUI-GGUF#readme" + +[tool.comfy] +PublisherId = "molbal" +DisplayName = "ComfyUI GGUF (molbal's fork)" +Icon = "https://raw.githubusercontent.com/molbal/ComfyUI-GGUF/refs/heads/main/icon.png" diff --git a/quant_ops.py b/quant_ops.py new file mode 100644 index 00000000..80d1adbb --- /dev/null +++ b/quant_ops.py @@ -0,0 +1,70 @@ +# GGML QuantizedTensor support for ComfyUI DynamicVRAM loading. +from dataclasses import dataclass + +import gguf +import torch + +from comfy_kitchen.tensor import ( + BaseLayoutParams, + QuantizedLayout, + QuantizedTensor, + register_layout_class, +) + +from .dequant import TORCH_COMPATIBLE_QTYPES, dequantize_functions + + +@dataclass(frozen=True) +class GGMLLayoutParams(BaseLayoutParams): + tensor_type: int + + +class GGMLLayout(QuantizedLayout): + Params = GGMLLayoutParams + + @classmethod + def quantize(cls, tensor, **kwargs): + raise NotImplementedError("Quantization to GGML format is not supported") + + @classmethod + def dequantize(cls, qdata, params): + qtype = gguf.GGMLQuantizationType(params.tensor_type) + orig_shape = params.orig_shape + + if qtype in TORCH_COMPATIBLE_QTYPES: + return qdata.reshape(orig_shape).to(params.orig_dtype) + + if qtype not in dequantize_functions: + dequantized = gguf.quants.dequantize(qdata.cpu().numpy(), qtype) + return torch.from_numpy(dequantized).reshape(orig_shape).to( + device=qdata.device, + dtype=params.orig_dtype, + ) + + _, type_size = gguf.GGML_QUANT_SIZES[qtype] + raw = qdata.reshape(-1).view(torch.uint8) + blocks = raw.reshape((raw.numel() // type_size, type_size)) + return dequantize_functions[qtype](blocks, *gguf.GGML_QUANT_SIZES[qtype], None).reshape( + orig_shape + ).to(params.orig_dtype) + + @classmethod + def get_plain_tensors(cls, qtensor): + return (qtensor._qdata,) + + @classmethod + def state_dict_tensors(cls, qdata, params): + return {"weight": qdata} + + +register_layout_class("GGMLLayout", GGMLLayout) + + +def make_quantized(qdata, tensor_type, tensor_shape, orig_dtype=torch.float16): + params = GGMLLayoutParams( + scale=torch.ones((), dtype=torch.float32), + orig_dtype=orig_dtype, + orig_shape=tuple(tensor_shape), + tensor_type=tensor_type.value if not isinstance(tensor_type, int) else tensor_type, + ) + return QuantizedTensor(qdata, "GGMLLayout", params) diff --git a/requirements.txt b/requirements.txt index f49905c7..91b8e7a1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,7 @@ gguf>=0.13.0 # optional - tokenizer sentencepiece protobuf +# optional - comfy_kitchen W4A4 INT4 tensor-core backend for Q4_CR_W4A4. +# ComfyUI bundles comfy_kitchen; pin the floor that ships a working ConvRot W4A4 +# int4 MMA so Q4_CR_W4A4 loads natively on CUDA (older 0.2.26 crashes on Ampere). +comfy-kitchen>=0.2.27 diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/test_targeted_quantization.py b/tests/test_targeted_quantization.py new file mode 100644 index 00000000..b557bb25 --- /dev/null +++ b/tests/test_targeted_quantization.py @@ -0,0 +1,3073 @@ +import unittest +import json +import gc +import os +from contextlib import nullcontext +from collections import OrderedDict +import importlib.util +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest import mock + +import gguf +import torch +import comfy.sd +from safetensors.torch import save_file + +from tools.convert import ( + MEBIBYTE, + ModelLTXV, + ModelLTXVUpsampler, + ModelMinimaxH3, + ModelMinimaxH3VAE, + ModelMiniMaxMusic3DiT, + ModelMiniMaxMusic3TextEncoder, + ModelTemplate, + _streamed_safetensors_layout, + convert_file, + convert_state_dict, + detect_arch, + plan_target_size_quantization, + quantize_int8_convrot, + quantize_int4_cr_w4a4, + resolve_quantization_device, +) +from dequant import dequantize, dequantize_functions, dequantize_tensor +from lora import ( + fuse_targets_into_state_dict, + load_gguf_lora, + load_lora, + materialize_int8_source_weights, + resolve_fusion_targets, +) + + +def load_gguf_loader(): + loader_path = Path(__file__).parents[1] / "loader.py" + package_name = "comfyui_gguf_test" + spec = importlib.util.spec_from_file_location( + f"{package_name}.loader", + loader_path, + submodule_search_locations=[str(loader_path.parent)], + ) + module = importlib.util.module_from_spec(spec) + import sys + sys.modules[package_name] = module + sys.modules[f"{package_name}.loader"] = module + spec.loader.exec_module(module) + return module + + +def ops_factory(): + """Load the repo's ops.py and return get_gguf_q4_w4a4_ops.""" + import sys + ops_path = Path(__file__).parents[1] / "ops.py" + package_name = "comfyui_gguf_test" + spec = importlib.util.spec_from_file_location( + f"{package_name}.ops", + ops_path, + submodule_search_locations=[str(ops_path.parent)], + ) + module = importlib.util.module_from_spec(spec) + sys.modules[f"{package_name}.ops"] = module + spec.loader.exec_module(module) + return module.get_gguf_q4_w4a4_ops(torch.bfloat16) + + +def nodes_factory(): + """Load the repo's nodes.py without requiring ComfyUI to load this plugin.""" + import sys + + nodes_path = Path(__file__).parents[1] / "nodes.py" + repository_path = nodes_path.parent.resolve() + original_sys_path = sys.path[:] + sys.path = [path for path in sys.path if Path(path).resolve() != repository_path] + package_name = "comfyui_gguf_nodes_test" + try: + package = sys.modules.get(package_name) + if package is None: + package = importlib.util.module_from_spec( + importlib.util.spec_from_loader(package_name, loader=None) + ) + package.__path__ = [str(nodes_path.parent)] + sys.modules[package_name] = package + spec = importlib.util.spec_from_file_location( + f"{package_name}.nodes", + nodes_path, + submodule_search_locations=[str(nodes_path.parent)], + ) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + finally: + sys.path = original_sys_path + + +class TargetSizeQuantizationTests(unittest.TestCase): + def setUp(self): + self.model_arch = ModelTemplate() + self.state_dict = OrderedDict( + (f"blocks.{index}.weight", torch.ones((4096, 32), dtype=torch.float32)) + for index in range(3) + ) + self.state_dict["normalization.weight"] = torch.ones((4096,), dtype=torch.float32) + + def test_reduces_center_core_layers_before_outer_layers(self): + plan, _, selected_size = plan_target_size_quantization( + self.state_dict, self.model_arch, 0.4 + ) + + self.assertEqual(plan["blocks.1.weight"], gguf.GGMLQuantizationType.Q5_0) + self.assertEqual(plan["blocks.0.weight"], gguf.GGMLQuantizationType.I8) + self.assertEqual(plan["blocks.2.weight"], gguf.GGMLQuantizationType.I8) + self.assertLessEqual(selected_size, int(0.4 * MEBIBYTE)) + + def test_reduces_one_dimensional_weights_only_after_all_core_layers(self): + plan, _, selected_size = plan_target_size_quantization( + self.state_dict, self.model_arch, 0.222 + ) + + for index in range(3): + self.assertEqual(plan[f"blocks.{index}.weight"], gguf.GGMLQuantizationType.Q4_0) + self.assertEqual(plan["normalization.weight"], gguf.GGMLQuantizationType.BF16) + self.assertLessEqual(selected_size, int(0.222 * MEBIBYTE)) + + def test_supports_standard_q8_baseline(self): + plan, _, _ = plan_target_size_quantization( + self.state_dict, + self.model_arch, + 2.0, + target_size_q8_type="Q8_0", + ) + + for index in range(3): + self.assertEqual(plan[f"blocks.{index}.weight"], gguf.GGMLQuantizationType.Q8_0) + + def test_reports_minimum_when_target_is_unsupported(self): + with self.assertRaisesRegex(ValueError, "smallest supported TARGET_SIZE output"): + plan_target_size_quantization(self.state_dict, self.model_arch, 0.1) + + +class Q8CRConversionDeviceTests(unittest.TestCase): + def test_auto_uses_cpu_without_cuda(self): + with mock.patch("tools.convert.torch.cuda.is_available", return_value=False): + self.assertEqual(resolve_quantization_device("auto").type, "cpu") + + def test_cuda_requires_available_device(self): + with mock.patch("tools.convert.torch.cuda.is_available", return_value=False): + with self.assertRaisesRegex(RuntimeError, "requires an available CUDA device"): + resolve_quantization_device("cuda") + + def test_cpu_quantization_stays_on_cpu(self): + qdata, scale, quant_conf, _ = quantize_int8_convrot( + torch.arange(256, dtype=torch.float32).reshape(1, 256), + device=torch.device("cpu"), + ) + + self.assertEqual(qdata.device.type, "cpu") + self.assertEqual(scale.device.type, "cpu") + self.assertTrue(quant_conf["weight_rotated"]) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is not available") + def test_cuda_quantization_has_cpu_equivalent_decode_error(self): + torch.manual_seed(0) + weight = torch.randn((32, 256), dtype=torch.float32) + cpu_qdata, cpu_scale, _, _ = quantize_int8_convrot(weight, device=torch.device("cpu")) + cuda_qdata, cuda_scale, _, _ = quantize_int8_convrot(weight, device=torch.device("cuda")) + + cpu_decoded = cpu_qdata.to(torch.float32) * cpu_scale + cuda_decoded = cuda_qdata.cpu().to(torch.float32) * cuda_scale.cpu() + self.assertTrue(torch.allclose(cpu_decoded, cuda_decoded, atol=1e-5, rtol=0)) + + def test_auto_device_serializes_q8_cr_layout(self): + state_dict = { + "video_patch_proj.weight": torch.ones((32, 32), dtype=torch.float32), + "audio_patch_proj.weight": torch.ones((32, 32), dtype=torch.float32), + "blocks.0.attn.qkv_proj.weight": torch.ones((96, 32), dtype=torch.float32), + "final_layer.video_out.weight": torch.ones((96, 32), dtype=torch.float32), + } + + with TemporaryDirectory() as temp_dir: + source_path = Path(temp_dir) / "minimax_h3.safetensors" + output_path = Path(temp_dir) / "minimax_h3-Q8_CR.gguf" + save_file(state_dict, str(source_path)) + + converted_path, _ = convert_file( + str(source_path), + str(output_path), + interact=False, + quant_type_name="Q8_CR", + quantization_device="auto", + ) + + reader = gguf.GGUFReader(converted_path) + tensor_types = {tensor.name: tensor.tensor_type for tensor in reader.tensors} + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + del reader + + self.assertEqual( + tensor_types["blocks.0.attn.qkv_proj.weight"], + gguf.GGMLQuantizationType.I8, + ) + self.assertEqual( + tensor_types["blocks.0.attn.qkv_proj.weight_scale"], + gguf.GGMLQuantizationType.F32, + ) + + @unittest.skipUnless(hasattr(torch, "float8_e4m3fn"), "PyTorch does not support FP8") + def test_streamed_layout_supports_safetensors_f8_e4m3(self): + with TemporaryDirectory() as temp_dir: + source_path = Path(temp_dir) / "fp8.safetensors" + save_file( + {"weight": torch.ones((2, 2), dtype=torch.float8_e4m3fn)}, + str(source_path), + ) + state_dict, source_keys = _streamed_safetensors_layout(str(source_path)) + + self.assertEqual(state_dict["weight"].dtype, torch.float8_e4m3fn) + self.assertEqual(source_keys["weight"], "weight") + + +class MiniMaxH3VAEConversionTests(unittest.TestCase): + def test_detects_video_vae_from_distinctive_decoder_and_encoder_keys(self): + state_dict = { + "decoder.transformer_blocks.0.scale1": torch.ones(32), + "decoder.x_embedder.weight": torch.ones((64, 32)), + "encoder.down.5.block.0.conv1.weight": torch.ones((2, 2, 3, 3, 3)), + } + + self.assertIsInstance(detect_arch(state_dict), ModelMinimaxH3VAE) + + def test_q8_cr_keeps_conv3d_fp16_and_restores_its_shape(self): + state_dict = { + "decoder.transformer_blocks.0.scale1": torch.ones(32, dtype=torch.float16), + "decoder.x_embedder.weight": torch.ones((64, 32), dtype=torch.float16), + "encoder.down.5.block.0.conv1.weight": torch.ones( + (2, 2, 3, 3, 3), dtype=torch.float16 + ), + } + + with TemporaryDirectory() as temp_dir: + source_path = Path(temp_dir) / "minimax_h3_vae.safetensors" + output_path = Path(temp_dir) / "minimax_h3_vae-Q8_CR.gguf" + save_file(state_dict, str(source_path)) + + converted_path, _ = convert_file( + str(source_path), + str(output_path), + interact=False, + quant_type_name="Q8_CR", + quantization_device="cpu", + ) + reader = gguf.GGUFReader(converted_path) + tensor_types = {tensor.name: tensor.tensor_type for tensor in reader.tensors} + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + del reader + + loader = load_gguf_loader() + loaded, extra = loader.gguf_sd_loader(converted_path, handle_prefix=None) + conv3d = loaded["encoder.down.5.block.0.conv1.weight"] + del conv3d.tensor_type + materialized = dequantize_tensor(conv3d, dtype=torch.Tensor(conv3d).dtype) + conv3d_shape = tuple(materialized.shape) + del materialized + del conv3d + del loaded + gc.collect() + + self.assertEqual(extra["arch_str"], "minimax_h3_vae") + self.assertEqual( + tensor_types["decoder.x_embedder.weight"], gguf.GGMLQuantizationType.I8 + ) + self.assertEqual( + tensor_types["encoder.down.5.block.0.conv1.weight"], + gguf.GGMLQuantizationType.F16, + ) + self.assertEqual( + conv3d_shape, + (2, 2, 3, 3, 3), + ) + + +class MiniMaxMusic3ConversionTests(unittest.TestCase): + def test_detects_dit_from_music_specific_conditioning_and_attention_keys(self): + state_dict = { + "cond_layer_logits": torch.ones(8), + "latent_conditioners.0.weight": torch.ones((2, 2, 1)), + "diffusion_transformer.preprocess_conv.weight": torch.ones((2, 2, 1)), + "diffusion_transformer.postprocess_conv.weight": torch.ones((2, 2, 1)), + "diffusion_transformer.timestep_features.weight": torch.ones((128, 1)), + "diffusion_transformer.transformer.rotary_pos_emb.inv_freq": torch.ones(16), + "diffusion_transformer.transformer.project_in.weight": torch.ones((64, 64)), + "diffusion_transformer.transformer.layers.0.self_attn.to_qkv.weight": torch.ones((192, 64)), + } + + self.assertIsInstance(detect_arch(state_dict), ModelMiniMaxMusic3DiT) + + def test_detects_pruned_text_encoder_from_audio_and_qwen_keys(self): + state_dict = { + "model.embed_tokens_prefill.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.embed_tokens_audio.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.audio_extra_embedding.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.audio_decoder.pos_embedding.weight": torch.ones((16, 32), dtype=torch.bfloat16), + "model.lm_head_pruned.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.audio_decoder.audio_heads.0.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.layers.0.self_attn.qkv_proj.weight": torch.ones((96, 32), dtype=torch.bfloat16), + } + + self.assertIsInstance(detect_arch(state_dict), ModelMiniMaxMusic3TextEncoder) + + def test_q8_cr_preserves_dit_convolutions(self): + state_dict = { + "cond_layer_logits": torch.ones(8), + "latent_conditioners.0.weight": torch.ones((2, 2, 1)), + "diffusion_transformer.preprocess_conv.weight": torch.ones((2, 2, 1)), + "diffusion_transformer.postprocess_conv.weight": torch.ones((2, 2, 1)), + "diffusion_transformer.timestep_features.weight": torch.ones((128, 1)), + "diffusion_transformer.transformer.rotary_pos_emb.inv_freq": torch.ones(16), + "diffusion_transformer.transformer.project_in.weight": torch.ones((64, 64)), + "diffusion_transformer.transformer.layers.0.self_attn.to_qkv.weight": torch.ones((192, 64)), + } + with TemporaryDirectory() as temp_dir: + output_path = Path(temp_dir) / "minimax_music3_dit-Q8_CR.gguf" + converted_path, model_arch = convert_state_dict( + state_dict, + str(output_path), + quant_type_name="Q8_CR", + quantization_device="cpu", + ) + reader = gguf.GGUFReader(converted_path) + tensor_types = {tensor.name: tensor.tensor_type for tensor in reader.tensors} + loaded, extra = load_gguf_loader().gguf_sd_loader(converted_path, handle_prefix=None) + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + del reader + del loaded + + self.assertIsInstance(model_arch, ModelMiniMaxMusic3DiT) + self.assertEqual(extra["arch_str"], "minimax_music3") + self.assertEqual( + tensor_types["diffusion_transformer.transformer.project_in.weight"], + gguf.GGMLQuantizationType.I8, + ) + self.assertEqual( + tensor_types["latent_conditioners.0.weight"], + gguf.GGMLQuantizationType.F32, + ) + + def test_q8_cr_preserves_text_embeddings(self): + state_dict = { + "model.embed_tokens_prefill.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.embed_tokens_audio.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.audio_extra_embedding.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.audio_decoder.pos_embedding.weight": torch.ones((16, 32), dtype=torch.bfloat16), + "model.lm_head_pruned.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.audio_decoder.audio_heads.0.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.layers.0.self_attn.qkv_proj.weight": torch.ones((96, 32), dtype=torch.bfloat16), + } + with TemporaryDirectory() as temp_dir: + source_path = Path(temp_dir) / "minimax_music3_text_encoder.safetensors" + output_path = Path(temp_dir) / "minimax_music3_text_encoder-Q8_CR.gguf" + save_file(state_dict, str(source_path)) + converted_path, model_arch = convert_file( + str(source_path), + str(output_path), + interact=False, + quant_type_name="Q8_CR", + quantization_device="cpu", + streamed=True, + ) + reader = gguf.GGUFReader(converted_path) + tensor_types = {tensor.name: tensor.tensor_type for tensor in reader.tensors} + loaded, extra = load_gguf_loader().gguf_sd_loader( + converted_path, + handle_prefix=None, + is_text_model=True, + ) + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + del reader + del loaded + gc.collect() + + self.assertIsInstance(model_arch, ModelMiniMaxMusic3TextEncoder) + self.assertEqual(extra["arch_str"], "minimax_music3") + self.assertEqual( + tensor_types["model.embed_tokens_prefill.weight"], + gguf.GGMLQuantizationType.BF16, + ) + self.assertEqual( + tensor_types["model.layers.0.self_attn.qkv_proj.weight"], + gguf.GGMLQuantizationType.I8, + ) + + def test_tokenizer_json_roundtrips_as_raw_bytes_for_all_quantizers(self): + tokenizer_json = b'{"text":"caf\xc3\xb6"}' + state_dict = { + "model.embed_tokens_prefill.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.embed_tokens_audio.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.audio_extra_embedding.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.audio_decoder.pos_embedding.weight": torch.ones((16, 32), dtype=torch.bfloat16), + "model.lm_head_pruned.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.audio_decoder.audio_heads.0.weight": torch.ones((64, 32), dtype=torch.bfloat16), + "model.layers.0.self_attn.qkv_proj.weight": torch.ones((96, 32), dtype=torch.bfloat16), + "tokenizer_json": torch.tensor(list(tokenizer_json), dtype=torch.uint8), + } + model_arch = detect_arch(state_dict) + plan, _, _ = plan_target_size_quantization( + state_dict, + model_arch, + max_size_mb=1, + ) + self.assertEqual(plan["tokenizer_json"], gguf.GGMLQuantizationType.I8) + + loader = load_gguf_loader() + for quant_type in ("Q4_0", "Q8_CR"): + with self.subTest(quant_type=quant_type), TemporaryDirectory() as temp_dir: + output_path = Path(temp_dir) / f"minimax_music3_text_encoder-{quant_type}.gguf" + converted_path, _ = convert_state_dict( + state_dict, + str(output_path), + quant_type_name=quant_type, + quantization_device="cpu", + ) + reader = gguf.GGUFReader(converted_path) + tensor_types = {tensor.name: tensor.tensor_type for tensor in reader.tensors} + self.assertEqual( + tensor_types["tokenizer_json"], + gguf.GGMLQuantizationType.I8, + ) + loaded = loader.gguf_clip_loader(converted_path) + self.assertEqual(loaded["tokenizer_json"].dtype, torch.uint8) + self.assertEqual(loaded["tokenizer_json"].tolist(), list(tokenizer_json)) + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + del reader + del loaded + gc.collect() + + def test_clip_loader_restores_legacy_float_tokenizer_bytes(self): + loader = load_gguf_loader() + tokenizer_json = b'{"text":"caf\xc3\xb6"}' + legacy = torch.tensor(list(tokenizer_json), dtype=torch.float32) + with mock.patch.object( + loader, + "gguf_sd_loader", + return_value=( + {"tokenizer_json": loader.GGMLTensor( + legacy, + tensor_type=gguf.GGMLQuantizationType.F32, + tensor_shape=legacy.shape, + )}, + {"arch_str": "minimax_music3"}, + ), + ): + loaded = loader.gguf_clip_loader("legacy-minimax-music3.gguf") + + self.assertEqual(loaded["tokenizer_json"].dtype, torch.uint8) + self.assertEqual(loaded["tokenizer_json"].tolist(), list(tokenizer_json)) + + +class LTX25ConversionTests(unittest.TestCase): + def test_detects_ltx25_audio_video_transformer(self): + state_dict = { + "adaln_single.emb.timestep_embedder.linear_2.weight": torch.ones((32, 16)), + "audio_adaln_single.linear.weight": torch.ones((32, 16)), + "transformer_blocks.27.scale_shift_table": torch.ones((6, 32)), + } + + self.assertIsInstance(detect_arch(state_dict), ModelLTXV) + + def test_converts_temporal_upscaler_without_quantizing_conv3d(self): + state_dict = { + "initial_conv.weight": torch.ones((32, 16, 3, 3, 3), dtype=torch.bfloat16), + "post_upsample_res_blocks.0.conv2.bias": torch.ones(32, dtype=torch.bfloat16), + "upsampler.0.weight": torch.ones((64, 32, 3, 3, 3), dtype=torch.bfloat16), + "final_conv.weight": torch.ones((16, 32, 3, 3, 3), dtype=torch.bfloat16), + } + metadata = { + "config": json.dumps( + { + "_class_name": "LatentUpsampler", + "in_channels": 16, + "mid_channels": 32, + "num_blocks_per_stage": 1, + "dims": 3, + "spatial_upsample": False, + "temporal_upsample": True, + } + ) + } + + with TemporaryDirectory() as temp_dir: + output_path = Path(temp_dir) / "ltx25-temporal-upscaler-Q8_CR.gguf" + converted_path, model_arch = convert_state_dict( + state_dict, + str(output_path), + source_metadata=metadata, + quant_type_name="Q8_CR", + quantization_device="cpu", + ) + reader = gguf.GGUFReader(converted_path) + tensor_types = {tensor.name: tensor.tensor_type for tensor in reader.tensors} + loader = load_gguf_loader() + loaded, extra = loader.gguf_sd_loader(converted_path, handle_prefix=None) + restored_shape = tuple( + loaded["upsampler.0.weight"].shape + ) + loaded_arch = extra["arch_str"] + config_dims = json.loads(extra["metadata"]["config"])["dims"] + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + del reader + del loaded + del extra + gc.collect() + + self.assertIsInstance(model_arch, ModelLTXVUpsampler) + self.assertEqual(loaded_arch, "ltxv_upscaler") + self.assertEqual( + tensor_types["upsampler.0.weight"], gguf.GGMLQuantizationType.BF16 + ) + self.assertEqual(restored_shape, (64, 32, 3, 3, 3)) + self.assertEqual(config_dims, 3) + + +class GGUFLoraTests(unittest.TestCase): + def _write_lora(self, path, down, up, alpha=4.0): + writer = gguf.GGUFWriter(path=None, arch="minimax_h3") + writer.add_string("general.type", "adapter") + writer.add_string("adapter.type", "lora") + writer.add_float32("adapter.lora.alpha", alpha) + writer.add_tensor( + "blocks.0.attn.qkv_proj.weight.lora_a", + down.numpy(), + raw_dtype=gguf.GGMLQuantizationType.F32, + ) + writer.add_tensor( + "blocks.0.attn.qkv_proj.weight.lora_b", + up.numpy(), + raw_dtype=gguf.GGMLQuantizationType.F32, + ) + writer.write_header_to_file(path=str(path)) + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + def test_imports_standard_gguf_factor_pair(self): + down = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + up = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]) + with TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "adapter.gguf" + self._write_lora(path, down, up) + lora, targets, metadata = load_gguf_lora(path) + + self.assertEqual(metadata["alpha"], 4.0) + self.assertIn("blocks.0.attn.qkv_proj.lora_A.weight", lora) + self.assertIn("blocks.0.attn.qkv_proj.lora_B.weight", lora) + self.assertTrue(torch.equal(targets["blocks.0.attn.qkv_proj"]["down"], down)) + self.assertTrue(torch.equal(targets["blocks.0.attn.qkv_proj"]["up"], up)) + + def test_fuses_lora_delta_in_selected_precision(self): + state_dict = { + "blocks.0.attn.qkv_proj.weight": torch.zeros((3, 2), dtype=torch.float16) + } + targets = { + "blocks.0.attn.qkv_proj": { + "base_name": "blocks.0.attn.qkv_proj.weight", + "down": torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + "up": torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]), + "alpha": 4.0, + } + } + + count = fuse_targets_into_state_dict( + state_dict, targets, strength=0.5, device=torch.device("cpu") + ) + + expected = torch.tensor([[1.0, 2.0], [3.0, 4.0], [4.0, 6.0]], dtype=torch.float16) + self.assertEqual(count, 1) + self.assertTrue(torch.equal(state_dict["blocks.0.attn.qkv_proj.weight"], expected)) + + def test_restores_per_row_scaled_int8_source_weights(self): + state_dict = { + "blocks.0.attn.qkv_proj.weight": torch.tensor( + [[10, -20], [30, -40]], dtype=torch.int8 + ), + "blocks.0.attn.qkv_proj.weight_scale": torch.tensor([0.1, 0.01]), + } + + restored_count = materialize_int8_source_weights(state_dict) + + self.assertEqual(restored_count, 1) + self.assertNotIn("blocks.0.attn.qkv_proj.weight_scale", state_dict) + self.assertEqual(state_dict["blocks.0.attn.qkv_proj.weight"].dtype, torch.float16) + self.assertTrue( + torch.equal( + state_dict["blocks.0.attn.qkv_proj.weight"], + torch.tensor([[1.0, -2.0], [0.3, -0.4]], dtype=torch.float16), + ) + ) + + def test_restores_convrot_int8_source_weights(self): + source = torch.tensor( + [[2.0, 2.0, 2.0, 2.0], [2.0, -2.0, 2.0, -2.0]], + dtype=torch.float32, + ) + qdata, scale, quant_conf, _ = quantize_int8_convrot( + source, convrot_groupsize=4, device=torch.device("cpu") + ) + state_dict = { + "blocks.0.proj.weight": qdata, + "blocks.0.proj.weight_scale": scale, + "blocks.0.proj.comfy_quant": torch.tensor( + list(json.dumps(quant_conf).encode("utf-8")), dtype=torch.uint8 + ), + } + + materialize_int8_source_weights(state_dict) + + self.assertNotIn("blocks.0.proj.weight_scale", state_dict) + self.assertNotIn("blocks.0.proj.comfy_quant", state_dict) + self.assertTrue( + torch.equal( + state_dict["blocks.0.proj.weight"], + source.to(dtype=torch.float16), + ) + ) + + def test_rejects_int8_source_weight_without_scale(self): + state_dict = { + "blocks.0.attn.qkv_proj.weight": torch.ones((2, 2), dtype=torch.int8), + } + + with self.assertRaisesRegex(ValueError, "missing its scale tensor"): + materialize_int8_source_weights(state_dict) + + def test_warns_and_skips_unmatched_lora_targets(self): + targets = { + "missing.layer": { + "base_name": "missing.layer.weight", + "down": torch.ones((1, 2)), + "up": torch.ones((2, 1)), + "alpha": 1.0, + } + } + + with self.assertLogs(level="WARNING") as logs: + resolved = resolve_fusion_targets({}, targets) + + self.assertEqual(resolved, {}) + self.assertIn("Skipping 1 LoRA target", logs.output[0]) + + +class OfflineLoraFusionTests(unittest.TestCase): + def test_imports_and_fuses_safetensors_factor_pair(self): + down = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + up = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]) + with TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "adapter.safetensors" + save_file( + { + "blocks.0.attn.qkv_proj.lora_A.weight": down, + "blocks.0.attn.qkv_proj.lora_B.weight": up, + "blocks.0.attn.qkv_proj.alpha": torch.tensor(4.0), + }, + str(path), + ) + _, targets, _ = load_lora(path) + + state_dict = { + "blocks.0.attn.qkv_proj.weight": torch.zeros((3, 2), dtype=torch.float32) + } + fuse_targets_into_state_dict(state_dict, targets, strength=0.5, device=torch.device("cpu")) + expected = torch.tensor([[1.0, 2.0], [3.0, 4.0], [4.0, 6.0]]) + self.assertTrue(torch.equal(state_dict["blocks.0.attn.qkv_proj.weight"], expected)) + + @mock.patch( + "lora._comfy_model_lora_target_map", + return_value={ + "transformer.single_transformer_blocks.0.attn.to_qkv_mlp_proj": ( + "diffusion_model.single_blocks.0.linear1.weight" + ) + }, + ) + def test_fuses_comfy_mapped_flux2_single_block_target(self, _): + state_dict = { + "double_blocks.0.img_attn.proj.weight": torch.zeros((2, 2)), + "single_blocks.0.linear1.weight": torch.zeros((2, 2)), + } + targets = { + "transformer.single_transformer_blocks.0.attn.to_qkv_mlp_proj": { + "base_name": ( + "transformer.single_transformer_blocks.0.attn." + "to_qkv_mlp_proj.weight" + ), + "down": torch.eye(2), + "up": torch.eye(2), + "alpha": 2.0, + } + } + + count = fuse_targets_into_state_dict( + state_dict, targets, strength=1.0, device=torch.device("cpu") + ) + + self.assertEqual(count, 1) + self.assertTrue(torch.equal(state_dict["single_blocks.0.linear1.weight"], torch.eye(2))) + + @mock.patch( + "lora._comfy_model_lora_target_map", + return_value={ + "transformer.transformer_blocks.0.attn.to_q": ( + "diffusion_model.double_blocks.0.img_attn.qkv.weight", + (0, 0, 2), + ), + "transformer.transformer_blocks.0.attn.to_v": ( + "diffusion_model.double_blocks.0.img_attn.qkv.weight", + (0, 4, 2), + ), + }, + ) + def test_fuses_comfy_mapped_qkv_slices(self, _): + state_dict = { + "double_blocks.0.img_attn.proj.weight": torch.zeros((2, 2)), + "single_blocks.0.linear1.weight": torch.zeros((2, 2)), + "double_blocks.0.img_attn.qkv.weight": torch.zeros((6, 2)), + } + targets = { + "transformer.transformer_blocks.0.attn.to_q": { + "base_name": "transformer.transformer_blocks.0.attn.to_q.weight", + "down": torch.eye(2), + "up": torch.eye(2), + "alpha": 2.0, + }, + "transformer.transformer_blocks.0.attn.to_v": { + "base_name": "transformer.transformer_blocks.0.attn.to_v.weight", + "down": torch.eye(2), + "up": torch.eye(2), + "alpha": 2.0, + }, + } + + count = fuse_targets_into_state_dict( + state_dict, targets, strength=1.0, device=torch.device("cpu") + ) + + self.assertEqual(count, 2) + self.assertTrue( + torch.equal( + state_dict["double_blocks.0.img_attn.qkv.weight"], + torch.tensor( + [[1.0, 0.0], [0.0, 1.0], [0.0, 0.0], [0.0, 0.0], [1.0, 0.0], [0.0, 1.0]] + ), + ) + ) + + @mock.patch( + "lora._comfy_model_lora_target_map", + return_value={ + "transformer.transformer_blocks.0.attn.to_q": ( + "diffusion_model.blocks.0.attn.wq.weight" + ) + }, + ) + def test_fuses_comfy_mapped_krea2_target(self, _): + state_dict = { + "blocks.0.attn.wq.weight": torch.zeros((2, 2)), + } + targets = { + "transformer.transformer_blocks.0.attn.to_q": { + "base_name": "transformer.transformer_blocks.0.attn.to_q.weight", + "down": torch.eye(2), + "up": torch.eye(2), + "alpha": 2.0, + } + } + + count = fuse_targets_into_state_dict( + state_dict, targets, strength=1.0, device=torch.device("cpu") + ) + + self.assertEqual(count, 1) + self.assertTrue(torch.equal(state_dict["blocks.0.attn.wq.weight"], torch.eye(2))) + + def test_imports_and_fuses_direct_lokr_adapter(self): + with TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "adapter.safetensors" + save_file( + { + "blocks.0.proj.lokr_w1": torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + "blocks.0.proj.lokr_w2": torch.tensor([[1.0, 0.0], [0.0, 1.0]]), + "blocks.0.proj.alpha": torch.tensor(100.0), + }, + str(path), + ) + _, targets, metadata = load_lora(path) + + state_dict = {"blocks.0.proj.weight": torch.zeros((4, 4))} + count = fuse_targets_into_state_dict( + state_dict, targets, strength=0.5, device=torch.device("cpu") + ) + + self.assertEqual(metadata["target_count"], 1) + self.assertEqual(count, 1) + self.assertTrue( + torch.equal( + state_dict["blocks.0.proj.weight"], + 0.5 * torch.kron( + torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + torch.eye(2), + ), + ) + ) + + @unittest.skipUnless(hasattr(torch, "float8_e4m3fn"), "PyTorch does not support FP8") + def test_fuses_scaled_fp8_target_as_fp16(self): + state_dict = { + "blocks.0.attn.qkv_proj.weight": torch.ones( + (2, 2), dtype=torch.float8_e4m3fn + ), + "blocks.0.attn.qkv_proj.weight_scale": torch.tensor(0.5), + } + targets = { + "blocks.0.attn.qkv_proj": { + "base_name": "blocks.0.attn.qkv_proj.weight", + "down": torch.eye(2), + "up": torch.eye(2), + "alpha": 2.0, + } + } + + fuse_targets_into_state_dict(state_dict, targets, strength=1.0, device=torch.device("cpu")) + + self.assertEqual(state_dict["blocks.0.attn.qkv_proj.weight"].dtype, torch.float16) + self.assertTrue( + torch.equal( + state_dict["blocks.0.attn.qkv_proj.weight"], + torch.tensor([[1.5, 0.5], [0.5, 1.5]], dtype=torch.float16), + ) + ) + + def test_converter_merges_safetensors_lora_before_gguf_export(self): + source = { + "video_patch_proj.weight": torch.zeros((32, 32), dtype=torch.float16), + "audio_patch_proj.weight": torch.zeros((32, 32), dtype=torch.float16), + "blocks.0.attn.qkv_proj.weight": torch.zeros((96, 32), dtype=torch.float16), + "final_layer.video_out.weight": torch.zeros((96, 32), dtype=torch.float16), + } + down = torch.zeros((2, 32), dtype=torch.float16) + down[:, 0] = 1 + up = torch.zeros((96, 2), dtype=torch.float16) + up[0] = 1 + + with TemporaryDirectory() as temp_dir: + source_path = Path(temp_dir) / "minimax_h3.safetensors" + lora_path = Path(temp_dir) / "adapter.safetensors" + output_path = Path(temp_dir) / "merged.gguf" + save_file(source, str(source_path)) + save_file( + { + "blocks.0.attn.qkv_proj.lora_A.weight": down, + "blocks.0.attn.qkv_proj.lora_B.weight": up, + "blocks.0.attn.qkv_proj.alpha": torch.tensor(2.0), + }, + str(lora_path), + ) + convert_file( + str(source_path), + str(output_path), + interact=False, + quant_type_name="F16", + lora_paths=[str(lora_path)], + ) + reader = gguf.GGUFReader(str(output_path)) + tensor = next( + tensor for tensor in reader.tensors + if tensor.name == "blocks.0.attn.qkv_proj.weight" + ) + merged = torch.from_numpy(tensor.data.copy()).view(torch.float16).reshape( + tuple(reversed(tensor.shape)) + ) + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + + self.assertTrue(torch.equal(merged[0, :2], torch.tensor([2.0, 0.0], dtype=torch.float16))) + + def test_streamed_converter_merges_lora_without_loading_state_dict(self): + source = { + "video_patch_proj.weight": torch.zeros((32, 32), dtype=torch.float16), + "audio_patch_proj.weight": torch.zeros((32, 32), dtype=torch.float16), + "blocks.0.attn.qkv_proj.weight": torch.zeros((96, 32), dtype=torch.float16), + "final_layer.video_out.weight": torch.zeros((96, 32), dtype=torch.float16), + } + down = torch.zeros((2, 32), dtype=torch.float16) + down[:, 0] = 1 + up = torch.zeros((96, 2), dtype=torch.float16) + up[0] = 1 + + with TemporaryDirectory() as temp_dir: + source_path = Path(temp_dir) / "minimax_h3.safetensors" + lora_path = Path(temp_dir) / "adapter.safetensors" + output_path = Path(temp_dir) / "merged.gguf" + save_file(source, str(source_path)) + save_file( + { + "blocks.0.attn.qkv_proj.lora_A.weight": down, + "blocks.0.attn.qkv_proj.lora_B.weight": up, + "blocks.0.attn.qkv_proj.alpha": torch.tensor(2.0), + }, + str(lora_path), + ) + with mock.patch("tools.convert.load_state_dict") as load_state_dict: + convert_file( + str(source_path), + str(output_path), + interact=False, + quant_type_name="F16", + lora_paths=[str(lora_path)], + streamed=True, + ) + load_state_dict.assert_not_called() + reader = gguf.GGUFReader(str(output_path)) + tensor = next( + tensor for tensor in reader.tensors + if tensor.name == "blocks.0.attn.qkv_proj.weight" + ) + merged = torch.from_numpy(tensor.data.copy()).view(torch.float16).reshape( + tuple(reversed(tensor.shape)) + ) + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + + self.assertTrue(torch.equal(merged[0, :2], torch.tensor([2.0, 0.0], dtype=torch.float16))) + + +class Qwen3VLDetectionMarkerTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.loader = load_gguf_loader() + + def test_uses_minimax_32b_detection_marker_for_5120_hidden_size(self): + state_dict = { + "model.layers.0.input_layernorm.weight": torch.zeros(5120), + "model.layers.49.self_attn.q_proj.weight": torch.zeros(1), + } + + self.loader.inject_qwen3vl_detection_markers(state_dict) + + self.assertEqual( + comfy.sd.detect_te_model(state_dict), + comfy.sd.TEModel.QWEN3VL_32B, + ) + self.assertIn("visual.deepstack_merger_list.0.norm.weight", state_dict) + self.assertNotIn("model.visual.deepstack_merger_list.0.norm.weight", state_dict) + self.assertNotIn("model.visual.merger.linear_fc2.weight", state_dict) + self.assertEqual( + state_dict["visual.deepstack_merger_list.0.norm.weight"].shape, + (4608,), + ) + + def test_uses_model_prefixed_detection_markers_for_8b(self): + state_dict = { + "model.layers.0.input_layernorm.weight": torch.zeros(4096), + } + + self.loader.inject_qwen3vl_detection_markers(state_dict) + + self.assertIn("model.visual.deepstack_merger_list.0.norm.weight", state_dict) + self.assertEqual( + state_dict["model.visual.merger.linear_fc2.weight"].shape, + (4096, 4608), + ) + + def test_pruned_32b_clip_loader_injects_marker(self): + with mock.patch.object( + self.loader, + "gguf_sd_loader", + return_value=( + { + "model.layers.0.input_layernorm.weight": torch.zeros(5120), + "model.layers.49.self_attn.q_proj.weight": torch.zeros(1), + }, + {"arch_str": "qwen3vl"}, + ), + ), mock.patch.object( + self.loader, + "gguf_mmproj_loader", + return_value={}, + ): + state_dict = self.loader.gguf_clip_loader("qwen3-vl-pruned.gguf") + + self.assertEqual( + state_dict["visual.deepstack_merger_list.0.norm.weight"].shape, + (4608,), + ) + + def test_pruned_32b_clip_loader_maps_matching_mmproj_to_visual_tower(self): + with mock.patch.object( + self.loader, + "gguf_sd_loader", + return_value=( + { + "model.layers.0.input_layernorm.weight": torch.zeros(5120), + "model.layers.49.self_attn.q_proj.weight": torch.zeros(1), + }, + {"arch_str": "qwen3vl"}, + ), + ), mock.patch.object( + self.loader, + "gguf_mmproj_loader", + return_value={ + "model.visual.deepstack_merger_list.0.norm.weight": torch.ones(4608), + }, + ): + state_dict = self.loader.gguf_clip_loader("qwen3-vl-pruned.gguf") + + self.assertIn("visual.deepstack_merger_list.0.norm.weight", state_dict) + self.assertNotIn("model.visual.deepstack_merger_list.0.norm.weight", state_dict) + self.assertTrue( + torch.equal( + state_dict["visual.deepstack_merger_list.0.norm.weight"], + torch.ones(4608), + ) + ) + + def test_maps_qwen3vl_mmproj_deepstack_tensors(self): + mapped = self.loader.sd_map_replace( + { + "v.deepstack.0.norm.weight": torch.ones(4608), + "v.deepstast.1.fc2.weight": torch.ones(5120, 4608), + "v.blk.0.attn_qkv.weight": torch.ones(3, 3), + }, + self.loader.CLIP_VISION_QWEN3_MAP, + ) + + self.assertIn( + "model.visual.deepstack_merger_list.0.norm.weight", + mapped, + ) + self.assertIn( + "model.visual.deepstack_merger_list.1.linear_fc2.weight", + mapped, + ) + self.assertIn("model.visual.blocks.0.attn.qkv.weight", mapped) + + def test_qwen3vl_mmproj_stacks_temporal_patch_embeddings(self): + with TemporaryDirectory() as temp_dir: + text_encoder = Path(temp_dir) / "qwen3-vl-IQ3_XXS.gguf" + mmproj = Path(temp_dir) / "qwen3-vl-mmproj-BF16.gguf" + text_encoder.touch() + mmproj.touch() + patch_a = torch.ones((2, 3, 2, 2)) + patch_b = torch.full((2, 3, 2, 2), 2.0) + + with mock.patch.object( + self.loader, + "gguf_sd_loader", + return_value=( + { + "v.patch_embd.weight": patch_a, + "v.patch_embd.weight.1": patch_b, + "v.deepstast.0.norm.weight": torch.ones(8), + }, + {}, + ), + ): + mapped = self.loader.gguf_mmproj_loader(str(text_encoder)) + + weight = mapped["model.visual.patch_embed.proj.weight"] + self.assertEqual(weight.shape, (2, 3, 2, 2, 2)) + self.assertTrue(torch.equal(weight[:, :, 0], patch_a)) + self.assertTrue(torch.equal(weight[:, :, 1], patch_b)) + + +class Qwen35GGUFLoaderTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.loader = load_gguf_loader() + + def test_qwen35_arch_is_whitelisted(self): + self.assertIn("qwen35", self.loader.TXT_ARCH_LIST) + + def test_maps_qwen35_tensor_layout_to_comfyui(self): + state_dict = { + "token_embd.weight": torch.zeros(1), + "output_norm.weight": torch.zeros(1), + "output.weight": torch.zeros(1), + "blk.0.attn_norm.weight": torch.zeros(1), + "blk.0.post_attention_norm.weight": torch.zeros(1), + "blk.0.attn_qkv.weight": torch.zeros(1), + "blk.0.attn_gate.weight": torch.zeros(1), + "blk.0.ssm_a": torch.zeros(1), + "blk.0.ssm_dt.bias": torch.zeros(1), + "blk.0.ssm_alpha.weight": torch.zeros(1), + "blk.0.ssm_beta.weight": torch.zeros(1), + "blk.0.ssm_conv1d.weight": torch.zeros(1), + "blk.0.ssm_norm.weight": torch.zeros(1), + "blk.0.ssm_out.weight": torch.zeros(1), + "blk.3.attn_q.weight": torch.zeros(1), + "blk.3.attn_k.weight": torch.zeros(1), + "blk.3.attn_v.weight": torch.zeros(1), + "blk.3.attn_output.weight": torch.zeros(1), + "blk.3.attn_q_norm.weight": torch.zeros(1), + "blk.3.attn_k_norm.weight": torch.zeros(1), + "blk.3.ffn_up.weight": torch.zeros(1), + "blk.3.ffn_gate.weight": torch.zeros(1), + "blk.3.ffn_down.weight": torch.zeros(1), + } + + mapped = self.loader.sd_map_replace(state_dict, self.loader.QWEN35_SD_MAP) + + self.assertIn("model.language_model.embed_tokens.weight", mapped) + self.assertIn("model.language_model.norm.weight", mapped) + self.assertIn("lm_head.weight", mapped) + self.assertIn("model.language_model.layers.0.input_layernorm.weight", mapped) + self.assertIn( + "model.language_model.layers.0.post_attention_layernorm.weight", + mapped, + ) + self.assertIn("model.language_model.layers.0.linear_attn.A_log", mapped) + self.assertIn( + "model.language_model.layers.0.linear_attn.in_proj_qkv.weight", mapped + ) + self.assertIn( + "model.language_model.layers.0.linear_attn.in_proj_z.weight", mapped + ) + self.assertIn( + "model.language_model.layers.0.linear_attn.in_proj_a.weight", mapped + ) + self.assertIn( + "model.language_model.layers.0.linear_attn.in_proj_b.weight", mapped + ) + self.assertIn( + "model.language_model.layers.0.linear_attn.dt_bias", mapped + ) + self.assertIn( + "model.language_model.layers.0.linear_attn.conv1d.weight", mapped + ) + self.assertIn( + "model.language_model.layers.0.linear_attn.norm.weight", mapped + ) + self.assertIn( + "model.language_model.layers.0.linear_attn.out_proj.weight", mapped + ) + self.assertIn("model.language_model.layers.3.self_attn.q_proj.weight", mapped) + self.assertIn("model.language_model.layers.3.self_attn.k_proj.weight", mapped) + self.assertIn("model.language_model.layers.3.self_attn.v_proj.weight", mapped) + self.assertIn("model.language_model.layers.3.self_attn.o_proj.weight", mapped) + self.assertIn("model.language_model.layers.3.self_attn.q_norm.weight", mapped) + self.assertIn("model.language_model.layers.3.self_attn.k_norm.weight", mapped) + self.assertIn("model.language_model.layers.3.mlp.up_proj.weight", mapped) + self.assertIn("model.language_model.layers.3.mlp.gate_proj.weight", mapped) + self.assertIn("model.language_model.layers.3.mlp.down_proj.weight", mapped) + + def test_qwen35_layout_is_detected_by_installed_comfyui(self): + for hidden_size, expected in ( + (1024, comfy.sd.TEModel.QWEN35_08B), + (2048, comfy.sd.TEModel.QWEN35_2B), + (2560, comfy.sd.TEModel.QWEN35_4B), + (4096, comfy.sd.TEModel.QWEN35_9B), + (5120, comfy.sd.TEModel.QWEN35_27B), + ): + with self.subTest(hidden_size=hidden_size): + state_dict = { + "model.language_model.layers.0.linear_attn.A_log": torch.zeros(1), + "model.language_model.layers.0.input_layernorm.weight": torch.zeros(hidden_size), + } + self.assertEqual( + comfy.sd.detect_te_model(state_dict), + expected, + ) + + def test_clip_loader_corrects_norms_and_alog(self): + norm_stored = torch.full((2560,), 2.0) # llama.cpp stores w + 1 + alog_stored = torch.full((32,), -2.0) # llama.cpp stores -exp(A_log) + + with mock.patch.object( + self.loader, + "gguf_sd_loader", + return_value=( + { + "token_embd.weight": torch.zeros((248320, 2560)), + "output_norm.weight": norm_stored.clone(), + "blk.0.attn_norm.weight": norm_stored.clone(), + "blk.0.post_attention_norm.weight": norm_stored.clone(), + "blk.0.ssm_norm.weight": torch.ones(128), + "blk.0.ssm_a": alog_stored.clone(), + "blk.0.attn_qkv.weight": torch.zeros((8192, 2560)), + "blk.3.attn_q_norm.weight": norm_stored.clone(), + "blk.3.attn_k_norm.weight": norm_stored.clone(), + }, + {"arch_str": "qwen35"}, + ), + ), mock.patch.object( + self.loader, + "gguf_mmproj_loader", + return_value={}, + ): + state_dict = self.loader.gguf_clip_loader("Qwen3.5-4B-BF16.gguf") + + expected_log = torch.log(torch.tensor(2.0)) + self.assertTrue( + torch.equal( + state_dict["model.language_model.norm.weight"], + torch.ones(2560), + ) + ) + self.assertTrue( + torch.equal( + state_dict["model.language_model.layers.0.input_layernorm.weight"], + torch.ones(2560), + ) + ) + self.assertTrue( + torch.equal( + state_dict["model.language_model.layers.0.post_attention_layernorm.weight"], + torch.ones(2560), + ) + ) + self.assertTrue( + torch.equal( + state_dict["model.language_model.layers.3.self_attn.q_norm.weight"], + torch.ones(2560), + ) + ) + self.assertTrue( + torch.equal( + state_dict["model.language_model.layers.3.self_attn.k_norm.weight"], + torch.ones(2560), + ) + ) + # linear_attn.norm is RMSNormGated and must NOT be shifted. + self.assertTrue( + torch.equal( + state_dict["model.language_model.layers.0.linear_attn.norm.weight"], + torch.ones(128), + ) + ) + # A_log is inverted back from -exp(A_log). + self.assertTrue( + torch.allclose( + state_dict["model.language_model.layers.0.linear_attn.A_log"], + torch.full((32,), expected_log), + ) + ) + + def test_conv1d_kernel_is_unsqueezed_to_depthwise_shape(self): + # V channels stored tiled so the corrected kernel equals arange() + tiled = [i * 2 for i in range(16)] + [i * 2 + 1 for i in range(16)] + with mock.patch.object( + self.loader, + "gguf_sd_loader", + return_value=( + { + "blk.0.attn_qkv.weight": torch.zeros((8192, 2560)), + "blk.0.attn_gate.weight": torch.zeros((4096, 2560)), + "blk.0.ssm_a": torch.full((32,), -1.0), + "blk.0.ssm_conv1d.weight": torch.cat( + [ + torch.arange(4096).unsqueeze(1).repeat(1, 4), + (torch.arange(4096) + 4096).reshape(32, 128)[tiled] + .reshape(-1) + .unsqueeze(1) + .repeat(1, 4), + ], + dim=0, + ).to(torch.float32), + }, + {"arch_str": "qwen35"}, + ), + ), mock.patch.object( + self.loader, + "gguf_mmproj_loader", + return_value={}, + ): + state_dict = self.loader.gguf_clip_loader("Qwen3.5-4B-BF16.gguf") + + conv = state_dict["model.language_model.layers.0.linear_attn.conv1d.weight"] + self.assertEqual(conv.shape, (8192, 1, 4)) + self.assertTrue( + torch.equal( + conv[:, 0, :], + torch.arange(8192).unsqueeze(1).repeat(1, 4), + ) + ) + + def test_reorders_tiled_v_heads_back_to_grouped_order(self): + # llama.cpp stores V heads tiled for k=2, v=4: [K0_v0, K1_v0, K0_v1, K1_v1] + # i.e. head order [h0, h2, h1, h3]; ComfyUI expects grouped [h0, h1, h2, h3]. + num_k_heads, num_v_heads, head_dim = 2, 4, 1 + value_dim = num_v_heads * head_dim + key_dim = num_k_heads * head_dim + conv_dim = 2 * key_dim + value_dim + tiled = [0, 2, 1, 3] + + def head_marked(rows, cols=1): + return torch.arange(rows, dtype=torch.float32).unsqueeze(1).repeat(1, cols) + + # stored rows/cols encode their V-head index in TILED order + v_stored = head_marked(value_dim, 3)[tiled] + prefix = "model.language_model.layers.0.linear_attn." + qkv = torch.cat( + [head_marked(key_dim, 3), head_marked(key_dim, 3), v_stored], dim=0 + ) + conv = torch.cat( + [ + torch.arange(2 * key_dim, dtype=torch.float32).unsqueeze(1).repeat(1, 2), + (torch.arange(value_dim, dtype=torch.float32) + 2 * key_dim) + [tiled].unsqueeze(1).repeat(1, 2), + ], + dim=0, + ) + sd = { + prefix + "in_proj_qkv.weight": qkv, + prefix + "in_proj_z.weight": v_stored, + prefix + "in_proj_a.weight": head_marked(num_v_heads)[tiled], + prefix + "in_proj_b.weight": head_marked(num_v_heads)[tiled], + prefix + "A_log": torch.full((num_v_heads,), -1.0), + prefix + "dt_bias": torch.arange(num_v_heads, dtype=torch.float32)[tiled], + prefix + "conv1d.weight": conv, + prefix + "out_proj.weight": head_marked(3, value_dim)[:, tiled], + "model.language_model.layers.0.input_layernorm.weight": torch.full((4,), 2.0), + } + + corrected = self.loader.qwen35_corrections(sd) + + qkv = corrected[prefix + "in_proj_qkv.weight"] + self.assertTrue(torch.equal(qkv[: 2 * key_dim], head_marked(key_dim, 3).repeat(2, 1))) + self.assertTrue(torch.equal(qkv[2 * key_dim:], head_marked(value_dim, 3))) + self.assertTrue(torch.equal(corrected[prefix + "in_proj_z.weight"], head_marked(value_dim, 3))) + self.assertTrue(torch.equal(corrected[prefix + "in_proj_a.weight"], head_marked(num_v_heads))) + self.assertTrue(torch.equal(corrected[prefix + "in_proj_b.weight"], head_marked(num_v_heads))) + self.assertTrue(torch.equal(corrected[prefix + "A_log"], torch.zeros(num_v_heads))) + self.assertTrue( + torch.equal( + corrected[prefix + "dt_bias"], + torch.arange(num_v_heads, dtype=torch.float32), + ) + ) + conv = corrected[prefix + "conv1d.weight"] + self.assertEqual(conv.shape, (conv_dim, 1, 2)) + self.assertTrue( + torch.equal( + conv[:, 0, :], + torch.arange(conv_dim, dtype=torch.float32).unsqueeze(1).repeat(1, 2), + ) + ) + self.assertTrue( + torch.equal(corrected[prefix + "out_proj.weight"], head_marked(3, value_dim)) + ) + # unshifted norm + self.assertTrue( + torch.equal( + corrected["model.language_model.layers.0.input_layernorm.weight"], + torch.ones(4), + ) + ) + + def test_reorders_quantized_tiled_v_heads(self): + num_k_heads, num_v_heads, head_dim = 2, 4, 1 + value_dim = num_v_heads * head_dim + tiled = [0, 2, 1, 3] + v_stored = torch.arange(value_dim, dtype=torch.float32).unsqueeze(1)[tiled] + q_tensor = self.loader.GGMLTensor( + v_stored.to(torch.bfloat16), + tensor_type=gguf.GGMLQuantizationType.BF16, + tensor_shape=v_stored.shape, + ) + reordered = self.loader._qwen35_v_reorder( + q_tensor, num_v_heads, num_k_heads, head_dim + ) + self.assertFalse(self.loader.is_quantized(reordered)) + self.assertTrue( + torch.equal( + reordered.float(), + torch.arange(value_dim, dtype=torch.float32).unsqueeze(1), + ) + ) + + def test_v_head_reorder_is_identity_for_balanced_heads(self): + num_k_heads, num_v_heads = 16, 16 + head_dim = 2 + value_dim = num_v_heads * head_dim + key_dim = num_k_heads * head_dim + conv_dim = 2 * key_dim + value_dim + prefix = "model.language_model.layers.0.linear_attn." + sd = { + prefix + "in_proj_qkv.weight": torch.arange(conv_dim * 3, dtype=torch.float32).reshape(conv_dim, 3), + prefix + "in_proj_z.weight": torch.arange(value_dim * 3, dtype=torch.float32).reshape(value_dim, 3), + prefix + "A_log": torch.full((num_v_heads,), -1.0), + prefix + "dt_bias": torch.arange(num_v_heads, dtype=torch.float32), + prefix + "conv1d.weight": torch.arange(conv_dim * 4, dtype=torch.float32).reshape(conv_dim, 4), + prefix + "out_proj.weight": torch.arange(3 * value_dim, dtype=torch.float32).reshape(3, value_dim), + } + + corrected = self.loader.qwen35_corrections(sd) + + conv = corrected[prefix + "conv1d.weight"] + self.assertEqual(conv.shape, (conv_dim, 1, 4)) + self.assertTrue( + torch.equal(conv[:, 0, :], torch.arange(conv_dim * 4).reshape(conv_dim, 4)) + ) + self.assertTrue( + torch.equal( + corrected[prefix + "in_proj_qkv.weight"], + sd[prefix + "in_proj_qkv.weight"], + ) + ) + self.assertTrue( + torch.equal(corrected[prefix + "out_proj.weight"], sd[prefix + "out_proj.weight"]) + ) + + def test_clip_loader_dequantizes_quantized_lm_head(self): + # BaseGenerate.logits() feeds lm_head straight to F.linear without + # dequantizing GGML tensors, so a Q8_0 head (raw bytes with scales + # interleaved) must arrive dequantized with its logical shape. + with mock.patch.object( + self.loader, + "gguf_sd_loader", + return_value=( + { + "token_embd.weight": torch.zeros((248320, 4096)), + "output.weight": self.loader.GGMLTensor( + torch.zeros((248320, 4352), dtype=torch.uint8), + tensor_type=gguf.GGMLQuantizationType.Q8_0, + tensor_shape=(248320, 4096), + ), + "blk.0.attn_qkv.weight": torch.zeros((8192, 4096)), + "blk.0.attn_gate.weight": torch.zeros((4096, 4096)), + "blk.0.ssm_a": torch.full((32,), -1.0), + }, + {"arch_str": "qwen35"}, + ), + ), mock.patch.object( + self.loader, + "gguf_mmproj_loader", + return_value={}, + ): + state_dict = self.loader.gguf_clip_loader("Qwen3.5-9B-Q8_0.gguf") + + head = state_dict["lm_head.weight"] + self.assertFalse(self.loader.is_quantized(head)) + self.assertEqual(head.shape, (248320, 4096)) + + def test_clip_loader_keeps_bf16_lm_head_quantized(self): + # BF16 storage keeps its logical shape, so Dynamic VRAM can keep it + # quantized and offloaded; only block-quantized heads are unsafe. + bf16_head = self.loader.GGMLTensor( + torch.zeros((248320, 4096), dtype=torch.bfloat16), + tensor_type=gguf.GGMLQuantizationType.BF16, + tensor_shape=(248320, 4096), + ) + with mock.patch.object( + self.loader, + "gguf_sd_loader", + return_value=( + { + "token_embd.weight": torch.zeros((248320, 4096)), + "output.weight": bf16_head, + "blk.0.attn_qkv.weight": torch.zeros((8192, 4096)), + "blk.0.attn_gate.weight": torch.zeros((4096, 4096)), + "blk.0.ssm_a": torch.full((32,), -1.0), + }, + {"arch_str": "qwen35"}, + ), + ), mock.patch.object( + self.loader, + "gguf_mmproj_loader", + return_value={}, + ): + state_dict = self.loader.gguf_clip_loader("Qwen3.5-9B-BF16.gguf") + + head = state_dict["lm_head.weight"] + self.assertTrue(self.loader.is_quantized(head)) + self.assertEqual(head.shape, (248320, 4096)) + + def test_mmproj_routes_fused_qkv_through_qwen3_vision_map(self): + with TemporaryDirectory() as temp_dir: + text_encoder = Path(temp_dir) / "Qwen3.5-4B-Q8_0.gguf" + mmproj = Path(temp_dir) / "mmproj-Qwen3.5-4B-BF16.gguf" + text_encoder.touch() + mmproj.touch() + + with mock.patch.object( + self.loader, + "gguf_sd_loader", + return_value=( + { + "v.blk.0.attn_qkv.weight": torch.ones((1024, 3072)), + "v.blk.0.attn_out.weight": torch.ones((1024, 1024)), + "v.blk.0.ln1.weight": torch.ones(1024), + "v.blk.0.ln2.weight": torch.ones(1024), + "v.blk.0.ffn_up.weight": torch.ones((1024, 4096)), + "v.blk.0.ffn_down.weight": torch.ones((4096, 1024)), + "v.patch_embd.weight": torch.ones((2, 3, 2, 2)), + "v.patch_embd.weight.1": torch.full((2, 3, 2, 2), 2.0), + "v.patch_embd.bias": torch.ones(1024), + "v.position_embd.weight": torch.ones((1024, 2304)), + "mm.0.weight": torch.ones(1), + "mm.2.weight": torch.ones(1), + "v.post_ln.weight": torch.ones(1024), + }, + {}, + ), + ): + mapped = self.loader.gguf_mmproj_loader(str(text_encoder)) + + self.assertIn("model.visual.blocks.0.attn.qkv.weight", mapped) + self.assertIn("model.visual.blocks.0.attn.proj.weight", mapped) + self.assertIn("model.visual.blocks.0.norm1.weight", mapped) + self.assertIn("model.visual.blocks.0.norm2.weight", mapped) + self.assertIn("model.visual.blocks.0.mlp.linear_fc1.weight", mapped) + self.assertIn("model.visual.blocks.0.mlp.linear_fc2.weight", mapped) + self.assertIn("model.visual.patch_embed.proj.weight", mapped) + self.assertIn("model.visual.patch_embed.proj.bias", mapped) + self.assertIn("visual.pos_embed.weight", mapped) + self.assertIn("model.visual.merger.linear_fc1.weight", mapped) + self.assertIn("model.visual.merger.linear_fc2.weight", mapped) + self.assertIn("model.visual.merger.norm.weight", mapped) + self.assertEqual( + mapped["model.visual.patch_embed.proj.weight"].shape, + (2, 3, 2, 2, 2), + ) + + +class Qwen3VLQuantizationTests(unittest.TestCase): + def test_supports_quant_types_used_by_pruned_32b_gguf(self): + for quant_name in ("IQ3_S", "IQ3_XXS", "IQ2_S", "IQ2_XS"): + quant_type = getattr(gguf.GGMLQuantizationType, quant_name) + _, type_size = gguf.GGML_QUANT_SIZES[quant_type] + output = dequantize( + torch.zeros((1, type_size), dtype=torch.uint8), + quant_type, + (256,), + dtype=torch.float32, + ) + + self.assertEqual(output.shape, (256,)) + self.assertEqual(output.dtype, torch.float32) + self.assertIn(quant_type, dequantize_functions) + + +class Gemma4GGUFLoaderTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.loader = load_gguf_loader() + + def test_maps_e4b_specific_tensor_layout_to_comfyui(self): + state_dict = { + "token_embd.weight": torch.zeros(1), + "per_layer_token_embd.weight": torch.zeros(1), + "per_layer_model_proj.weight": torch.zeros(1), + "per_layer_proj_norm.weight": torch.zeros(1), + "blk.0.inp_gate.weight": torch.zeros(1), + "blk.0.proj.weight": torch.zeros(1), + "blk.0.layer_output_scale.weight": torch.zeros(1), + "blk.0.post_norm.weight": torch.zeros(1), + "blk.0.post_ffw_norm.weight": torch.zeros(1), + } + + mapped = self.loader.sd_map_replace(state_dict, self.loader.GEMMA4_SD_MAP) + + self.assertEqual( + set(mapped), + { + "model.embed_tokens.weight", + "model.embed_tokens_per_layer.weight", + "model.per_layer_model_projection.weight", + "model.per_layer_projection_norm.weight", + "model.layers.0.per_layer_input_gate.weight", + "model.layers.0.per_layer_projection.weight", + "model.layers.0.layer_scalar", + "model.layers.0.post_per_layer_input_norm.weight", + "model.layers.0.post_feedforward_layernorm.weight", + }, + ) + + def test_recreates_gemma4_bpe_tokenizer_json(self): + tokenizer_json = self.loader.gemma4_tokenizer_json( + ["", "", "", "", "hello", "\u2581world"], + ["h e", "he llo"], + [3, 3, 3, 3, 1, 1], + ) + tokenizer = json.loads(bytes(tokenizer_json.tolist())) + + self.assertEqual(tokenizer["model"]["type"], "BPE") + self.assertEqual(tokenizer["model"]["vocab"]["hello"], 4) + self.assertEqual(tokenizer["pre_tokenizer"]["type"], "Metaspace") + self.assertEqual( + [token["content"] for token in tokenizer["added_tokens"]], + ["", "", "", ""], + ) + + def test_e4b_layout_is_detected_by_installed_comfyui(self): + state_dict = { + "model.layers.0.post_feedforward_layernorm.weight": torch.zeros(2560), + "model.layers.41.self_attn.q_norm.weight": torch.zeros(256), + } + + self.assertEqual( + comfy.sd.detect_te_model(state_dict), + comfy.sd.TEModel.GEMMA_4_E4B, + ) + + +class MinimaxH3DetectionTests(unittest.TestCase): + def test_detects_native_minimax_h3_checkpoint_layout(self): + checkpoint_keys = { + "video_patch_proj.weight", + "audio_patch_proj.weight", + "blocks.0.attn.qkv_proj.weight", + "final_layer.video_out.weight", + } + + model_arch = detect_arch(checkpoint_keys) + + self.assertIsInstance(model_arch, ModelMinimaxH3) + self.assertEqual(model_arch.arch, "minimax_h3") + + def test_keeps_adaln_curve_table_in_full_precision(self): + model_arch = ModelMinimaxH3() + + self.assertIn("adaln_t_table", model_arch.keys_hiprec) + + def test_q4_preserves_sensitive_transformer_components(self): + state_dict = { + "blocks.0.attn.qkv_proj.weight": torch.ones((512, 512), dtype=torch.float16), + "blocks.0.adaln_proj.linear.weight": torch.ones((512, 8), dtype=torch.float16), + "blocks.0.attn.norm1.weight": torch.ones((512,), dtype=torch.bfloat16), + "blocks.0.mlp.norm.weight": torch.ones((512,), dtype=torch.bfloat16), + "blocks.0.modulation.weight": torch.ones((512, 512), dtype=torch.float16), + "condition_proj.weight": torch.ones((512, 512), dtype=torch.bfloat16), + "final_layer.video_out.weight": torch.ones((96, 512), dtype=torch.float32), + } + + with TemporaryDirectory() as temp_dir: + source_path = Path(temp_dir) / "minimax_h3.safetensors" + output_path = Path(temp_dir) / "minimax_h3-Q4_CR_W4A4.gguf" + save_file(state_dict, str(source_path)) + + converted_path, _ = convert_file( + str(source_path), + str(output_path), + interact=False, + quant_type_name="Q4_CR_W4A4", + quantization_device="cpu", + ) + + reader = gguf.GGUFReader(converted_path) + tensor_types = {tensor.name: tensor.tensor_type for tensor in reader.tensors} + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + del reader + + self.assertEqual( + tensor_types["blocks.0.attn.qkv_proj.weight"], + gguf.GGMLQuantizationType.F32, + ) + self.assertEqual( + tensor_types["blocks.0.adaln_proj.linear.weight"], + gguf.GGMLQuantizationType.F32, + ) + self.assertEqual( + tensor_types["blocks.0.attn.norm1.weight"], + gguf.GGMLQuantizationType.BF16, + ) + self.assertEqual( + tensor_types["blocks.0.modulation.weight"], + gguf.GGMLQuantizationType.F32, + ) + self.assertEqual( + tensor_types["condition_proj.weight"], + gguf.GGMLQuantizationType.BF16, + ) + + def test_q4_preserves_conditioning_refiner_and_output_paths(self): + state_dict = { + "video_patch_proj.weight": torch.ones((512, 96), dtype=torch.float32), + "audio_patch_proj.weight": torch.ones((512, 32), dtype=torch.float32), + "blocks.0.attn.qkv_proj.weight": torch.ones((512, 512), dtype=torch.float16), + "condition_proj.weight": torch.ones((512, 512), dtype=torch.bfloat16), + "token_refiner.blocks.0.attn.qkv_proj.weight": torch.ones( + (512, 512), dtype=torch.bfloat16 + ), + "final_layer.video_out.weight": torch.ones((96, 512), dtype=torch.float32), + "adaln_t_table": torch.ones((32, 8), dtype=torch.float32), + } + + with TemporaryDirectory() as temp_dir: + source_path = Path(temp_dir) / "minimax_h3.safetensors" + output_path = Path(temp_dir) / "minimax_h3-Q4_CR_W4A4.gguf" + save_file(state_dict, str(source_path)) + + converted_path, _ = convert_file( + str(source_path), + str(output_path), + interact=False, + quant_type_name="Q4_CR_W4A4", + quantization_device="cpu", + ) + + reader = gguf.GGUFReader(converted_path) + tensor_types = {tensor.name: tensor.tensor_type for tensor in reader.tensors} + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + del reader + + self.assertEqual( + tensor_types["blocks.0.attn.qkv_proj.weight"], + gguf.GGMLQuantizationType.F32, + ) + self.assertEqual( + tensor_types["condition_proj.weight"], + gguf.GGMLQuantizationType.BF16, + ) + self.assertEqual( + tensor_types["token_refiner.blocks.0.attn.qkv_proj.weight"], + gguf.GGMLQuantizationType.BF16, + ) + self.assertEqual( + tensor_types["final_layer.video_out.weight"], + gguf.GGMLQuantizationType.F32, + ) + + def test_converts_to_minimax_h3_gguf_with_full_precision_adaln_table(self): + state_dict = { + "video_patch_proj.weight": torch.ones((32, 32), dtype=torch.float16), + "audio_patch_proj.weight": torch.ones((32, 32), dtype=torch.float16), + "blocks.0.attn.qkv_proj.weight": torch.ones((96, 32), dtype=torch.float16), + "final_layer.video_out.weight": torch.ones((96, 32), dtype=torch.float16), + "adaln_t_table": torch.ones((32, 32), dtype=torch.float32), + } + + with TemporaryDirectory() as temp_dir: + source_path = Path(temp_dir) / "minimax_h3.safetensors" + output_path = Path(temp_dir) / "minimax_h3-Q8_0.gguf" + save_file(state_dict, str(source_path)) + + converted_path, model_arch = convert_file( + str(source_path), + str(output_path), + interact=False, + quant_type_name="Q8_0", + ) + + reader = gguf.GGUFReader(converted_path) + tensor_types = {tensor.name: tensor.tensor_type for tensor in reader.tensors} + architecture = reader.get_field("general.architecture") + architecture_name = str(architecture.parts[architecture.data[-1]], "utf-8") + del architecture + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + del reader + + self.assertEqual(model_arch.arch, "minimax_h3") + self.assertEqual(architecture_name, "minimax_h3") + self.assertEqual( + tensor_types["blocks.0.attn.qkv_proj.weight"], + gguf.GGMLQuantizationType.Q8_0, + ) + self.assertEqual( + tensor_types["adaln_t_table"], + gguf.GGMLQuantizationType.F32, + ) + + +class Q4CRW4A4QuantizationTests(unittest.TestCase): + def test_int4_cr_w4a4_packs_kitchen_native_layout(self): + torch.manual_seed(0) + weight = torch.randn(64, 512, dtype=torch.float32) + packed, wscales, quant_conf, orig_shape = quantize_int4_cr_w4a4( + weight, convrot_groupsize=256, quant_group_size=64, device=torch.device("cpu") + ) + + # packed (N, K//2) int8, scales (N,) float32 (per-output-row) + self.assertEqual(packed.shape, (64, 256)) + self.assertEqual(packed.dtype, torch.int8) + self.assertEqual(wscales.shape, (64,)) + self.assertEqual(wscales.dtype, torch.float32) + self.assertEqual(orig_shape, (64, 512)) + self.assertEqual(quant_conf["format"], "int4_cr") + self.assertEqual(quant_conf["backing"], "w4a4") + self.assertEqual(quant_conf["convrot_groupsize"], 256) + self.assertEqual(quant_conf["quant_group_size"], 64) + self.assertTrue(quant_conf["sym"]) + + def test_int4_cr_w4a4_roundtrips_through_dequant(self): + torch.manual_seed(0) + weight = torch.randn(64, 1024, dtype=torch.float32) + packed, wscales, quant_conf, _ = quantize_int4_cr_w4a4( + weight, convrot_groupsize=256, quant_group_size=64, device=torch.device("cpu") + ) + + # Mirror the ops._dequantized_weight math: signed two's-complement int4 + # per-output-row scale, then un-rotate by the block-diagonal Hadamard. + n, k = weight.shape + x32 = packed.to(torch.int32) + lo = (x32 & 0x0F).to(torch.float32) + hi = ((x32 >> 4) & 0x0F).to(torch.float32) + nibbles = torch.stack([lo, hi], dim=-1).reshape(n, k) + nibbles = torch.where(nibbles >= 8, nibbles - 16, nibbles) + w_rot = nibbles * wscales.reshape(-1, 1) + + # Un-rotate to the original basis. + cg = 256 + from tools.convert import _build_regular_hadamard + h = _build_regular_hadamard(cg, dtype=torch.float32, device="cpu") + ng = k // cg + w_rec = (w_rot.reshape(n, ng, cg) @ h).reshape(n, k) + + relative = (w_rec - weight).abs().amax().item() / weight.abs().amax().item() + self.assertLess(relative, 0.3) + + def test_int4_cr_w4a4_rejects_k_not_divisible_by_convrot(self): + # K=300 is not divisible by the 256 convrot group -> must reject. + with self.assertRaisesRegex(ValueError, "must divide input features"): + quantize_int4_cr_w4a4( + torch.randn(64, 300), convrot_groupsize=256, quant_group_size=64, + device=torch.device("cpu") + ) + with self.assertRaisesRegex(ValueError, "must divide input features"): + quantize_int4_cr_w4a4( + torch.randn(64, 224), convrot_groupsize=256, quant_group_size=64, + device=torch.device("cpu") + ) + + def test_int4_cr_w4a4_serializes_kitchen_metadata(self): + state_dict = { + "video_patch_proj.weight": torch.ones((32, 32), dtype=torch.float32), + "audio_patch_proj.weight": torch.ones((32, 32), dtype=torch.float32), + "blocks.0.mlp.fc1.weight": torch.ones((128, 512), dtype=torch.float32), + "final_layer.video_out.weight": torch.ones((96, 512), dtype=torch.float32), + } + + with TemporaryDirectory() as temp_dir: + source_path = Path(temp_dir) / "minimax_h3.safetensors" + output_path = Path(temp_dir) / "minimax_h3-Q4_CR_W4A4.gguf" + save_file(state_dict, str(source_path)) + + converted_path, _ = convert_file( + str(source_path), + str(output_path), + interact=False, + quant_type_name="Q4_CR_W4A4", + quantization_device="cpu", + ) + + reader = gguf.GGUFReader(converted_path) + tensor_types = {tensor.name: tensor.tensor_type for tensor in reader.tensors} + names = set(tensor_types.keys()) + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + del reader + + # Weight stays I8 (2 uint4 per byte) and per-row scale stored as F16. + self.assertEqual(tensor_types["blocks.0.mlp.fc1.weight"], gguf.GGMLQuantizationType.I8) + self.assertEqual(tensor_types["blocks.0.mlp.fc1.weight_scale"], gguf.GGMLQuantizationType.F16) + self.assertIn("blocks.0.mlp.fc1.weight_scale", names) + + def test_int4_cr_w4a4_loader_routes_to_w4a4_ops(self): + state_dict = { + "video_patch_proj.weight": torch.randn(32, 32, dtype=torch.float32), + "audio_patch_proj.weight": torch.randn(32, 32, dtype=torch.float32), + "blocks.0.mlp.fc1.weight": torch.randn(128, 512, dtype=torch.float32), + "final_layer.video_out.weight": torch.randn(96, 512, dtype=torch.float32), + } + + with TemporaryDirectory() as temp_dir: + source_path = Path(temp_dir) / "minimax_h3.safetensors" + output_path = Path(temp_dir) / "minimax_h3-Q4_CR_W4A4.gguf" + save_file(state_dict, str(source_path)) + + converted_path, _ = convert_file( + str(source_path), + str(output_path), + interact=False, + quant_type_name="Q4_CR_W4A4", + quantization_device="cpu", + ) + + loader = load_gguf_loader() + sd, extra = loader.gguf_sd_loader(converted_path) + + w = sd["blocks.0.mlp.fc1.weight"] + sd_key = "blocks.0.mlp.fc1" + qraw = sd[f"{sd_key}.comfy_quant"] + quant_conf = json.loads(bytes(qraw.tolist()).decode("utf-8")) + + w_detached = w.detach().clone().tolist() + del w + del qraw + del sd + del extra + gc.collect() + + self.assertEqual(quant_conf["format"], "int4_cr") + self.assertEqual(quant_conf["backing"], "w4a4") + self.assertEqual(quant_conf["orig_shape"], [128, 512]) + + def test_minimax_h3_q4_keeps_residual_projections_in_full_precision(self): + state_dict = { + "video_patch_proj.weight": torch.ones((32, 32), dtype=torch.float32), + "audio_patch_proj.weight": torch.ones((32, 32), dtype=torch.float32), + "blocks.0.attn.qkv_proj.weight": torch.ones((96, 512), dtype=torch.float32), + "blocks.0.attn.out_proj.weight": torch.ones((96, 512), dtype=torch.float32), + "blocks.0.mlp.fc1.weight": torch.ones((128, 512), dtype=torch.float32), + "blocks.0.mlp.fc2.weight": torch.ones((96, 512), dtype=torch.float32), + "final_layer.video_out.weight": torch.ones((96, 32), dtype=torch.float32), + } + + with TemporaryDirectory() as temp_dir: + source_path = Path(temp_dir) / "minimax_h3.safetensors" + output_path = Path(temp_dir) / "minimax_h3-Q4_CR_W4A4.gguf" + save_file(state_dict, str(source_path)) + + converted_path, _ = convert_file( + str(source_path), + str(output_path), + interact=False, + quant_type_name="Q4_CR_W4A4", + quantization_device="cpu", + ) + + reader = gguf.GGUFReader(converted_path) + tensor_types = {tensor.name: tensor.tensor_type for tensor in reader.tensors} + quant_config_names = { + field_name[len("comfy.gguf.quant."):] + for field_name in reader.fields + if field_name.startswith("comfy.gguf.quant.") + } + reader.tensors.clear() + reader.fields.clear() + reader.data._mmap.close() + del reader + + self.assertEqual( + tensor_types["blocks.0.attn.qkv_proj.weight"], + gguf.GGMLQuantizationType.F32, + ) + self.assertEqual( + tensor_types["blocks.0.attn.out_proj.weight"], + gguf.GGMLQuantizationType.F32, + ) + self.assertEqual( + tensor_types["blocks.0.mlp.fc2.weight"], + gguf.GGMLQuantizationType.F32, + ) + self.assertNotIn("blocks.0.attn.qkv_proj.weight", quant_config_names) + self.assertNotIn("blocks.0.attn.out_proj.weight", quant_config_names) + self.assertNotIn("blocks.0.mlp.fc2.weight", quant_config_names) + + def test_int4_cr_w4a4_ops_empty_patch_keeps_native_kernel(self): + # With no weight_function (the normal DynamicVRAM case when no LoRA is active), + # the native W4A4 kernel runs and the device-resident QuantizedTensor is cached. + # A real patch (LoRA) must never be probed against the packed int8, because a + # LoRA adapter swallows its addmm-on-int8 error and silently returns the int8 + # weight unchanged, which would drop the adapter delta and misroute to native. + ops = ops_factory() + lin = ops.Linear(64, 64, bias=True) + weight2 = torch.randn(64, 512, dtype=torch.float32) + packed, wscales, quant_conf, orig_shape = quantize_int4_cr_w4a4( + weight2, convrot_groupsize=256, quant_group_size=64, device=torch.device("cpu") + ) + quant_conf["orig_shape"] = [64, 512] + lin._load_from_state_dict( + { + "weight": packed, + "weight_scale": wscales.half(), + "comfy_quant": torch.tensor(list(json.dumps(quant_conf).encode("utf-8"))), + "bias": torch.randn(64, dtype=torch.float32), + }, + prefix="", + local_metadata={}, + strict=True, + missing_keys=[], + unexpected_keys=[], + error_msgs=[], + ) + self.assertTrue(lin._quantized) + + x = torch.randn(8, 512, dtype=torch.float32) + + # No patch: use the cached QuantizedTensor fast path. + lin.weight_function = [] + lin.bias_function = [] + lin._quantized_weight = None + lin._quantized_weight_device = None + out = lin(x) + self.assertEqual(out.shape, (8, 64)) + self.assertIsNotNone(lin._quantized_weight) + self.assertEqual(lin._quantized_weight_device, str(x.device)) + + # A real (value-changing) patch: cache the patched floating-point weight so + # the adapter delta is not erased by INT4 re-quantization. + def patch(t, *a, **kw): + return t.to(torch.float32) + 1.0 + lin.weight_function = [patch] + lin._quantized_weight = None + lin._quantized_weight_device = None + out = lin(x) + self.assertEqual(out.shape, (8, 64)) + self.assertIsNone(lin._quantized_weight) + self.assertIsNotNone(lin._fused_weight) + self.assertFalse(isinstance(lin._fused_weight, torch.Tensor) and lin._fused_weight.dtype == torch.int8) + + def test_int4_cr_w4a4_ops_value_patch_is_cached_in_compute_dtype(self): + # A genuine value-changing patch is cached in the compute dtype rather than + # re-quantized, because INT4 rounding can erase a small adapter delta. + ops = ops_factory() + lin = ops.Linear(64, 64, bias=True) + weight = torch.randn(64, 512, dtype=torch.float32) + packed, wscales, quant_conf, orig_shape = quantize_int4_cr_w4a4( + weight, convrot_groupsize=256, quant_group_size=64, device=torch.device("cpu") + ) + quant_conf["orig_shape"] = [64, 512] + lin._load_from_state_dict( + { + "weight": packed, + "weight_scale": wscales.half(), + "comfy_quant": torch.tensor(list(json.dumps(quant_conf).encode("utf-8"))), + "bias": torch.randn(64, dtype=torch.float32), + }, + prefix="", + local_metadata={}, + strict=True, + missing_keys=[], + unexpected_keys=[], + error_msgs=[], + ) + x = torch.randn(8, 512, dtype=torch.float32) + + # A patch that returns fp32 must be fused and retained as a floating weight. + def patch(t, *a, **kw): + return t.to(torch.float32) + lin.weight_function = [patch] + lin._quantized_weight = None + lin._quantized_weight_device = None + out = lin(x) + self.assertEqual(out.shape, (8, 64)) + self.assertIsNotNone(lin._fused_weight) + self.assertTrue(isinstance(lin._fused_weight, torch.Tensor)) + self.assertEqual(lin._fused_weight.dtype, lin._compute_dtype) + self.assertIsNone(lin._quantized_weight) + + def test_int4_cr_w4a4_ops_eager_fusion_and_cache_eviction(self): + ops = ops_factory() + lin = ops.Linear(64, 64, bias=True) + weight = torch.randn(64, 512, dtype=torch.float32) + packed, wscales, quant_conf, _ = quantize_int4_cr_w4a4( + weight, convrot_groupsize=256, quant_group_size=64, device=torch.device("cpu") + ) + quant_conf["orig_shape"] = [64, 512] + lin._load_from_state_dict( + { + "weight": packed, + "weight_scale": wscales.half(), + "comfy_quant": torch.tensor(list(json.dumps(quant_conf).encode("utf-8"))), + "bias": torch.zeros(64, dtype=torch.float32), + }, + prefix="", + local_metadata={}, + strict=True, + missing_keys=[], + unexpected_keys=[], + error_msgs=[], + ) + + calls = [] + + def patch(t, *args, **kwargs): + calls.append(None) + return t + 0.25 + + lin.weight_function = [patch] + lin._get_cached_quantized_weight(torch.device("cpu")) + self.assertIsNotNone(lin._quantized_weight) + self.assertTrue(lin.prepare_fused_weight(torch.device("cpu"))) + self.assertEqual(len(calls), 1) + self.assertIsNone(lin._quantized_weight) + first_fused = id(lin._fused_weight) + first_patch_id = lin._fused_patch_id + self.assertEqual(lin._fused_weight.device, torch.device("cpu")) + + # A forward on another device must use a temporary copy, not promote the + # persistent offload cache and retain it on CUDA for subsequent layers. + cached = lin._get_cached_fused_weight(torch.device("cuda:0")) + self.assertEqual(id(cached), first_fused) + self.assertEqual(cached.device, torch.device("cpu")) + + lin(torch.randn(2, 512)) + self.assertEqual(len(calls), 1) + self.assertEqual(id(lin._fused_weight), first_fused) + + lin.weight_function = [lambda t, *args, **kwargs: t - 0.25] + self.assertTrue(lin.prepare_fused_weight(torch.device("cpu"))) + self.assertIsNone(lin._quantized_weight) + self.assertNotEqual(lin._fused_patch_id, first_patch_id) + + lin.evict_quantized_caches() + self.assertIsNone(lin._quantized_weight) + self.assertIsNone(lin._fused_weight) + self.assertIsNone(lin._fused_bias) + + def test_int4_cr_w4a4_ops_non_native_patch_is_applied(self): + # A non-native patch must be applied to the dequantized, un-rotated weight + # and retained in the fallback cache. Standard LoRA uses the native low-rank + # bypass exercised by the following test instead. + ops = ops_factory() + lin = ops.Linear(64, 64, bias=True) + weight = torch.randn(64, 512, dtype=torch.float32) + packed, wscales, quant_conf, orig_shape = quantize_int4_cr_w4a4( + weight, convrot_groupsize=256, quant_group_size=64, device=torch.device("cpu") + ) + quant_conf["orig_shape"] = [64, 512] + lin._load_from_state_dict( + { + "weight": packed, + "weight_scale": wscales.half(), + "comfy_quant": torch.tensor(list(json.dumps(quant_conf).encode("utf-8"))), + "bias": torch.randn(64, dtype=torch.float32), + }, + prefix="", + local_metadata={}, + strict=True, + missing_keys=[], + unexpected_keys=[], + error_msgs=[], + ) + self.assertTrue(lin._quantized) + x = torch.randn(8, 512, dtype=torch.float32) + + # Low-rank LoRA-style delta encoded as a weight_function. + r = 8 + a = torch.randn(64, r, dtype=torch.float32) + b = torch.randn(r, 512, dtype=torch.float32) + delta = a @ b + + def lora_patch(t, *args, **kwargs): + # Only valid on a full-precision, un-rotated weight of the base shape. + self.assertEqual(t.shape, (64, 512)) + self.assertNotEqual(t.dtype, torch.int8) + return t + delta + + lin.weight_function = [lora_patch] + lin._quantized_weight = None + lin._quantized_weight_device = None + # Fusion keeps the adapter's complete delta in a cached compute-dtype weight; + # the plain packed cache is not retained. + out = lin(x) + self.assertEqual(out.shape, (8, 64)) + self.assertIsNotNone(lin._fused_weight) + self.assertIsNone(lin._quantized_weight) + + # Reference: apply the same patch to the dequantized weight. The cached fused + # weight should match it exactly in the compute dtype. + fused = lin._dequantized_weight(torch.device("cpu"), lin._compute_dtype) + fused = lora_patch(fused).to(dtype=lin._compute_dtype) + torch.testing.assert_close(lin._fused_weight, fused) + ref_out = torch.nn.functional.linear(x, fused.to(dtype=x.dtype)) + ref_bias = lin.bias.to(device=x.device, dtype=x.dtype) + if ref_bias is not None: + ref_out = ref_out + ref_bias + torch.testing.assert_close(out, ref_out, atol=1e-4, rtol=1e-3) + base_weight = lin._dequantized_weight(torch.device("cpu"), x.dtype) + base_out = torch.nn.functional.linear(x, base_weight, ref_bias) + self.assertFalse(torch.allclose(out, base_out, atol=1e-4, rtol=1e-3)) + + def test_int4_cr_w4a4_ops_standard_lora_uses_native_base_and_residual(self): + from comfy.weight_adapter import LoRAAdapter + + ops = ops_factory() + lin = ops.Linear(64, 64, bias=True) + weight = torch.randn(64, 512, dtype=torch.float32) + packed, wscales, quant_conf, orig_shape = quantize_int4_cr_w4a4( + weight, convrot_groupsize=256, quant_group_size=64, device=torch.device("cpu") + ) + quant_conf["orig_shape"] = [64, 512] + lin._load_from_state_dict( + { + "weight": packed, + "weight_scale": wscales.half(), + "comfy_quant": torch.tensor(list(json.dumps(quant_conf).encode("utf-8"))), + "bias": torch.randn(64, dtype=torch.float32), + }, + prefix="", + local_metadata={}, + strict=True, + missing_keys=[], + unexpected_keys=[], + error_msgs=[], + ) + + rank = 8 + up = torch.randn(64, rank, dtype=torch.float32) + down = torch.randn(rank, 512, dtype=torch.float32) + alpha = 4.0 + adapter = LoRAAdapter(set(), (up, down, alpha, None, None, None)) + key = "blocks.1.weight" + patches = {key: [(0.75, adapter, 1.0, None, None)]} + + class FakeLowVramPatch: + is_lowvram_patch = True + + def __init__(self, key, patches): + self.key = key + self.patches = patches + self.prepared_patches = None + + def __call__(self, value): + return value + + x = torch.randn(8, 512, dtype=torch.float32) + base_out = lin(x) + lin.weight_function = [FakeLowVramPatch(key, patches)] + + # The supported adapter must not trigger full-matrix fusion or retain a + # second full-precision weight; the native packed base remains cached. + self.assertIsNotNone(lin._native_lora_patch_entries()) + self.assertFalse(lin.prepare_fused_weight(torch.device("cpu"))) + self.assertIsNotNone(lin._quantized_weight) + out = lin(x) + scale = alpha / rank + expected = base_out + 0.75 * torch.nn.functional.linear( + torch.nn.functional.linear(x, down), up + ) * scale + torch.testing.assert_close(out, expected, atol=1e-4, rtol=1e-3) + self.assertIsNone(lin._fused_weight) + self.assertIsNotNone(lin._quantized_weight) + + def test_int4_cr_w4a4_ops_cpu_fallback_applies_native_lora(self): + from comfy.weight_adapter import LoRAAdapter + + ops = ops_factory() + lin = ops.Linear(64, 64, bias=True) + weight = torch.randn(64, 512, dtype=torch.float32) + packed, wscales, quant_conf, orig_shape = quantize_int4_cr_w4a4( + weight, convrot_groupsize=256, quant_group_size=64, device=torch.device("cpu") + ) + quant_conf["orig_shape"] = [64, 512] + lin._load_from_state_dict( + { + "weight": packed, + "weight_scale": wscales.half(), + "comfy_quant": torch.tensor(list(json.dumps(quant_conf).encode("utf-8"))), + "bias": torch.randn(64, dtype=torch.float32), + }, + prefix="", + local_metadata={}, + strict=True, + missing_keys=[], + unexpected_keys=[], + error_msgs=[], + ) + + rank = 8 + up = torch.randn(64, rank, dtype=torch.float32) + down = torch.randn(rank, 512, dtype=torch.float32) + adapter = LoRAAdapter(set(), (up, down, 4.0, None, None, None)) + key = "blocks.1.weight" + patches = {key: [(0.75, adapter, 1.0, None, None)]} + + class FakeLowVramPatch: + is_lowvram_patch = True + + def __init__(self, key, patches): + self.key = key + self.patches = patches + self.prepared_patches = None + + def __call__(self, value): + return value + + x = torch.randn(8, 512, dtype=torch.float32) + lin.weight_function = [FakeLowVramPatch(key, patches)] + native_entries = lin._native_lora_patch_entries() + + fallback = lin._cpu_forward_fallback(x, lin.bias, native_entries=native_entries) + expected = torch.nn.functional.linear( + x, lin._dequantized_weight(torch.device("cpu"), x.dtype), lin.bias + ) + 0.75 * torch.nn.functional.linear( + torch.nn.functional.linear(x, down), up + ) * (4.0 / rank) + torch.testing.assert_close(fallback, expected, atol=1e-4, rtol=1e-3) + self.assertEqual(fallback.device, torch.device("cpu")) + + def test_int4_cr_w4a4_ops_dispatches_to_cpu_on_cuda_oom(self): + ops = ops_factory() + lin = ops.Linear(4, 4, bias=False) + lin.weight = torch.zeros((4, 4), dtype=torch.float32) + lin._quantized = True + lin.weight_function = [lambda value: value] + input_tensor = mock.Mock(device=torch.device("cuda"), dtype=torch.float16) + sentinel = object() + + with mock.patch.object( + lin, "_get_cached_quantized_weight", side_effect=torch.OutOfMemoryError("test") + ), mock.patch.object(lin, "_cpu_forward_fallback", return_value=sentinel) as fallback: + result = lin.forward_comfy_cast_weights(input_tensor) + + self.assertIs(result, sentinel) + fallback.assert_called_once_with( + input_tensor, None, native_entries=None, use_fused_weight=True + ) + + def test_int4_cr_w4a4_ops_lokr_uses_native_base_and_residual(self): + from comfy.weight_adapter.lokr import LoKrAdapter + + ops = ops_factory() + lin = ops.Linear(64, 64, bias=True) + weight = torch.randn(64, 512, dtype=torch.float32) + packed, wscales, quant_conf, orig_shape = quantize_int4_cr_w4a4( + weight, convrot_groupsize=256, quant_group_size=64, device=torch.device("cpu") + ) + quant_conf["orig_shape"] = [64, 512] + lin._load_from_state_dict( + { + "weight": packed, + "weight_scale": wscales.half(), + "comfy_quant": torch.tensor(list(json.dumps(quant_conf).encode("utf-8"))), + "bias": torch.randn(64, dtype=torch.float32), + }, + prefix="", + local_metadata={}, + strict=True, + missing_keys=[], + unexpected_keys=[], + error_msgs=[], + ) + + # This matches the direct lokr_w1/lokr_w2 layout used by the Krea2 LoKr + # adapter. Keep the factors on CPU, as they are after normal LoRA loading. + w1 = torch.randn(8, 8, dtype=torch.float32) + w2 = torch.randn(8, 64, dtype=torch.float32) + adapter = LoKrAdapter( + set(), (w1, w2, None, None, None, None, None, None, None) + ) + key = "blocks.1.weight" + patches = {key: [(0.75, adapter, 1.0, None, None)]} + + class FakeLowVramPatch: + is_lowvram_patch = True + + def __init__(self, key, patches): + self.key = key + self.patches = patches + self.prepared_patches = None + + def __call__(self, value): + return value + + x = torch.randn(8, 512, dtype=torch.float32) + base_out = lin(x) + lin.weight_function = [FakeLowVramPatch(key, patches)] + + self.assertIsNotNone(lin._native_lora_patch_entries()) + self.assertFalse(lin.prepare_fused_weight(torch.device("cpu"))) + with mock.patch.object( + lin, "_get_cached_fused_weight", side_effect=AssertionError("LoKr fallback") + ): + out = lin(x) + + expected = base_out + 0.75 * torch.nn.functional.linear( + x, torch.kron(w1, w2) + ) + torch.testing.assert_close(out, expected, atol=1e-4, rtol=1e-3) + self.assertEqual(w1.device, torch.device("cpu")) + self.assertEqual(w2.device, torch.device("cpu")) + self.assertIsNone(lin._fused_weight) + self.assertIsNotNone(lin._quantized_weight) + + def test_int4_lora_factor_cache_reuses_cuda_factors_with_headroom(self): + ops = ops_factory() + lin = ops.Linear(64, 64, bias=False) + factors = [torch.randn(8, 64), torch.randn(64, 8), torch.randn(4, 4)] + cuda = torch.device("cuda:0") + reserve = 1024**3 + + with mock.patch.object( + torch.cuda, "mem_get_info", return_value=(2 * reserve, 8 * reserve) + ), mock.patch.object( + comfy.model_management, "extra_reserved_memory", return_value=reserve + ), mock.patch.object( + comfy.model_management, + "cast_to_device", + side_effect=lambda factor, device, dtype: factor, + ) as cast: + for factor in factors: + self.assertIs( + lin._get_cached_lora_factor(factor, cuda, torch.bfloat16), factor + ) + self.assertIs( + lin._get_cached_lora_factor(factor, cuda, torch.bfloat16), factor + ) + + self.assertEqual(cast.call_count, len(factors)) + self.assertEqual(len(lin._lora_factor_cache), len(factors)) + lin.evict_quantized_caches() + self.assertFalse(lin._lora_factor_cache) + + def test_int4_lora_factor_cache_streams_at_cuda_reserve_boundary(self): + ops = ops_factory() + lin = ops.Linear(64, 64, bias=False) + factor = torch.randn(8, 64) + cuda = torch.device("cuda:0") + reserve = 1024**3 + + with mock.patch.object( + torch.cuda, "mem_get_info", return_value=(reserve, 8 * reserve) + ), mock.patch.object( + comfy.model_management, "extra_reserved_memory", return_value=reserve + ), mock.patch.object( + comfy.model_management, + "cast_to_device", + side_effect=lambda factor, device, dtype: factor, + ) as cast: + lin._get_cached_lora_factor(factor, cuda, torch.bfloat16) + lin._get_cached_lora_factor(factor, cuda, torch.bfloat16) + + self.assertEqual(cast.call_count, 2) + self.assertFalse(lin._lora_factor_cache) + + def test_int4_lora_factor_cache_logs_loaded_models_and_reraises_cuda_oom(self): + ops = ops_factory() + lin = ops.Linear(64, 64, bias=False) + factor = torch.randn(8, 64) + cuda = torch.device("cuda:0") + loaded_model = type( + "FakeLoadedModel", + (), + { + "model": type( + "FakePatcher", + (), + {"model": type("CachedModel", (), {"name": "base"})(), "size": 2 * 1024**2}, + )() + }, + )() + missing_size_model = type( + "FakeLoadedModel", + (), + {"model": type("FakePatcher", (), {"model": type("AdapterModel", (), {})()})()}, + )() + + with mock.patch.object( + torch.cuda, "mem_get_info", return_value=(2 * 1024**3, 8 * 1024**3) + ), mock.patch.object( + comfy.model_management, + "current_loaded_models", + [loaded_model, missing_size_model], + ), mock.patch.object( + comfy.model_management, + "cast_to_device", + side_effect=torch.OutOfMemoryError("forced cache OOM"), + ), self.assertLogs(level="ERROR") as logs: + with self.assertRaisesRegex(torch.OutOfMemoryError, "forced cache OOM"): + lin._get_cached_lora_factor(factor, cuda, torch.bfloat16) + + self.assertIn("caching INT4 LoRA/LoKr factors", logs.output[0]) + self.assertIn("CachedModel (base): 2.0 MiB", logs.output[1]) + self.assertIn("AdapterModel: size unavailable", logs.output[2]) + + def test_int4_cr_w4a4_ops_fused_cache_survives_lowvram_patch_rebind(self): + # Dynamic VRAM (GGUFModelPatcherDynamic.load) promotes a freshly created + # LowVramPatch into weight_function every time the model is moved to device. + # That patch object is recreated on each reload, so the fused-weight cache must + # NOT be keyed on id(patch_object) — otherwise the expensive Hadamard rotate + + # int4 pack runs on every forward (a 12x slowdown measured on large layers). + # This test verifies the cache survives patch-object rebinding and only re-fuses + # when the underlying patch *content* (key / patches dict) actually changes. + ops = ops_factory() + lin = ops.Linear(64, 64, bias=True) + weight = torch.randn(64, 512, dtype=torch.float32) + packed, wscales, quant_conf, orig_shape = quantize_int4_cr_w4a4( + weight, convrot_groupsize=256, quant_group_size=64, device=torch.device("cpu") + ) + quant_conf["orig_shape"] = [64, 512] + lin._load_from_state_dict( + { + "weight": packed, + "weight_scale": wscales.half(), + "comfy_quant": torch.tensor(list(json.dumps(quant_conf).encode("utf-8"))), + "bias": torch.randn(64, dtype=torch.float32), + }, + prefix="", + local_metadata={}, + strict=True, + missing_keys=[], + unexpected_keys=[], + error_msgs=[], + ) + self.assertTrue(lin._quantized) + x = torch.randn(8, 512, dtype=torch.float32) + + # Mimic comfy.model_patcher.LowVramPatch: an object with is_lowvram_patch, + # a stable tensor key, and a pointer to the model-patcher-held patches dict. + class FakeLowVramPatch: + is_lowvram_patch = True + def __init__(self, key, patches): + self.key = key + self.patches = patches + self.prepared_patches = None + def __call__(self, weight): + # A genuine value-changing adapter, so it must fuse. + return weight + 0.1 + + # The dict/list live on the model patcher and persist across reloads. + key = "blocks.1.weight" + patches = {key: [(1.0, torch.full((64, 512), 0.5), 1.0, None, None)]} + + # First forward builds the fused weight. + lin.weight_function = [FakeLowVramPatch(key, patches)] + out = lin(x) + self.assertEqual(out.shape, (8, 64)) + self.assertIsNotNone(lin._fused_weight) + first_id = id(lin._fused_weight) + sig1 = lin._fused_patch_signature() + + # Simulate repeated DynamicVRAM reloads: a NEW patch object each forward, but + # the same underlying patches dict/list (which persists on the model patcher). + for _ in range(3): + lin.weight_function = [FakeLowVramPatch(key, patches)] + out = lin(x) + self.assertEqual(out.shape, (8, 64)) + # The cache must have been reused (same fused-QuantizedTensor, same signature). + self.assertEqual(id(lin._fused_weight), first_id) + self.assertEqual(lin._fused_patch_signature(), sig1) + + # A real patch change (new patches dict → new list) must invalidate the cache. + changed_patches = {key: [(1.0, torch.full((64, 512), 0.25), 1.0, None, None)]} + lin.weight_function = [FakeLowVramPatch(key, changed_patches)] + out = lin(x) + self.assertEqual(out.shape, (8, 64)) + self.assertIsNotNone(lin._fused_weight) + self.assertNotEqual(lin._fused_patch_signature(), sig1) + + def test_int4_lora_fusion_defaults_to_execution_device(self): + nodes_module = nodes_factory() + calls = [] + moves = [] + + modules = [ + type( + "FakeQuantizedLayer", + (), + { + "weight_function": [object()], + "bias_function": [], + "_fused_patch_signature": lambda self: ("patch",), + "prepare_fused_weight": lambda self, device: calls.append(device) or True, + "move_fused_caches": lambda self, device: moves.append(device), + }, + )() + ] + patcher = type( + "FakePatcher", + (), + { + "model": type("FakeModel", (), {"named_modules": lambda self: [("layer", modules[0])]} )(), + "load_device": torch.device("cuda:0"), + "offload_device": torch.device("cpu"), + "_gguf_patch_layout_signature": None, + }, + )() + + with mock.patch.object(nodes_module.comfy.model_management, "cuda_device_context", return_value=nullcontext()), \ + mock.patch.object(nodes_module.comfy.model_management, "throw_exception_if_processing_interrupted", None): + prepared = nodes_module.GGUFModelPatcher._prepare_gguf_quantized_weights(patcher) + + self.assertFalse(hasattr(nodes_module, "tqdm")) + self.assertEqual(prepared, 1) + self.assertEqual(calls, [torch.device("cuda:0")]) + self.assertEqual(moves, [torch.device("cpu")]) + + def test_int4_patch_preparation_evicts_caches_after_interrupt(self): + nodes_module = nodes_factory() + calls = [] + evictions = [] + modules = [ + type( + "FakeQuantizedLayer", + (), + { + "weight_function": [object()], + "bias_function": [], + "_fused_patch_signature": lambda self: ("patch",), + "prepare_fused_weight": lambda self, device: calls.append(device) or True, + }, + )(), + type( + "FakeQuantizedLayer", + (), + { + "weight_function": [object()], + "bias_function": [], + "_fused_patch_signature": lambda self: ("patch",), + "prepare_fused_weight": lambda self, device: calls.append(device) or True, + }, + )(), + ] + patcher = type( + "FakePatcher", + (), + { + "model": type( + "FakeModel", + (), + { + "named_modules": lambda self: [ + ("first", modules[0]), + ("second", modules[1]), + ] + }, + )(), + "load_device": torch.device("cpu"), + "offload_device": torch.device("cpu"), + "_gguf_patch_layout_signature": None, + "_evict_gguf_quantized_caches": lambda self: evictions.append(True), + }, + )() + with mock.patch.object( + nodes_module.comfy.model_management, + "throw_exception_if_processing_interrupted", + side_effect=RuntimeError("interrupted"), + ), mock.patch.object( + nodes_module.comfy.model_management, + "cuda_device_context", + return_value=nullcontext(), + ): + with self.assertRaisesRegex(RuntimeError, "interrupted"): + nodes_module.GGUFModelPatcher._prepare_gguf_quantized_weights(patcher) + + self.assertEqual(calls, []) + self.assertEqual(evictions, [True]) + + def test_dynamic_vram_warns_for_unique_offloaded_lora_factors(self): + nodes_module = nodes_factory() + factor = torch.ones((2, 2)) + adapter = type( + "FakeLoRAAdapter", + (), + {"name": "lora", "weights": (factor, factor, 1.0, None, None, None)}, + )() + key = "blocks.1.weight" + patch_function = type( + "FakeLowVramPatch", + (), + { + "is_lowvram_patch": True, + "key": key, + "patches": {key: [(1.0, adapter, 1.0, None, None)]}, + "prepared_patches": None, + }, + )() + module = type( + "FakeLayer", + (), + {"weight_function": [patch_function], "bias_function": []}, + )() + patcher = type( + "FakeDynamicPatcher", + (), + { + "model": type( + "FakeModel", + (), + {"named_modules": lambda self: [("layer", module)]}, + )(), + "load_device": torch.device("cuda:0"), + "offload_device": torch.device("cpu"), + }, + )() + + with mock.patch.dict( + os.environ, {"COMFYUI_GGUF_INT4_LORA_OFFLOAD": "yes"} + ), mock.patch.object( + nodes_module, "_DYNAMIC_VRAM_LORA_WARNING_MIN_BYTES", 1 + ), mock.patch.object(torch.cuda, "is_available", return_value=True), mock.patch.object( + torch.cuda, "mem_get_info", return_value=(2 * 1024**3, 8 * 1024**3) + ), self.assertLogs(level="WARNING") as logs: + warned = nodes_module.GGUFModelPatcherDynamic._warn_dynamic_vram_lora_streaming( + patcher + ) + + self.assertTrue(warned) + self.assertIn("1 unique offloaded LoRA/LoKr factor tensors", logs.output[0]) + self.assertIn("from cpu to cuda:0", logs.output[0]) + self.assertIn("CUDA free/total 2.0/8.0 GiB", logs.output[0]) + self.assertIn("does not guarantee an OOM", logs.output[0]) + + def test_dynamic_vram_preloads_unique_lora_and_lokr_factors(self): + nodes_module = nodes_factory() + shared = torch.ones((2, 2)) + lora_only = torch.ones((2, 3)) + lokr_only = torch.ones((3, 2)) + lora = type( + "FakeLoRAAdapter", + (), + {"name": "lora", "weights": (shared, lora_only, 4.0, None, "preserve")}, + )() + lokr = type( + "FakeLoKrAdapter", + (), + {"name": "lokr", "weights": (lokr_only, shared, None, 2.0)}, + )() + key = "blocks.1.weight" + patch_function = type( + "FakeLowVramPatch", + (), + { + "is_lowvram_patch": True, + "key": key, + "patches": { + key: [ + (1.0, lora, 1.0, None, None), + (1.0, lokr, 1.0, None, None), + ] + }, + "prepared_patches": None, + }, + )() + module = type( + "FakeLayer", + (), + {"weight_function": [patch_function], "bias_function": []}, + )() + replacements = {} + moved = [] + + def move_factor(factor, device): + moved.append((factor, device)) + replacement = factor.clone() + replacements[id(factor)] = replacement + return replacement + + patcher = type( + "FakeDynamicPatcher", + (), + { + "model": type( + "FakeModel", + (), + {"named_modules": lambda self: [("layer", module)]}, + )(), + "load_device": torch.device("cuda:0"), + "offload_device": torch.device("cpu"), + "_move_dynamic_vram_lora_factor": staticmethod(move_factor), + }, + )() + factor_bytes = sum( + factor.numel() * factor.element_size() + for factor in (shared, lora_only, lokr_only) + ) + reserve = 123456 + + with mock.patch.object(torch.cuda, "is_available", return_value=True), mock.patch.object( + torch.cuda, "mem_get_info", return_value=(reserve + factor_bytes, 8 * reserve) + ), mock.patch.object( + nodes_module.comfy.model_management, "extra_reserved_memory", return_value=reserve + ): + preloaded = nodes_module.GGUFModelPatcherDynamic._preload_dynamic_vram_lora_factors( + patcher + ) + + self.assertTrue(preloaded) + self.assertEqual(len(moved), 3) + self.assertEqual( + {id(factor) for factor, _ in moved}, + {id(shared), id(lora_only), id(lokr_only)}, + ) + self.assertTrue(all(device == torch.device("cuda:0") for _, device in moved)) + self.assertIs(lora.weights[0], replacements[id(shared)]) + self.assertIs(lora.weights[1], replacements[id(lora_only)]) + self.assertIs(lokr.weights[0], replacements[id(lokr_only)]) + self.assertIs(lokr.weights[1], replacements[id(shared)]) + self.assertEqual(lora.weights[2:], (4.0, None, "preserve")) + self.assertEqual(lokr.weights[2:], (None, 2.0)) + + def test_dynamic_vram_keeps_lora_factors_on_cuda_without_offload_enabled(self): + nodes_module = nodes_factory() + lora_factor = torch.ones((2, 2)) + lokr_factor = torch.ones((2, 3)) + lora = type( + "FakeLoRAAdapter", (), {"name": "lora", "weights": (lora_factor, 4.0)} + )() + lokr = type( + "FakeLoKrAdapter", (), {"name": "lokr", "weights": (lokr_factor, None)} + )() + original_lora_weights = lora.weights + original_lokr_weights = lokr.weights + key = "blocks.1.weight" + patch_function = type( + "FakeLowVramPatch", + (), + { + "is_lowvram_patch": True, + "key": key, + "patches": { + key: [ + (1.0, lora, 1.0, None, None), + (1.0, lokr, 1.0, None, None), + ] + }, + "prepared_patches": None, + }, + )() + module = type( + "FakeLayer", + (), + {"weight_function": [patch_function], "bias_function": []}, + )() + moved = [] + patcher = type( + "FakeDynamicPatcher", + (), + { + "model": type( + "FakeModel", + (), + {"named_modules": lambda self: [("layer", module)]}, + )(), + "load_device": torch.device("cuda:0"), + "offload_device": torch.device("cpu"), + "_move_dynamic_vram_lora_factor": staticmethod( + lambda factor, device: moved.append((factor, device)) or factor.clone() + ), + }, + )() + factor_bytes = ( + lora_factor.numel() * lora_factor.element_size() + + lokr_factor.numel() * lokr_factor.element_size() + ) + reserve = 123456 + + with mock.patch.dict( + os.environ, {"COMFYUI_GGUF_INT4_LORA_OFFLOAD": "false"} + ), mock.patch.object(torch.cuda, "is_available", return_value=True), mock.patch.object( + torch.cuda, "mem_get_info", return_value=(reserve + factor_bytes - 1, 8 * reserve) + ), mock.patch.object( + nodes_module.comfy.model_management, "extra_reserved_memory", return_value=reserve + ): + preloaded = nodes_module.GGUFModelPatcherDynamic._preload_dynamic_vram_lora_factors( + patcher + ) + + self.assertTrue(preloaded) + self.assertEqual(len(moved), 2) + self.assertIsNot(lora.weights, original_lora_weights) + self.assertIsNot(lokr.weights, original_lokr_weights) + + def test_dynamic_vram_offloads_lora_factors_only_when_enabled(self): + nodes_module = nodes_factory() + factor = torch.ones((2, 2)) + adapter = type("FakeLoRAAdapter", (), {"name": "lora", "weights": (factor, 4.0)})() + key = "blocks.1.weight" + patch_function = type( + "FakeLowVramPatch", + (), + { + "is_lowvram_patch": True, + "key": key, + "patches": {key: [(1.0, adapter, 1.0, None, None)]}, + "prepared_patches": None, + }, + )() + module = type( + "FakeLayer", (), {"weight_function": [patch_function], "bias_function": []} + )() + moved = [] + patcher = type( + "FakeDynamicPatcher", + (), + { + "model": type( + "FakeModel", (), {"named_modules": lambda self: [("layer", module)]} + )(), + "load_device": torch.device("cuda:0"), + "offload_device": torch.device("cpu"), + "_move_dynamic_vram_lora_factor": staticmethod( + lambda value, device: moved.append((value, device)) or value.clone() + ), + }, + )() + reserve = 123456 + + with mock.patch.dict( + os.environ, {"COMFYUI_GGUF_INT4_LORA_OFFLOAD": "ON"} + ), mock.patch.object(torch.cuda, "is_available", return_value=True), mock.patch.object( + torch.cuda, "mem_get_info", return_value=(reserve + factor.numel() * factor.element_size() - 1, 8 * reserve) + ), mock.patch.object( + nodes_module.comfy.model_management, "extra_reserved_memory", return_value=reserve + ): + preloaded = nodes_module.GGUFModelPatcherDynamic._preload_dynamic_vram_lora_factors( + patcher + ) + + self.assertFalse(preloaded) + self.assertFalse(moved) + self.assertIs(adapter.weights[0], factor) + + def test_int4_lora_offload_environment_values(self): + nodes_module = nodes_factory() + for value in ("true", "TRUE", "1", "yes", "On"): + with mock.patch.dict(os.environ, {"COMFYUI_GGUF_INT4_LORA_OFFLOAD": value}): + self.assertTrue(nodes_module.int4_lora_offload_enabled()) + for value in ("", "false", "0", "no", "off", "unexpected"): + with mock.patch.dict(os.environ, {"COMFYUI_GGUF_INT4_LORA_OFFLOAD": value}): + self.assertFalse(nodes_module.int4_lora_offload_enabled()) + with mock.patch.dict(os.environ, {}, clear=True): + self.assertFalse(nodes_module.int4_lora_offload_enabled()) + + def test_dynamic_vram_preload_logs_loaded_models_and_reraises_cuda_oom(self): + nodes_module = nodes_factory() + factor = torch.ones((2, 2)) + adapter = type("FakeLoRAAdapter", (), {"name": "lora", "weights": (factor, 4.0)})() + key = "blocks.1.weight" + patch_function = type( + "FakeLowVramPatch", + (), + { + "is_lowvram_patch": True, + "key": key, + "patches": {key: [(1.0, adapter, 1.0, None, None)]}, + "prepared_patches": None, + }, + )() + module = type( + "FakeLayer", (), {"weight_function": [patch_function], "bias_function": []} + )() + patcher = type( + "FakeDynamicPatcher", + (), + { + "model": type( + "FakeModel", (), {"named_modules": lambda self: [("layer", module)]} + )(), + "load_device": torch.device("cuda:0"), + "offload_device": torch.device("cpu"), + "_move_dynamic_vram_lora_factor": staticmethod( + mock.Mock(side_effect=torch.OutOfMemoryError("forced preload OOM")) + ), + }, + )() + loaded_model = type( + "FakeLoadedModel", + (), + { + "model": type( + "FakePatcher", + (), + {"model": type("DiffusionModel", (), {"name": "gguf"})(), "size": 3 * 1024**2}, + )() + }, + )() + + with mock.patch.dict( + os.environ, {"COMFYUI_GGUF_INT4_LORA_OFFLOAD": "false"} + ), mock.patch.object(torch.cuda, "is_available", return_value=True), mock.patch.object( + torch.cuda, "mem_get_info", return_value=(2 * 1024**3, 8 * 1024**3) + ), mock.patch.object( + nodes_module.comfy.model_management, "extra_reserved_memory", return_value=0 + ), mock.patch.object( + nodes_module.comfy.model_management, "current_loaded_models", [loaded_model] + ), self.assertLogs(level="ERROR") as logs: + with self.assertRaisesRegex(torch.OutOfMemoryError, "forced preload OOM"): + nodes_module.GGUFModelPatcherDynamic._preload_dynamic_vram_lora_factors( + patcher + ) + + self.assertIn("preloading Dynamic VRAM INT4 LoRA/LoKr factors", logs.output[0]) + self.assertIn("DiffusionModel (gguf): 3.0 MiB", logs.output[1]) + + def test_dynamic_vram_does_not_warn_for_static_lora_patch(self): + nodes_module = nodes_factory() + module = type( + "FakeLayer", + (), + {"weight_function": [object()], "bias_function": []}, + )() + patcher = type( + "FakeDynamicPatcher", + (), + { + "model": type( + "FakeModel", + (), + {"named_modules": lambda self: [("layer", module)]}, + )(), + "load_device": torch.device("cuda:0"), + "offload_device": torch.device("cpu"), + }, + )() + + self.assertFalse( + nodes_module.GGUFModelPatcherDynamic._warn_dynamic_vram_lora_streaming( + patcher + ) + ) + + def test_performance_log_records_quantized_forward(self): + import sys + + ops_factory() + ops_module = sys.modules["comfyui_gguf_test.ops"] + with TemporaryDirectory() as temp_dir: + log_path = Path(temp_dir) / "performance.log" + with mock.patch.dict(os.environ, {ops_module._PERF_LOG_ENV: str(log_path)}): + logger = ops_module._configure_perf_logger() + self.assertIsNotNone(logger) + handler = logger.handlers[-1] + previous_logger = ops_module._PERF_LOGGER + ops_module._PERF_LOGGER = logger + try: + class FakeLayer: + in_features = 16 + out_features = 32 + + input_tensor = torch.zeros((2, 16), dtype=torch.float32) + result = ops_module._perf_forward( + "int4_convrot_w4a4", + FakeLayer(), + input_tensor, + lambda: input_tensor + 1, + ) + self.assertTrue(torch.equal(result, torch.ones_like(input_tensor))) + handler.flush() + contents = log_path.read_text(encoding="utf-8") + finally: + ops_module._PERF_LOGGER = previous_logger + logger.removeHandler(handler) + handler.close() + + self.assertIn("forward mode=int4_convrot_w4a4", contents) + self.assertIn("elapsed_ms=", contents) + self.assertIn("input_shape=(2, 16)", contents) + + def test_performance_log_is_opt_in(self): + import sys + + ops_factory() + ops_module = sys.modules["comfyui_gguf_test.ops"] + with mock.patch.dict(os.environ, {}, clear=True): + self.assertIsNone(ops_module._configure_perf_logger()) + with mock.patch.dict(os.environ, {ops_module._PERF_LOG_ENV: ""}): + self.assertIsNone(ops_module._configure_perf_logger()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/conversion_webui.html b/tools/conversion_webui.html new file mode 100644 index 00000000..b9c77180 --- /dev/null +++ b/tools/conversion_webui.html @@ -0,0 +1,364 @@ + + + + + + + GGUF Conversion Dashboard + + + +
+
+

GGUF conversion dashboard

+

Local-only job control for tools\convert.py. Enter existing filesystem paths; checkpoints never leave this machine.

+
+
+
+

New conversion

+

Jobs run one at a time. Choose a new output filename; the converter will not replace an existing GGUF.

+
+
+
+ + +

Supports .safetensors, .ckpt, .pt, .pth, and .bin.

+
+
+ + +
+
+ Conversion plan +
+ + +
+
+
+ + +
+ + +
+ + +
+
+
+ + +
+
+ +
+ +

Optional .safetensors or .gguf adapters. Fusion happens before export.

+
+
+ + +
+
+
+ +
+
+ + + + diff --git a/tools/conversion_webui.py b/tools/conversion_webui.py new file mode 100644 index 00000000..465012ce --- /dev/null +++ b/tools/conversion_webui.py @@ -0,0 +1,351 @@ +"""Local web dashboard for tools/convert.py. + +Run from the repository root with: + python tools\\conversion_webui.py +""" + +from __future__ import annotations + +import argparse +import json +import os +import queue +import subprocess +import sys +import threading +import uuid +import webbrowser +from datetime import datetime, timezone +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + + +ROOT = Path(__file__).resolve().parent.parent +CONVERTER = ROOT / "tools" / "convert.py" +PAGE = Path(__file__).with_suffix(".html") +SOURCE_EXTENSIONS = {".safetensors", ".ckpt", ".pt", ".pth", ".bin"} +LORA_EXTENSIONS = {".safetensors", ".gguf"} +QUANT_TYPES = {"source", "Q8_0", "Q5_1", "Q5_0", "Q4_1", "Q4_0", "Q8_CR", "Q4_CR_W4A4"} +Q8_TYPES = {"Q8_CR", "Q8_0"} +DEVICES = {"auto", "cpu", "cuda"} + + +def now() -> str: + return datetime.now(timezone.utc).isoformat(timespec="seconds") + + +def resolve_path(value: str) -> Path: + path = Path(value).expanduser() + return path if path.is_absolute() else (ROOT / path).resolve() + + +def validate_path(value: Any, label: str, extensions: set[str]) -> Path: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{label} is required.") + path = resolve_path(value.strip()) + if not path.is_file(): + raise ValueError(f"{label} does not exist: {path}") + if path.suffix.lower() not in extensions: + allowed = ", ".join(sorted(extensions)) + raise ValueError(f"{label} must use one of: {allowed}") + return path + + +class ConversionManager: + def __init__(self) -> None: + self.jobs: dict[str, dict[str, Any]] = {} + self.pending: queue.Queue[str] = queue.Queue() + self.lock = threading.RLock() + self.active_process: subprocess.Popen[str] | None = None + self.active_job_id: str | None = None + self.worker = threading.Thread(target=self._work, daemon=True, name="gguf-conversion-worker") + self.worker.start() + + def create(self, payload: dict[str, Any]) -> dict[str, Any]: + source = validate_path(payload.get("source"), "Source checkpoint", SOURCE_EXTENSIONS) + destination_value = payload.get("destination") + if not isinstance(destination_value, str) or not destination_value.strip(): + raise ValueError("Destination path is required.") + destination = resolve_path(destination_value.strip()) + if destination.suffix.lower() != ".gguf": + raise ValueError("Destination path must end with .gguf.") + if not destination.parent.is_dir(): + raise ValueError(f"Destination folder does not exist: {destination.parent}") + if destination.exists(): + raise ValueError(f"Destination already exists: {destination}") + + mode = payload.get("mode", "quant") + quant_type = payload.get("quant_type", "Q8_CR") + if mode not in {"quant", "target"}: + raise ValueError("Conversion mode is invalid.") + if quant_type not in QUANT_TYPES: + raise ValueError("Quantization type is invalid.") + + target_size = payload.get("target_size_mb") + target_q8_type = payload.get("target_size_q8_type", "Q8_CR") + if mode == "target": + try: + target_size = float(target_size) + except (TypeError, ValueError) as error: + raise ValueError("Target size must be a number of MiB.") from error + if target_size <= 0: + raise ValueError("Target size must be greater than zero.") + if target_q8_type not in Q8_TYPES: + raise ValueError("Target-size Q8 baseline is invalid.") + + device = payload.get("device", "auto") + if device not in DEVICES: + raise ValueError("Quantization device is invalid.") + + loras: list[tuple[Path, float]] = [] + for index, entry in enumerate(payload.get("loras", []), start=1): + if not isinstance(entry, dict): + raise ValueError(f"LoRA {index} is invalid.") + path_value = entry.get("path", "") + if not str(path_value).strip(): + continue + lora_path = validate_path(path_value, f"LoRA {index}", LORA_EXTENSIONS) + try: + strength = float(entry.get("strength", 1)) + except (TypeError, ValueError) as error: + raise ValueError(f"LoRA {index} strength must be a number.") from error + loras.append((lora_path, strength)) + + command = [ + sys.executable, + str(CONVERTER), + "--src", + str(source), + "--dst", + str(destination), + "--quantization-device", + device, + ] + if mode == "target": + command.extend( + [ + "--max-size-mb", + str(target_size), + "--target-size-q8-type", + target_q8_type, + ] + ) + elif quant_type != "source": + command.extend(["--quant-type", quant_type]) + for lora_path, strength in loras: + command.extend(["--lora", str(lora_path), "--lora-strength", str(strength)]) + if payload.get("streamed"): + command.append("--streamed") + + job_id = uuid.uuid4().hex[:8] + job = { + "id": job_id, + "created_at": now(), + "started_at": None, + "finished_at": None, + "status": "queued", + "source": str(source), + "destination": str(destination), + "mode": mode, + "quant_type": quant_type if mode == "quant" else None, + "target_size_mb": target_size if mode == "target" else None, + "target_size_q8_type": target_q8_type if mode == "target" else None, + "device": device, + "streamed": bool(payload.get("streamed")), + "loras": [{"path": str(path), "strength": strength} for path, strength in loras], + "command": command, + "log": [], + "error": None, + "return_code": None, + } + with self.lock: + self.jobs[job_id] = job + self.pending.put(job_id) + return self.public(job) + + def cancel(self, job_id: str) -> dict[str, Any]: + with self.lock: + job = self._job(job_id) + if job["status"] == "queued": + job["status"] = "cancelled" + job["finished_at"] = now() + job["log"].append("Cancelled before conversion started.") + elif job["status"] == "running": + job["status"] = "cancelling" + job["log"].append("Cancellation requested.") + if self.active_job_id == job_id and self.active_process is not None: + self.active_process.terminate() + else: + raise ValueError("Only queued or running conversions can be cancelled.") + return self.public(job) + + def list_jobs(self) -> list[dict[str, Any]]: + with self.lock: + return [self.public(job) for job in reversed(list(self.jobs.values()))] + + def _job(self, job_id: str) -> dict[str, Any]: + try: + return self.jobs[job_id] + except KeyError as error: + raise ValueError("Conversion job was not found.") from error + + def public(self, job: dict[str, Any]) -> dict[str, Any]: + return {key: value for key, value in job.items() if key != "command"} + + def _work(self) -> None: + while True: + job_id = self.pending.get() + with self.lock: + job = self.jobs.get(job_id) + if job is None or job["status"] != "queued": + self.pending.task_done() + continue + job["status"] = "running" + job["started_at"] = now() + job["log"].append("Starting conversion.") + + try: + environment = os.environ.copy() + environment["PYTHONUNBUFFERED"] = "1" + process = subprocess.Popen( + job["command"], + cwd=ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + encoding="utf-8", + errors="replace", + bufsize=1, + env=environment, + ) + with self.lock: + self.active_process = process + self.active_job_id = job_id + assert process.stdout is not None + for line in process.stdout: + with self.lock: + job["log"].append(line.rstrip()) + return_code = process.wait() + with self.lock: + job["return_code"] = return_code + if job["status"] == "cancelling": + job["status"] = "cancelled" + elif return_code == 0: + job["status"] = "completed" + else: + job["status"] = "failed" + job["error"] = f"Converter exited with code {return_code}." + except OSError as error: + with self.lock: + job["status"] = "failed" + job["error"] = f"Could not start converter: {error}" + job["log"].append(job["error"]) + finally: + with self.lock: + job["finished_at"] = now() + self.active_process = None + self.active_job_id = None + self.pending.task_done() + + +MANAGER = ConversionManager() + + +class RequestHandler(BaseHTTPRequestHandler): + server_version = "ComfyUI-GGUF Conversion Dashboard" + + def do_GET(self) -> None: + path = urlparse(self.path).path + if path == "/": + self._send_page() + elif path == "/api/jobs": + self._send_json(HTTPStatus.OK, {"jobs": MANAGER.list_jobs()}) + else: + self._send_json(HTTPStatus.NOT_FOUND, {"error": "Not found."}) + + def do_POST(self) -> None: + path = urlparse(self.path).path + try: + payload = self._read_json() + if path == "/api/jobs": + self._send_json(HTTPStatus.CREATED, {"job": MANAGER.create(payload)}) + elif path.startswith("/api/jobs/") and path.endswith("/cancel"): + job_id = path.removeprefix("/api/jobs/").removesuffix("/cancel").strip("/") + self._send_json(HTTPStatus.OK, {"job": MANAGER.cancel(job_id)}) + else: + self._send_json(HTTPStatus.NOT_FOUND, {"error": "Not found."}) + except ValueError as error: + self._send_json(HTTPStatus.BAD_REQUEST, {"error": str(error)}) + + def _read_json(self) -> dict[str, Any]: + length_header = self.headers.get("Content-Length", "0") + try: + length = int(length_header) + except ValueError as error: + raise ValueError("Invalid request length.") from error + if length <= 0 or length > 131072: + raise ValueError("Request body must be between 1 and 131072 bytes.") + try: + payload = json.loads(self.rfile.read(length)) + except json.JSONDecodeError as error: + raise ValueError("Request body must be valid JSON.") from error + if not isinstance(payload, dict): + raise ValueError("Request body must be a JSON object.") + return payload + + def _send_page(self) -> None: + try: + content = PAGE.read_bytes() + except OSError as error: + self._send_json(HTTPStatus.INTERNAL_SERVER_ERROR, {"error": f"Could not read UI: {error}"}) + return + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.send_header("Content-Length", str(len(content))) + self.end_headers() + self.wfile.write(content) + + def _send_json(self, status: HTTPStatus, payload: dict[str, Any]) -> None: + content = json.dumps(payload).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json; charset=utf-8") + self.send_header("Content-Length", str(len(content))) + self.end_headers() + self.wfile.write(content) + + def log_message(self, format: str, *args: object) -> None: + return + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run the local ComfyUI-GGUF conversion dashboard.") + parser.add_argument("--port", type=int, default=8189, help="Local TCP port to use (default: 8189).") + parser.add_argument("--no-browser", action="store_true", help="Do not open the dashboard automatically.") + args = parser.parse_args() + if not 1 <= args.port <= 65535: + parser.error("--port must be between 1 and 65535.") + return args + + +def main() -> None: + args = parse_args() + address = ("127.0.0.1", args.port) + httpd = ThreadingHTTPServer(address, RequestHandler) + url = f"http://{address[0]}:{address[1]}" + print(f"ComfyUI-GGUF conversion dashboard: {url}") + print("Press Ctrl+C to stop the server.") + if not args.no_browser: + webbrowser.open(url) + try: + httpd.serve_forever() + except KeyboardInterrupt: + print("\nStopping dashboard.") + finally: + httpd.server_close() + + +if __name__ == "__main__": + main() diff --git a/tools/convert.py b/tools/convert.py index 5029c874..6c5eef89 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -1,365 +1,1626 @@ -# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) -import os -import gguf -import torch -import logging -import argparse -from tqdm import tqdm -from safetensors.torch import load_file, save_file - -QUANTIZATION_THRESHOLD = 1024 -REARRANGE_THRESHOLD = 512 -MAX_TENSOR_NAME_LENGTH = 127 -MAX_TENSOR_DIMS = 4 - -class ModelTemplate: - arch = "invalid" # string describing architecture - shape_fix = False # whether to reshape tensors - keys_detect = [] # list of lists to match in state dict - keys_banned = [] # list of keys that should mark model as invalid for conversion - keys_hiprec = [] # list of keys that need to be kept in fp32 for some reason - keys_ignore = [] # list of strings to ignore keys by when found - - def handle_nd_tensor(self, key, data): - raise NotImplementedError(f"Tensor detected that exceeds dims supported by C++ code! ({key} @ {data.shape})") - -class ModelFlux(ModelTemplate): - arch = "flux" - keys_detect = [ - ("transformer_blocks.0.attn.norm_added_k.weight",), - ("double_blocks.0.img_attn.proj.weight",), - ] - keys_banned = ["transformer_blocks.0.attn.norm_added_k.weight",] - -class ModelSD3(ModelTemplate): - arch = "sd3" - keys_detect = [ - ("transformer_blocks.0.attn.add_q_proj.weight",), - ("joint_blocks.0.x_block.attn.qkv.weight",), - ] - keys_banned = ["transformer_blocks.0.attn.add_q_proj.weight",] - -class ModelAura(ModelTemplate): - arch = "aura" - keys_detect = [ - ("double_layers.3.modX.1.weight",), - ("joint_transformer_blocks.3.ff_context.out_projection.weight",), - ] - keys_banned = ["joint_transformer_blocks.3.ff_context.out_projection.weight",] - -class ModelHiDream(ModelTemplate): - arch = "hidream" - keys_detect = [ - ( - "caption_projection.0.linear.weight", - "double_stream_blocks.0.block.ff_i.shared_experts.w3.weight" - ) - ] - keys_hiprec = [ - # nn.parameter, can't load from BF16 ver - ".ff_i.gate.weight", - "img_emb.emb_pos" - ] - -class CosmosPredict2(ModelTemplate): - arch = "cosmos" - keys_detect = [ - ( - "blocks.0.mlp.layer1.weight", - "blocks.0.adaln_modulation_cross_attn.1.weight", - ) - ] - keys_hiprec = ["pos_embedder"] - keys_ignore = ["_extra_state", "accum_"] - -class ModelHyVid(ModelTemplate): - arch = "hyvid" - keys_detect = [ - ( - "double_blocks.0.img_attn_proj.weight", - "txt_in.individual_token_refiner.blocks.1.self_attn_qkv.weight", - ) - ] - - def handle_nd_tensor(self, key, data): - # hacky but don't have any better ideas - path = f"./fix_5d_tensors_{self.arch}.safetensors" # TODO: somehow get a path here?? - if os.path.isfile(path): - raise RuntimeError(f"5D tensor fix file already exists! {path}") - fsd = {key: torch.from_numpy(data)} - tqdm.write(f"5D key found in state dict! Manual fix required! - {key} {data.shape}") - save_file(fsd, path) - -class ModelWan(ModelHyVid): - arch = "wan" - keys_detect = [ - ( - "blocks.0.self_attn.norm_q.weight", - "text_embedding.2.weight", - "head.modulation", - ) - ] - keys_hiprec = [ - ".modulation" # nn.parameter, can't load from BF16 ver - ] - -class ModelLTXV(ModelTemplate): - arch = "ltxv" - keys_detect = [ - ( - "adaln_single.emb.timestep_embedder.linear_2.weight", - "transformer_blocks.27.scale_shift_table", - "caption_projection.linear_2.weight", - ) - ] - keys_hiprec = [ - "scale_shift_table" # nn.parameter, can't load from BF16 base quant - ] - -class ModelSDXL(ModelTemplate): - arch = "sdxl" - shape_fix = True - keys_detect = [ - ("down_blocks.0.downsamplers.0.conv.weight", "add_embedding.linear_1.weight",), - ( - "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", - "output_blocks.2.2.conv.weight", "output_blocks.5.2.conv.weight", - ), # Non-diffusers - ("label_emb.0.0.weight",), - ] - -class ModelSD1(ModelTemplate): - arch = "sd1" - shape_fix = True - keys_detect = [ - ("down_blocks.0.downsamplers.0.conv.weight",), - ( - "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", "input_blocks.9.0.op.weight", - "output_blocks.2.1.conv.weight", "output_blocks.5.2.conv.weight", "output_blocks.8.2.conv.weight" - ), # Non-diffusers - ] - -class ModelLumina2(ModelTemplate): - arch = "lumina2" - keys_detect = [ - ("cap_embedder.1.weight", "context_refiner.0.attention.qkv.weight") - ] - -arch_list = [ModelFlux, ModelSD3, ModelAura, ModelHiDream, CosmosPredict2, - ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1, ModelLumina2] - -def is_model_arch(model, state_dict): - # check if model is correct - matched = False - invalid = False - for match_list in model.keys_detect: - if all(key in state_dict for key in match_list): - matched = True - invalid = any(key in state_dict for key in model.keys_banned) - break - assert not invalid, "Model architecture not allowed for conversion! (i.e. reference VS diffusers format)" - return matched - -def detect_arch(state_dict): - model_arch = None - for arch in arch_list: - if is_model_arch(arch, state_dict): - model_arch = arch() - break - assert model_arch is not None, "Unknown model architecture!" - return model_arch - -def parse_args(): - parser = argparse.ArgumentParser(description="Generate F16 GGUF files from single UNET") - parser.add_argument("--src", required=True, help="Source model ckpt file.") - parser.add_argument("--dst", help="Output unet gguf file.") - args = parser.parse_args() - - if not os.path.isfile(args.src): - parser.error("No input provided!") - - return args - -def strip_prefix(state_dict): - # prefix for mixed state dict - prefix = None - for pfx in ["model.diffusion_model.", "model."]: - if any([x.startswith(pfx) for x in state_dict.keys()]): - prefix = pfx - break - - # prefix for uniform state dict - if prefix is None: - for pfx in ["net."]: - if all([x.startswith(pfx) for x in state_dict.keys()]): - prefix = pfx - break - - # strip prefix if found - if prefix is not None: - logging.info(f"State dict prefix found: '{prefix}'") - sd = {} - for k, v in state_dict.items(): - if prefix not in k: - continue - k = k.replace(prefix, "") - sd[k] = v - else: - logging.debug("State dict has no prefix") - sd = state_dict - - return sd - -def load_state_dict(path): - if any(path.endswith(x) for x in [".ckpt", ".pt", ".bin", ".pth"]): - state_dict = torch.load(path, map_location="cpu", weights_only=True) - for subkey in ["model", "module"]: - if subkey in state_dict: - state_dict = state_dict[subkey] - break - if len(state_dict) < 20: - raise RuntimeError(f"pt subkey load failed: {state_dict.keys()}") - else: - state_dict = load_file(path) - - return strip_prefix(state_dict) - -def handle_tensors(writer, state_dict, model_arch): - name_lengths = tuple(sorted( - ((key, len(key)) for key in state_dict.keys()), - key=lambda item: item[1], - reverse=True, - )) - if not name_lengths: - return - max_name_len = name_lengths[0][1] - if max_name_len > MAX_TENSOR_NAME_LENGTH: - bad_list = ", ".join(f"{key!r} ({namelen})" for key, namelen in name_lengths if namelen > MAX_TENSOR_NAME_LENGTH) - raise ValueError(f"Can only handle tensor names up to {MAX_TENSOR_NAME_LENGTH} characters. Tensors exceeding the limit: {bad_list}") - for key, data in tqdm(state_dict.items()): - old_dtype = data.dtype - - if any(x in key for x in model_arch.keys_ignore): - tqdm.write(f"Filtering ignored key: '{key}'") - continue - - if data.dtype == torch.bfloat16: - data = data.to(torch.float32).numpy() - # this is so we don't break torch 2.0.X - elif data.dtype in [getattr(torch, "float8_e4m3fn", "_invalid"), getattr(torch, "float8_e5m2", "_invalid")]: - data = data.to(torch.float16).numpy() - else: - data = data.numpy() - - n_dims = len(data.shape) - data_shape = data.shape - if old_dtype == torch.bfloat16: - data_qtype = gguf.GGMLQuantizationType.BF16 - # elif old_dtype == torch.float32: - # data_qtype = gguf.GGMLQuantizationType.F32 - else: - data_qtype = gguf.GGMLQuantizationType.F16 - - # The max no. of dimensions that can be handled by the quantization code is 4 - if len(data.shape) > MAX_TENSOR_DIMS: - model_arch.handle_nd_tensor(key, data) - continue # needs to be added back later - - # get number of parameters (AKA elements) in this tensor - n_params = 1 - for dim_size in data_shape: - n_params *= dim_size - - if old_dtype in (torch.float32, torch.bfloat16): - if n_dims == 1: - # one-dimensional tensors should be kept in F32 - # also speeds up inference due to not dequantizing - data_qtype = gguf.GGMLQuantizationType.F32 - - elif n_params <= QUANTIZATION_THRESHOLD: - # very small tensors - data_qtype = gguf.GGMLQuantizationType.F32 - - elif any(x in key for x in model_arch.keys_hiprec): - # tensors that require max precision - data_qtype = gguf.GGMLQuantizationType.F32 - - if (model_arch.shape_fix # NEVER reshape for models such as flux - and n_dims > 1 # Skip one-dimensional tensors - and n_params >= REARRANGE_THRESHOLD # Only rearrange tensors meeting the size requirement - and (n_params / 256).is_integer() # Rearranging only makes sense if total elements is divisible by 256 - and not (data.shape[-1] / 256).is_integer() # Only need to rearrange if the last dimension is not divisible by 256 - ): - orig_shape = data.shape - data = data.reshape(n_params // 256, 256) - writer.add_array(f"comfy.gguf.orig_shape.{key}", tuple(int(dim) for dim in orig_shape)) - - try: - data = gguf.quants.quantize(data, data_qtype) - except (AttributeError, gguf.QuantError) as e: - tqdm.write(f"falling back to F16: {e}") - data_qtype = gguf.GGMLQuantizationType.F16 - data = gguf.quants.quantize(data, data_qtype) - - new_name = key # do we need to rename? - - shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}" - tqdm.write(f"{f'%-{max_name_len + 4}s' % f'{new_name}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}") - - writer.add_tensor(new_name, data, raw_dtype=data_qtype) - -def convert_file(path, dst_path=None, interact=True, overwrite=False): - # load & run model detection logic - state_dict = load_state_dict(path) - model_arch = detect_arch(state_dict) - logging.info(f"* Architecture detected from input: {model_arch.arch}") - - # detect & set dtype for output file - dtypes = [x.dtype for x in state_dict.values()] - dtypes = {x:dtypes.count(x) for x in set(dtypes)} - main_dtype = max(dtypes, key=dtypes.get) - - if main_dtype == torch.bfloat16: - ftype_name = "BF16" - ftype_gguf = gguf.LlamaFileType.MOSTLY_BF16 - # elif main_dtype == torch.float32: - # ftype_name = "F32" - # ftype_gguf = None - else: - ftype_name = "F16" - ftype_gguf = gguf.LlamaFileType.MOSTLY_F16 - - if dst_path is None: - dst_path = f"{os.path.splitext(path)[0]}-{ftype_name}.gguf" - elif "{ftype}" in dst_path: # lcpp logic - dst_path = dst_path.replace("{ftype}", ftype_name) - - if os.path.isfile(dst_path) and not overwrite: - if interact: - input("Output exists enter to continue or ctrl+c to abort!") - else: - raise OSError("Output exists and overwriting is disabled!") - - # handle actual file - writer = gguf.GGUFWriter(path=None, arch=model_arch.arch) - writer.add_quantization_version(gguf.GGML_QUANT_VERSION) - if ftype_gguf is not None: - writer.add_file_type(ftype_gguf) - - handle_tensors(writer, state_dict, model_arch) - writer.write_header_to_file(path=dst_path) - writer.write_kv_data_to_file() - writer.write_tensors_to_file(progress=True) - writer.close() - - fix = f"./fix_5d_tensors_{model_arch.arch}.safetensors" - if os.path.isfile(fix): - logging.warning(f"\n### Warning! Fix file found at '{fix}'") - logging.warning(" you most likely need to run 'fix_5d_tensors.py' after quantization.") - - return dst_path, model_arch - -if __name__ == "__main__": - args = parse_args() - convert_file(args.src, args.dst) - +# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import os +import math +import gguf +import json +import numpy as np +import torch +import logging +import argparse +import sys +import tempfile +from collections import OrderedDict +from fnmatch import fnmatchcase +from tqdm import tqdm +from safetensors import safe_open +from safetensors.torch import save_file + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from lora import ( + fuse_target_entries_into_tensor, + fuse_targets_into_state_dict, + load_lora, + materialize_int8_source_weights, + resolve_fusion_targets, +) + +QUANTIZATION_THRESHOLD = 1024 +REARRANGE_THRESHOLD = 512 +MAX_TENSOR_NAME_LENGTH = 127 +MAX_TENSOR_DIMS = 4 +_FP8_DTYPES = { + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e5m2", None), +} - {None} +RAW_BYTE_TENSOR_KEYS = frozenset(("tokenizer_json", "spiece_model", "tekken_model")) + +class ModelTemplate: + arch = "invalid" # string describing architecture + shape_fix = False # whether to reshape tensors + preserve_nd_shapes = False + keys_detect = [] # list of lists to match in state dict + keys_banned = [] # list of keys that should mark model as invalid for conversion + keys_hiprec = [] # list of keys that need to be kept in fp32 for some reason + keys_noquant = [] # list of keys that must retain their source precision + keys_ignore = [] # list of strings to ignore keys by when found + + def handle_nd_tensor(self, key, data): + raise NotImplementedError(f"Tensor detected that exceeds dims supported by C++ code! ({key} @ {data.shape})") + +def key_matches(key, patterns): + """ + Match a tensor name against a list of patterns. + + Plain patterns match anywhere in the key (today's behavior, unchanged -- + e.g. "pos_embedder", "scale_shift_table", ".modulation"). + + A pattern prefixed with '^' matches only at the START of the key, e.g. + "^tmlp." matches "tmlp.0.weight" but NOT "blocks.5.txtmlp.0.weight" -- + use this for short/generic fragments that would otherwise collide with + an unrelated, similarly-named submodule elsewhere in the tensor name + (see ModelKrea2: bare "tmlp."/"tproj." also matched every per-block + "txtmlp."/"txtproj." tensor, silently forcing far more of the model to + F32 than intended). Patterns containing '*' or '?' use shell-style + wildcards, which can express a path segment without matching nested + submodules. + """ + for pattern in patterns: + if "*" in pattern or "?" in pattern: + if fnmatchcase(key, pattern[1:] if pattern.startswith("^") else pattern): + return True + continue + if pattern.startswith("^"): + if key.startswith(pattern[1:]): + return True + elif pattern in key: + return True + return False + +class ModelFlux(ModelTemplate): + arch = "flux" + keys_detect = [ + ("transformer_blocks.0.attn.norm_added_k.weight",), + ("double_blocks.0.img_attn.proj.weight",), + ] + keys_banned = ["transformer_blocks.0.attn.norm_added_k.weight",] + +class ModelSD3(ModelTemplate): + arch = "sd3" + keys_detect = [ + ("transformer_blocks.0.attn.add_q_proj.weight",), + ("joint_blocks.0.x_block.attn.qkv.weight",), + ] + keys_banned = ["transformer_blocks.0.attn.add_q_proj.weight",] + +class ModelAura(ModelTemplate): + arch = "aura" + keys_detect = [ + ("double_layers.3.modX.1.weight",), + ("joint_transformer_blocks.3.ff_context.out_projection.weight",), + ] + keys_banned = ["joint_transformer_blocks.3.ff_context.out_projection.weight",] + +class ModelHiDream(ModelTemplate): + arch = "hidream" + keys_detect = [ + ( + "caption_projection.0.linear.weight", + "double_stream_blocks.0.block.ff_i.shared_experts.w3.weight" + ) + ] + keys_hiprec = [ + # nn.parameter, can't load from BF16 ver + ".ff_i.gate.weight", + "img_emb.emb_pos" + ] + +class CosmosPredict2(ModelTemplate): + arch = "cosmos" + keys_detect = [ + ( + "blocks.0.mlp.layer1.weight", + "blocks.0.adaln_modulation_cross_attn.1.weight", + ) + ] + keys_hiprec = ["pos_embedder"] + keys_ignore = ["_extra_state", "accum_"] + +class ModelHyVid(ModelTemplate): + arch = "hyvid" + keys_detect = [ + ( + "double_blocks.0.img_attn_proj.weight", + "txt_in.individual_token_refiner.blocks.1.self_attn_qkv.weight", + ) + ] + + def handle_nd_tensor(self, key, data): + # hacky but don't have any better ideas + path = f"./fix_5d_tensors_{self.arch}.safetensors" # TODO: somehow get a path here?? + if os.path.isfile(path): + raise RuntimeError(f"5D tensor fix file already exists! {path}") + fsd = {key: torch.from_numpy(data)} + tqdm.write(f"5D key found in state dict! Manual fix required! - {key} {data.shape}") + save_file(fsd, path) + +class ModelWan(ModelHyVid): + arch = "wan" + keys_detect = [ + ( + "blocks.0.self_attn.norm_q.weight", + "text_embedding.2.weight", + "head.modulation", + ) + ] + keys_hiprec = [ + ".modulation" # nn.parameter, can't load from BF16 ver + ] + +class ModelLTXV(ModelTemplate): + arch = "ltxv" + keys_detect = [ + ( + "adaln_single.emb.timestep_embedder.linear_2.weight", + "transformer_blocks.27.scale_shift_table", + "caption_projection.linear_2.weight", + ), + # LTX 2.3 audio-video checkpoints replace the video-only caption + # projection with audio/video connector modules. + ( + "adaln_single.emb.timestep_embedder.linear_2.weight", + "transformer_blocks.27.scale_shift_table", + "audio_adaln_single.linear.weight", + ), + ] + keys_hiprec = [ + "scale_shift_table", # nn.Parameter, can't load from BF16 base quant + "learnable_registers", # Connector nn.Parameter, not a Linear weight + ] + # LTX's native INT8 ConvRot checkpoints intentionally leave these + # projections in BF16. The many 32-output gate logits are particularly + # small, so ConvRot dispatch overhead outweighs their INT8 benefit. + keys_noquant = [ + "adaln_single", + "patchify_proj", + "proj_out", + "to_gate_logits", + ] + +class ModelLTXVUpsampler(ModelTemplate): + arch = "ltxv_upscaler" + preserve_nd_shapes = True + keys_detect = [ + ( + "initial_conv.weight", + "post_upsample_res_blocks.0.conv2.bias", + "upsampler.0.weight", + "final_conv.weight", + ) + ] + # LTX 2.5 latent upscalers are entirely convolutional. No GGUF runtime + # quantized convolution path exists, so retain the source precision. + keys_noquant = ["^"] + +class ModelSDXL(ModelTemplate): + arch = "sdxl" + shape_fix = True + keys_detect = [ + ("down_blocks.0.downsamplers.0.conv.weight", "add_embedding.linear_1.weight",), + ( + "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", + "output_blocks.2.2.conv.weight", "output_blocks.5.2.conv.weight", + ), # Non-diffusers + ("label_emb.0.0.weight",), + ] + +class ModelSD1(ModelTemplate): + arch = "sd1" + shape_fix = True + keys_detect = [ + ("down_blocks.0.downsamplers.0.conv.weight",), + ( + "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", "input_blocks.9.0.op.weight", + "output_blocks.2.1.conv.weight", "output_blocks.5.2.conv.weight", "output_blocks.8.2.conv.weight" + ), # Non-diffusers + ] + +class ModelLumina2(ModelTemplate): + arch = "lumina2" + keys_detect = [ + ("cap_embedder.1.weight", "context_refiner.0.attention.qkv.weight") + ] + +class ModelIdeogram(ModelTemplate): + arch = "ideogram" + keys_detect = [ + ( + "t_embedding.mlp_in.weight", + "layers.0.attention.qkv.weight", + "final_layer.linear.weight", + ) + ] + +class ModelKrea2(ModelTemplate): + """ + Krea-2 is a novel architecture from krea.ai — NOT Ideogram4. + Key structure (verified from Krea2_Turbo_fp8mixed.safetensors header): + blocks.N.attn.{wq,wk,wv,wo,gate} — separate Q/K/V/O projections + gating + blocks.N.attn.qknorm.{qnorm,knorm} — Q/K norms + blocks.N.mlp.{up,gate,down} — SwiGLU-style MLP + blocks.N.mod.lin — per-block modulation + blocks.N.{pre,post}norm.scale — RMSNorm scales + txtfusion.layerwise_blocks.N.* — layerwise text-image cross-attention + txtfusion.refiner_blocks.N.* — refiner text-image cross-attention + txtfusion.projector — text projector + first.weight / last.* — input / output projections + tmlp.N / tproj.N — timestep MLP / projection + + NOTE: ComfyUI core must have krea2 diffusion model support + (i.e. a detection branch for 'blocks.0.attn.wq.weight' in model_detection.py + and a matching supported_models entry) for the GGUF to load correctly. + Krea-2 was released 2026-06-22; verify your ComfyUI build is up to date. + """ + arch = "krea2" + keys_detect = [ + ( + "blocks.0.attn.wq.weight", + "txtfusion.projector.weight", + "first.weight", + ), + ( + "blocks.0.attn.wq.weight", + "txtfusion.layerwise_blocks.0.attn.wq.weight", + "last.linear.weight", + ), + ] + keys_hiprec = [ + "^first.", + "^last.", + "^tproj.", + "^tmlp.", + "^txtmlp.", + "^txtfusion.projector.", + ] + + +class ModelMinimaxH3(ModelTemplate): + arch = "minimax_h3" + keys_detect = [ + ( + "video_patch_proj.weight", + "audio_patch_proj.weight", + "blocks.0.attn.qkv_proj.weight", + "final_layer.video_out.weight", + ) + ] + # The timestep table and final projections are used directly by the + # conditioning/output paths, rather than as ordinary transformer Linear + # weights. Keep them in FP32 for numerical stability. + keys_hiprec = [ + "adaln_t_table", + "^video_patch_proj.", + "^audio_patch_proj.", + "^final_layer.", + "adaln_proj", + "modulation", + "^blocks.*.attn.qkv_proj.", + "^blocks.*.attn.out_proj.", + "^blocks.*.mlp.fc2.", + ] + # These paths are already BF16 in the reference checkpoint. Quantizing the + # conditioning projection or token refiner to W4A4 costs quality while + # saving little compared with the 50 main transformer blocks. + keys_noquant = [ + "^condition_proj.", + "^token_refiner.", + "norm", + ] + + +class ModelMinimaxH3VAE(ModelTemplate): + arch = "minimax_h3_vae" + preserve_nd_shapes = True + keys_detect = [ + ( + "decoder.transformer_blocks.0.scale1", + "decoder.x_embedder.weight", + "encoder.down.5.block.0.conv1.weight", + ) + ] + + +class ModelMiniMaxMusic3DiT(ModelTemplate): + arch = "minimax_music3" + keys_detect = [ + ( + "cond_layer_logits", + "latent_conditioners.0.weight", + "diffusion_transformer.transformer.layers.0.self_attn.to_qkv.weight", + "diffusion_transformer.transformer.project_in.weight", + ) + ] + # These are Fourier/rotary buffers and 1x1 convolutional paths. The + # MiniMax Music3 runtime is FP32 and has no low-bit convolution kernel. + keys_hiprec = [ + "cond_layer_", + "latent_conditioners.", + "preprocess_conv.", + "postprocess_conv.", + "timestep_features", + "rotary_pos_emb", + ] + + +class ModelMiniMaxMusic3TextEncoder(ModelTemplate): + arch = "minimax_music3" + keys_detect = [ + ( + "model.embed_tokens_prefill.weight", + "model.embed_tokens_audio.weight", + "model.lm_head_pruned.weight", + "model.audio_decoder.audio_heads.0.weight", + "model.layers.0.self_attn.qkv_proj.weight", + ) + ] + # The native Q8_CR path accelerates Linear only. Keep lookup tables in + # BF16 so ComfyUI's Embedding operations retain their normal behavior. + keys_noquant = [ + "embed_tokens_prefill", + "embed_tokens_audio", + "audio_extra_embedding", + "audio_decoder.pos_embedding", + ] + + +arch_list = [ModelFlux, ModelSD3, ModelAura, ModelHiDream, CosmosPredict2, + ModelLTXV, ModelLTXVUpsampler, ModelHyVid, ModelWan, ModelSDXL, ModelSD1, ModelLumina2, + ModelKrea2, ModelIdeogram, ModelMinimaxH3, ModelMinimaxH3VAE, + ModelMiniMaxMusic3DiT, ModelMiniMaxMusic3TextEncoder] + +def is_model_arch(model, state_dict): + # check if model is correct + matched = False + invalid = False + for match_idx, match_list in enumerate(model.keys_detect): + if all(key in state_dict for key in match_list): + matched = True + invalid = any(key in state_dict for key in model.keys_banned) + if len(model.keys_detect) > 1: + # Multiple detect variants usually mean multiple known checkpoint + # exports of the same architecture (e.g. different key subsets + # across releases). Logging which one matched makes it obvious + # which variant you're actually converting. + logging.info(f"* Matched keys_detect variant #{match_idx} for '{model.arch}': {match_list}") + break + assert not invalid, "Model architecture not allowed for conversion! (i.e. reference VS diffusers format)" + return matched + +def detect_arch(state_dict): + model_arch = None + for arch in arch_list: + if is_model_arch(arch, state_dict): + model_arch = arch() + break + assert model_arch is not None, "Unknown model architecture!" + return model_arch + +def validate_key_patterns(model_arch, state_dict): + """ + Warn if a configured key pattern (keys_hiprec / keys_noquant / keys_ignore) + matches no tensor at all in this checkpoint. + + These patterns are substring matches against tensor names (`any(x in key ...)`), + hand-written against one specific checkpoint export. If the upstream model + renames a layer in a later release, the pattern silently stops firing -- + no error, no warning, just quietly reduced precision/behavior on tensors + that were meant to be protected. This check surfaces that case early, + at conversion time, instead of relying on someone noticing degraded output + later. + + This is intentionally a warning, not an assert: a pattern legitimately + matching nothing can happen for known reasons too, e.g. a checkpoint + variant that simply doesn't include that sub-module (see ModelKrea2's + two keys_detect alternatives, which exist for exactly this reason). + """ + for attr in ("keys_hiprec", "keys_noquant", "keys_ignore"): + for pattern in getattr(model_arch, attr, []): + if not any(key_matches(key, [pattern]) for key in state_dict.keys()): + logging.warning( + f"[{model_arch.arch}] '{pattern}' in {attr} matched no tensor in this " + f"checkpoint -- possible naming drift in the source model, or an " + f"intentionally absent sub-module for this checkpoint variant. " + f"Verify this is expected before trusting the output precision." + ) + +QUANT_TYPE_MAP = { + "F16": (gguf.GGMLQuantizationType.F16, gguf.LlamaFileType.MOSTLY_F16), + "BF16": (gguf.GGMLQuantizationType.BF16, gguf.LlamaFileType.MOSTLY_BF16), + "Q8_0": (gguf.GGMLQuantizationType.Q8_0, gguf.LlamaFileType.MOSTLY_Q8_0), + "Q5_1": (gguf.GGMLQuantizationType.Q5_1, gguf.LlamaFileType.MOSTLY_Q5_1), + "Q5_0": (gguf.GGMLQuantizationType.Q5_0, gguf.LlamaFileType.MOSTLY_Q5_0), + "Q4_1": (gguf.GGMLQuantizationType.Q4_1, gguf.LlamaFileType.MOSTLY_Q4_1), + "Q4_0": (gguf.GGMLQuantizationType.Q4_0, gguf.LlamaFileType.MOSTLY_Q4_0), + "Q8_CR": (gguf.GGMLQuantizationType.I8, None), # INT8 ConvRot (ComfyUI native) + # Q4_CR_W4A4 is a custom W4A4 INT4 format backed by comfy_kitchen's fast + # ConvRot int4 tensor-core MMA. Stored as kitchen-native packed int4 (N, K//2) + # + per-output-row fp scales (no per-group scales/zeros). + "Q4_CR_W4A4": (gguf.GGMLQuantizationType.I8, None), + # Q4_PT is retired pending a performant Ampere W4A16 backend. + # "Q4_PT": (gguf.GGMLQuantizationType.I8, None), +} + +TARGET_SIZE_QUANT_TYPE = "TARGET_SIZE" +TARGET_SIZE_Q8_TYPES = ("Q8_CR", "Q8_0") +DEFAULT_TARGET_SIZE_Q8_TYPE = "Q8_CR" +Q4_CR_W4A4_CONVROT_GROUP_SIZE = 256 +Q4_CR_W4A4_QUANT_GROUP_SIZE = 64 +QUANTIZATION_DEVICE_OPTIONS = ("auto", "cpu", "cuda") +MEBIBYTE = 1024 * 1024 + + +def _tensor_size_bytes(shape, quant_type): + """Return the GGUF payload size for a tensor with the given quantization.""" + n_params = 1 + for dim_size in shape: + n_params *= dim_size + + if quant_type == gguf.GGMLQuantizationType.I8: + return n_params + + block_size, type_size = gguf.constants.GGML_QUANT_SIZES[quant_type] + if n_params % block_size: + raise ValueError( + f"{quant_type.name} requires a tensor size divisible by {block_size}, " + f"got shape {tuple(shape)} ({n_params} elements)." + ) + return n_params // block_size * type_size + + +def _default_qtype(data_or_dtype): + return ( + gguf.GGMLQuantizationType.BF16 + if getattr(data_or_dtype, "dtype", data_or_dtype) == torch.bfloat16 + else gguf.GGMLQuantizationType.F16 + ) + + +def _is_raw_byte_tensor(key, data_or_dtype, ndim=None): + dtype = getattr(data_or_dtype, "dtype", data_or_dtype) + if ndim is None: + ndim = len(data_or_dtype.shape) + return key in RAW_BYTE_TENSOR_KEYS and ndim == 1 and dtype in (torch.uint8, torch.int8) + + +def _is_target_core_tensor(key, data, model_arch): + if len(data.shape) != 2: + return False + if data.numel() <= QUANTIZATION_THRESHOLD: + return False + if key_matches(key, model_arch.keys_hiprec) or key_matches(key, model_arch.keys_noquant): + return False + return True + + +def _can_use_q4_0(data): + block_size, _ = gguf.constants.GGML_QUANT_SIZES[gguf.GGMLQuantizationType.Q4_0] + return data.shape[-1] % block_size == 0 + + +def plan_target_size_quantization( + state_dict, + model_arch, + max_size_mb, + target_size_q8_type=DEFAULT_TARGET_SIZE_Q8_TYPE, +): + """ + Select per-tensor types that fit a maximum serialized payload size. + + Core 2-D tensors start in the selected Q8 type. The center of their + checkpoint order is downgraded to Q5_0 first, then to Q4_0 only when + needed, leaving the beginning and end at higher precision as long as + possible. + Once every Q4-compatible core tensor is Q4_0, ordinary 1-D tensors may be + reduced to BF16. Protected tensors always remain F32. + """ + if max_size_mb <= 0: + raise ValueError("--max-size-mb must be greater than zero.") + if target_size_q8_type not in TARGET_SIZE_Q8_TYPES: + raise ValueError( + f"--target-size-q8-type must be one of {', '.join(TARGET_SIZE_Q8_TYPES)}, " + f"got {target_size_q8_type!r}." + ) + target_q8_type = QUANT_TYPE_MAP[target_size_q8_type][0] + + plan = {} + core_tensors = [] + one_dimensional_tensors = [] + + for key, data in state_dict.items(): + if key_matches(key, model_arch.keys_ignore): + continue + if key.endswith(".comfy_quant") or key.endswith("_scale") and len(data.shape) == 0: + continue + if len(data.shape) == 0 or len(data.shape) > MAX_TENSOR_DIMS: + continue + + n_params = data.numel() + if _is_raw_byte_tensor(key, data): + # GGUF has no U8 tensor type. I8 is used as a byte container and + # the loader restores the unsigned view without changing bits. + plan[key] = gguf.GGMLQuantizationType.I8 + elif len(data.shape) == 1 or n_params <= QUANTIZATION_THRESHOLD or key_matches(key, model_arch.keys_hiprec): + plan[key] = gguf.GGMLQuantizationType.F32 + if len(data.shape) == 1 and not key_matches(key, model_arch.keys_hiprec): + one_dimensional_tensors.append((key, data)) + elif key_matches(key, model_arch.keys_noquant): + plan[key] = _default_qtype(data) + elif len(data.shape) == 4 and "conv" in key.lower(): + plan[key] = gguf.GGMLQuantizationType.F16 + elif _is_target_core_tensor(key, data, model_arch): + plan[key] = target_q8_type + if _can_use_q4_0(data): + core_tensors.append((key, data)) + else: + plan[key] = _default_qtype(data) + + def plan_size(): + total = 0 + for key, data in state_dict.items(): + if key not in plan: + continue + qtype = plan[key] + total += _tensor_size_bytes(data.shape, qtype) + if qtype == gguf.GGMLQuantizationType.I8 and not _is_raw_byte_tensor(key, data): + # Q8_CR stores a F32 scale for every output row. + total += data.shape[0] * 4 + return total + + target_size = int(max_size_mb * MEBIBYTE) + maximum_size = plan_size() + if maximum_size <= target_size: + return plan, maximum_size, maximum_size + + center = (len(core_tensors) - 1) / 2 + center_first_core_tensors = [ + (key, data) + for _, (key, data) in sorted( + enumerate(core_tensors), + key=lambda item: (abs(item[0] - center), item[0]), + ) + ] + for target_qtype in ( + gguf.GGMLQuantizationType.Q5_0, + gguf.GGMLQuantizationType.Q4_0, + ): + for key, _ in center_first_core_tensors: + plan[key] = target_qtype + current_size = plan_size() + if current_size <= target_size: + return plan, maximum_size, current_size + + for key, _ in one_dimensional_tensors: + plan[key] = gguf.GGMLQuantizationType.BF16 + current_size = plan_size() + if current_size <= target_size: + return plan, maximum_size, current_size + + minimum_size = plan_size() + raise ValueError( + f"Cannot shrink this model to {max_size_mb:g} MiB. " + f"The smallest supported TARGET_SIZE output is {minimum_size / MEBIBYTE:.2f} MiB " + f"(all Q4_0-compatible core matrices at Q4_0 and ordinary 1-D tensors at BF16). " + "Q3 and lower quantization are not supported." + ) + + +def _validate_quantization_device(device): + if device not in QUANTIZATION_DEVICE_OPTIONS: + raise ValueError( + f"--quantization-device must be one of {', '.join(QUANTIZATION_DEVICE_OPTIONS)}, " + f"got {device!r}." + ) + + +def resolve_quantization_device(device): + _validate_quantization_device(device) + if device == "cpu": + return torch.device("cpu") + if not torch.cuda.is_available(): + if device == "cuda": + raise RuntimeError("--quantization-device cuda requires an available CUDA device.") + return torch.device("cpu") + return torch.device("cuda") + + +def _can_use_cuda_q8_cr(data, device): + # ConvRot needs the uploaded source plus F32 rotation and quantization workspaces. + required_bytes = data.numel() * 16 + data.shape[0] * 4 + free_bytes, _ = torch.cuda.mem_get_info(device) + return required_bytes <= free_bytes + + +def quantize_int8_convrot(weight, convrot_groupsize=256, device=None): + """ + Quantize a 2D Linear weight to INT8 with ConvRot grouping. + Uses per-output-channel scales to match ComfyUI's TensorWiseINT8Layout. + """ + if device is not None: + weight = weight.to(device) + weight = weight.to(torch.float32) + orig_shape = tuple(weight.shape) + groupsize = next( + ( + size + for size in (convrot_groupsize, 64, 16, 4) + if size <= weight.shape[1] and weight.shape[1] % size == 0 + ), + None, + ) + if groupsize is not None: + from comfy_kitchen.tensor.int8_utils import _build_hadamard, _rotate_weight + + hadamard = _build_hadamard(groupsize, device=weight.device, dtype=weight.dtype) + weight = _rotate_weight(weight, hadamard, groupsize) + + scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) / 127.0 + qdata = (weight / scale).round().clamp(-128, 127).to(torch.int8) + quant_conf = { + "format": "int8_tensorwise", + "convrot": groupsize is not None, + "weight_rotated": groupsize is not None, + "per_row": True, + } + if groupsize is not None: + quant_conf["convrot_groupsize"] = groupsize + return qdata, scale, quant_conf, orig_shape + + +def quantize_int4_cr_w4a4( + weight, + convrot_groupsize=256, + quant_group_size=64, + device=None, + dtype=torch.bfloat16, +): + """ + Quantize a 2D Linear weight to the ConvRot W4A4 path (Q4_CR_W4A4, backed by + comfy_kitchen's fast int4 tensor-core MMA). + + Storage matches comfy_kitchen's TensorCoreConvRotW4A4Layout contract so the + loader can rebuild a QuantizedTensor without re-packing: + qweight: (N, K//2) int8 two signed int4 per byte, row-major + bits 0..3 -> column 2j (low nibble) + bits 4..7 -> column 2j+1 (high nibble) + wscales: (N,) float32 per-output-row symmetric scale + + The weight is rotated by a block-diagonal regular Hadamard (group + ``convrot_groupsize``) along K before quantization, matching the ConvRot + activation rotation the kernel applies at runtime. Dequant (after the + runtime rotates the activation back) is symmetric about zero: 7-bit signed + emission range [-7, 7] with scale = absmax / 7. + """ + if device is not None: + weight = weight.to(device) + if quant_group_size != 64: + raise ValueError("Q4_CR W4A4 requires quant_group_size 64 (int4 MMA kernel contract).") + weight = weight.to(torch.float32) + orig_shape = tuple(weight.shape) + n, k = orig_shape + + if k % convrot_groupsize != 0: + raise ValueError( + f"Q4_CR W4A4 convrot group size {convrot_groupsize} must divide input " + f"features {k} for tensor {orig_shape}." + ) + + h = _build_regular_hadamard(convrot_groupsize, dtype=torch.float32, device=weight.device) + n_groups = k // convrot_groupsize + weight_grouped = weight.reshape(n, n_groups, convrot_groupsize) + weight_rotated = torch.matmul(weight_grouped, h.T).reshape(n, k).float() + + absmax = weight_rotated.abs().amax(dim=1, keepdim=True).clamp_min(1e-10) + scales = absmax / 7.0 + q = (weight_rotated / scales).round().clamp(-7, 7).to(torch.int8) + + q_flat = q.to(torch.int32) + lo = q_flat[:, 0::2] & 0x0F + hi = q_flat[:, 1::2] & 0x0F + packed = (lo | (hi << 4)).to(torch.int8) + + wscales = scales.reshape(n).to(torch.float32) + + quant_conf = { + "format": "int4_cr", + "backing": "w4a4", + "convrot_groupsize": convrot_groupsize, + "quant_group_size": quant_group_size, + "orig_shape": orig_shape, + "sym": True, + } + return packed, wscales, quant_conf, orig_shape + + +def _build_regular_hadamard(size, dtype=torch.float32, device="cpu"): + """Build a normalized regular (Sylvester) Hadamard of a power-of-4 size.""" + if size < 4 or (size & (size - 1)) != 0 or not math.log(size, 4).is_integer(): + raise ValueError(f"Regular Hadamard size must be a power of 4, got {size}") + h4 = torch.tensor( + [[1, 1, 1, -1], [1, 1, -1, 1], [1, -1, 1, 1], [-1, 1, 1, 1]], + dtype=dtype, + device=device, + ) + h = h4 + current_size = 4 + while current_size < size: + h = torch.kron(h, h4) + current_size *= 4 + return h / (size ** 0.5) + + +def retired_quantize_int4_pytorch(weight, group_size=64): + """ + Quantize a 2D Linear weight for PyTorch's native INT4 kernel. + Weights are serialized as packed uint8 [n, k//2]. The runtime converts this + portable representation to PyTorch's device-specific INT4 layout. + """ + orig_shape = tuple(weight.shape) + n, k = orig_shape + weight = weight.to(torch.float32) + + pad = 0 + if k % group_size != 0: + pad = group_size - (k % group_size) + weight = torch.nn.functional.pad(weight, (0, pad)) + k = k + pad + + w_grouped = weight.reshape(n, k // group_size, group_size) + w_min = w_grouped.amin(dim=-1, keepdim=True) + w_max = w_grouped.amax(dim=-1, keepdim=True) + scale = (w_max - w_min) / 15.0 + scale = scale.clamp_min(1e-9) + q = ((w_grouped - w_min) / scale).round().clamp(0, 15).to(torch.uint8) + + q_flat = q.reshape(n, k) + packed = ((q_flat[:, 0::2] << 4) | q_flat[:, 1::2]).to(torch.uint8) + + qsz = torch.zeros(k // group_size, n, 2, dtype=torch.float32) + qsz[:, :, 0] = scale.reshape(n, k // group_size).t() + # _weight_int4pack_mm dequantizes as (q - 8) * scale + offset. + qsz[:, :, 1] = (w_min + 8 * scale).reshape(n, k // group_size).t() + + quant_conf = { + "format": "int4_compact_gemm", + "group_size": group_size, + "orig_shape": orig_shape, + "pad": pad, + } + return packed, qsz, quant_conf, orig_shape + +def parse_args(): + parser = argparse.ArgumentParser( + description="Convert diffusion model safetensors/ckpt to GGUF." + " By default produces an F16/BF16 GGUF; use --quant-type to quantize." + ) + parser.add_argument("--src", required=True, help="Source model ckpt/safetensors file.") + parser.add_argument("--dst", help="Output GGUF file path.") + parser.add_argument( + "--lora", + action="append", + default=[], + metavar="PATH", + help="LoRA .safetensors or .gguf adapter to merge before export. Repeat for multiple adapters.", + ) + parser.add_argument( + "--lora-strength", + action="append", + type=float, + default=[], + metavar="VALUE", + help="Merge strength for each --lora, in the same order. Defaults to 1.0 for every adapter.", + ) + parser.add_argument( + "--streamed", + action="store_true", + help=( + "Stream safetensors tensors through quantization and continuously flush " + "disk-backed GGUF payload staging to reduce RAM use." + ), + ) + parser.add_argument( + "--quant-type", + choices=list(QUANT_TYPE_MAP.keys()), + default=None, + help="Target quantization type for eligible 2-D+ tensors " + "(1-D biases/scales stay F32). Defaults to F16/BF16 matching the source dtype.", + ) + parser.add_argument( + "--max-size-mb", + type=float, + default=None, + help=( + "Maximum output payload size in MiB. Selects TARGET_SIZE quantization: " + "selected-Q8 core weights are progressively changed to Q5_0 then Q4_0 " + "from the model center outward, then ordinary 1-D tensors to BF16 if necessary." + ), + ) + parser.add_argument( + "--target-size-q8-type", + choices=TARGET_SIZE_Q8_TYPES, + default=DEFAULT_TARGET_SIZE_Q8_TYPE, + help=( + "Q8 representation used by --max-size-mb before core matrices are reduced to Q4_0. " + "Q8_CR uses native INT8 ConvRot; Q8_0 uses standard GGUF Q8." + ), + ) + parser.add_argument( + "--quantization-device", + choices=QUANTIZATION_DEVICE_OPTIONS, + default="auto", + help=( + "Device for Q8_CR conversion. auto uses CUDA when available; CPU remains " + "the fallback for individual matrices that cannot fit in available VRAM." + ), + ) + args = parser.parse_args() + + if not os.path.isfile(args.src): + parser.error("No input provided!") + + return args + +def strip_prefix(state_dict): + # MiniMax Music3's standalone text encoder owns the `model.` namespace. + # It is not a ComfyUI wrapper prefix: the upstream runtime loads + # `model.embed_tokens_*`, `model.layers.*`, and `model.audio_decoder.*`. + minimax_music3_text_encoder = all( + key in state_dict + for key in ( + "model.embed_tokens_prefill.weight", + "model.embed_tokens_audio.weight", + "model.lm_head_pruned.weight", + "model.audio_decoder.audio_heads.0.weight", + ) + ) + + # prefix for mixed state dict + prefix = None + for pfx in ["model.diffusion_model.", "model."]: + if pfx == "model." and minimax_music3_text_encoder: + continue + if any([x.startswith(pfx) for x in state_dict.keys()]): + prefix = pfx + break + + # prefix for uniform state dict + if prefix is None: + for pfx in ["net."]: + if all([x.startswith(pfx) for x in state_dict.keys()]): + prefix = pfx + break + + # strip prefix if found + if prefix is not None: + logging.info(f"State dict prefix found: '{prefix}'") + sd = {} + for k, v in state_dict.items(): + if prefix not in k: + continue + k = k.replace(prefix, "") + sd[k] = v + else: + logging.debug("State dict has no prefix") + sd = state_dict + + return sd + +def load_state_dict(path, progress_callback=None): + if any(path.endswith(x) for x in [".ckpt", ".pt", ".bin", ".pth"]): + with tqdm(total=1, desc="Reading checkpoint", unit="file") as progress: + state_dict = torch.load(path, map_location="cpu", weights_only=True) + progress.update() + if progress_callback is not None: + progress_callback("read", 1, 1) + for subkey in ["model", "module"]: + if subkey in state_dict: + state_dict = state_dict[subkey] + break + if len(state_dict) < 20: + raise RuntimeError(f"pt subkey load failed: {state_dict.keys()}") + else: + with safe_open(path, framework="pt", device="cpu") as checkpoint: + keys = list(checkpoint.keys()) + state_dict = {} + for index, key in enumerate(tqdm(keys, desc="Reading tensors", unit="tensor"), start=1): + state_dict[key] = checkpoint.get_tensor(key) + if progress_callback is not None: + progress_callback("read", index, len(keys)) + + return strip_prefix(state_dict) + +def load_safetensors_metadata(path): + if not path.endswith(".safetensors"): + return {} + with safe_open(path, framework="pt", device="cpu") as checkpoint: + return checkpoint.metadata() or {} + +def handle_tensors( + writer, + state_dict, + model_arch, + quant_type=None, + quant_type_name=None, + quantization_plan=None, + quantization_device="auto", + progress_callback=None, + fp8_scales=None, + progress_offset=0, + progress_total=None, + show_progress=True, + verbose=True, +): + # Pre-collect per-tensor FP8 scales (0-dim float32 tensors named "{key}_scale"). + # These must be applied to their FP8 weight tensors before GGUF quantization. + # The actual weight value is fp8_value * scale; ignoring scale produces wrong magnitudes. + fp8_scales = fp8_scales if fp8_scales is not None else { + k[:-len("_scale")]: v.item() + for k, v in state_dict.items() + if k.endswith("_scale") and len(v.shape) == 0 and v.dtype == torch.float32 + } + if fp8_scales and verbose: + tqdm.write(f"Found {len(fp8_scales)} FP8 per-tensor scale(s); will apply before quantization.") + + name_lengths = tuple(sorted( + ((key, len(key)) for key in state_dict.keys()), + key=lambda item: item[1], + reverse=True, + )) + if not name_lengths: + return + max_name_len = name_lengths[0][1] + if max_name_len > MAX_TENSOR_NAME_LENGTH: + bad_list = ", ".join(f"{key!r} ({namelen})" for key, namelen in name_lengths if namelen > MAX_TENSOR_NAME_LENGTH) + raise ValueError(f"Can only handle tensor names up to {MAX_TENSOR_NAME_LENGTH} characters. Tensors exceeding the limit: {bad_list}") + _validate_quantization_device(quantization_device) + q8_cr_device = None + q4_cr_device = None + tensor_items = tqdm(state_dict.items()) if show_progress else state_dict.items() + for tensor_index, (key, data) in enumerate(tensor_items, start=1): + old_dtype = data.dtype + + if key_matches(key, model_arch.keys_ignore): + if verbose: + tqdm.write(f"Filtering ignored key: '{key}'") + continue + + # comfy_quant tensors are FP8 scale factors specific to ComfyUI's custom FP8 format. + # weight_scale tensors are 0-dim per-tensor FP8 scales (e.g. from torchao/fp8 fine-tunes). + # Both are meaningless after GGUF re-quantization and must be dropped so the loader + # does not try to apply them to already-GGUF-dequantized weights. + if key.endswith(".comfy_quant") or key.endswith("_scale") and len(data.shape) == 0: + if verbose: + tqdm.write(f"Dropping FP8 scale tensor: '{key}'") + continue + + # 0-dim (scalar) tensors cannot be stored in GGUF and have no meaningful weight data. + if len(data.shape) == 0: + if verbose: + tqdm.write(f"Skipping 0-dim scalar tensor: '{key}'") + continue + + if data.dtype == torch.bfloat16: + data = data.to(torch.float32).numpy() + # this is so we don't break torch 2.0.X + elif data.dtype in [getattr(torch, "float8_e4m3fn", "_invalid"), getattr(torch, "float8_e5m2", "_invalid")]: + data = data.to(torch.float32) + if key in fp8_scales: + data = data * fp8_scales[key] # apply per-tensor dequantization scale + data = data.numpy() + else: + data = data.numpy() + + n_dims = len(data.shape) + data_shape = data.shape + data_qtype = _default_qtype(old_dtype) + + # GGUF supports at most four dimensions. VAE Conv3d weights preserve + # their original shape in metadata and combine their final dimensions + # only for storage. + if n_dims > MAX_TENSOR_DIMS: + if not model_arch.preserve_nd_shapes: + model_arch.handle_nd_tensor(key, data) + continue + orig_shape = data.shape + data = data.reshape(*data.shape[:MAX_TENSOR_DIMS - 1], -1) + data_shape = data.shape + n_dims = len(data_shape) + writer.add_array( + f"comfy.gguf.orig_shape.{key}", + tuple(int(dim) for dim in orig_shape), + ) + + n_params = 1 + for dim_size in data_shape: + n_params *= dim_size + + apply_quantization_rules = ( + quant_type_name == "Q8_CR" + or ( + quant_type_name in QUANT_TYPE_MAP + and quant_type_name not in {"F16", "BF16"} + ) + or old_dtype in (torch.float32, torch.bfloat16) + or old_dtype in _FP8_DTYPES + ) + raw_byte_tensor = _is_raw_byte_tensor(key, old_dtype, n_dims) + if raw_byte_tensor: + data_qtype = gguf.GGMLQuantizationType.I8 + # GGUF exposes only signed I8, but the payload is intentionally + # reinterpreted rather than converted so bytes >= 128 survive. + data = np.ascontiguousarray(data).view(np.int8) + elif quantization_plan is not None and key in quantization_plan: + data_qtype = quantization_plan[key] + elif apply_quantization_rules: + if n_dims == 1: + # One-dimensional tensors should be kept in F32. This is a + # universal safety net and must take priority over + # keys_noquant -- a broad keys_noquant prefix (e.g. Krea2's + # "^last.") would otherwise also match that submodule's 1D + # bias/scale tensors and silently downgrade them from the + # F32 they'd normally always get to whatever the generic + # default happens to be (F16/BF16). + data_qtype = gguf.GGMLQuantizationType.F32 + elif n_params <= QUANTIZATION_THRESHOLD: + data_qtype = gguf.GGMLQuantizationType.F32 + elif key_matches(key, model_arch.keys_hiprec): + # More specific than keys_noquant by design: keys_hiprec + # forces F32 even when a broader keys_noquant pattern for + # the same submodule would also match (e.g. Krea2's + # "last.modulation.lin" needs full F32, even though the + # broader "^last." keys_noquant entry also matches it). + data_qtype = gguf.GGMLQuantizationType.F32 + elif key_matches(key, model_arch.keys_noquant): + pass + elif n_dims == 4 and "conv" in key.lower(): + # Native quantized paths are Linear-only. + data_qtype = gguf.GGMLQuantizationType.F16 + elif quant_type is not None: + data_qtype = quant_type + + if ( + quant_type_name == "Q8_CR" + and n_dims > 1 + and n_dims != 2 + and not key_matches(key, model_arch.keys_hiprec) + and not key_matches(key, model_arch.keys_noquant) + ): + # Custom native layouts only represent Linear matrices. + data_qtype = gguf.GGMLQuantizationType.F16 + # Q4_PT layout restrictions are retained with + # retired_quantize_int4_pytorch and are intentionally not selectable. + + # Q4_CR_W4A4 only supports Linear matrices whose K dimension is divisible + # by the ConvRot group size (default 256). Anything else falls back to F16 + # to avoid a conversion-time crash on awkward shapes. + if ( + quant_type_name == "Q4_CR_W4A4" + and ( + n_dims != 2 + or data.shape[1] % Q4_CR_W4A4_CONVROT_GROUP_SIZE != 0 + ) + and not key_matches(key, model_arch.keys_hiprec) + and not key_matches(key, model_arch.keys_noquant) + and not raw_byte_tensor + ): + data_qtype = gguf.GGMLQuantizationType.F16 + + if raw_byte_tensor: + if verbose: + tqdm.write( + f"{f'%-{max_name_len + 4}s' % key} " + f"{old_dtype} --> {data_qtype.name}, shape = " + f"{{{', '.join(str(n) for n in reversed(data.shape))}}}" + ) + writer.add_tensor(key, data, raw_dtype=data_qtype) + if progress_callback is not None: + progress_callback("quantize", progress_offset + tensor_index, progress_total or len(state_dict)) + continue + + # Q8_CR is the supported custom quantization path. + if ( + data_qtype == gguf.GGMLQuantizationType.I8 + and n_dims == 2 + and quant_type_name == "Q8_CR" + ): + quantization_tensor = torch.from_numpy(data) + if q8_cr_device is None: + q8_cr_device = resolve_quantization_device(quantization_device) + device = q8_cr_device + if device.type == "cuda" and not _can_use_cuda_q8_cr(quantization_tensor, device): + logging.warning( + "Q8_CR CUDA fallback for %s: insufficient free VRAM for this matrix.", + key, + ) + device = torch.device("cpu") + try: + qdata, scale, quant_conf, orig_shape = quantize_int8_convrot( + quantization_tensor, + device=device, + ) + except torch.OutOfMemoryError: + if device.type != "cuda": + raise + torch.cuda.empty_cache() + logging.warning( + "Q8_CR CUDA fallback for %s: CUDA ran out of memory while quantizing.", + key, + ) + qdata, scale, quant_conf, orig_shape = quantize_int8_convrot( + quantization_tensor, + device=torch.device("cpu"), + ) + writer.add_tensor( + key, + qdata.cpu().numpy(), + raw_dtype=gguf.GGMLQuantizationType.I8, + ) + writer.add_tensor( + f"{key}_scale", + scale.cpu().numpy(), + raw_dtype=gguf.GGMLQuantizationType.F32, + ) + writer.add_string(f"comfy.gguf.quant.{key}", json.dumps(quant_conf)) + if progress_callback is not None: + progress_callback("quantize", progress_offset + tensor_index, progress_total or len(state_dict)) + continue + + # Q4_CR_W4A4: custom W4A4 INT4 backed by comfy_kitchen's fast ConvRot + # int4 tensor-core MMA. Serializes kitchen-native packed int4 (N, K//2) + # + per-output-row fp scales (no per-group scales/zeros). + if ( + quant_type_name == "Q4_CR_W4A4" + and data_qtype == gguf.GGMLQuantizationType.I8 + and n_dims == 2 + and data.shape[1] % Q4_CR_W4A4_CONVROT_GROUP_SIZE == 0 + and not key_matches(key, model_arch.keys_hiprec) + and not key_matches(key, model_arch.keys_noquant) + ): + qdata_q4 = torch.from_numpy(data) + if q4_cr_device is None: + q4_cr_device = resolve_quantization_device(quantization_device) + device = q4_cr_device + try: + qdata, wscales, quant_conf, orig_shape = quantize_int4_cr_w4a4( + qdata_q4, + convrot_groupsize=Q4_CR_W4A4_CONVROT_GROUP_SIZE, + quant_group_size=Q4_CR_W4A4_QUANT_GROUP_SIZE, + device=device, + ) + except torch.OutOfMemoryError: + if device.type != "cuda": + raise + torch.cuda.empty_cache() + logging.warning( + "Q4_CR_W4A4 CUDA fallback for %s: CUDA ran out of memory while quantizing.", + key, + ) + qdata, wscales, quant_conf, orig_shape = quantize_int4_cr_w4a4( + qdata_q4, + convrot_groupsize=Q4_CR_W4A4_CONVROT_GROUP_SIZE, + quant_group_size=Q4_CR_W4A4_QUANT_GROUP_SIZE, + device=torch.device("cpu"), + ) + writer.add_tensor( + key, + qdata.cpu().numpy(), + raw_dtype=gguf.GGMLQuantizationType.I8, + ) + writer.add_tensor( + f"{key}_scale", + wscales.cpu().half().numpy(), + raw_dtype=gguf.GGMLQuantizationType.F16, + ) + writer.add_string(f"comfy.gguf.quant.{key}", json.dumps(quant_conf)) + if progress_callback is not None: + progress_callback("quantize", progress_offset + tensor_index, progress_total or len(state_dict)) + continue + + if (model_arch.shape_fix # NEVER reshape for models such as flux + and n_dims > 1 # Skip one-dimensional tensors + and n_params >= REARRANGE_THRESHOLD # Only rearrange tensors meeting the size requirement + and (n_params / 256).is_integer() # Rearranging only makes sense if total elements is divisible by 256 + and not (data.shape[-1] / 256).is_integer() # Only need to rearrange if the last dimension is not divisible by 256 + ): + orig_shape = data.shape + data = data.reshape(n_params // 256, 256) + writer.add_array(f"comfy.gguf.orig_shape.{key}", tuple(int(dim) for dim in orig_shape)) + + try: + data = gguf.quants.quantize(data, data_qtype) + except (AttributeError, gguf.QuantError) as e: + if verbose: + tqdm.write(f"falling back to F16: {e}") + data_qtype = gguf.GGMLQuantizationType.F16 + data = gguf.quants.quantize(data, data_qtype) + + new_name = key # do we need to rename? + + shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}" + if verbose: + tqdm.write( + f"{f'%-{max_name_len + 4}s' % f'{new_name}'} " + f"{old_dtype} --> {data_qtype.name}, shape = {shape_str}" + ) + + writer.add_tensor(new_name, data, raw_dtype=data_qtype) + if progress_callback is not None: + progress_callback("quantize", progress_offset + tensor_index, progress_total or len(state_dict)) + +def convert_file( + path, + dst_path=None, + interact=True, + overwrite=False, + quant_type_name=None, + max_size_mb=None, + target_size_q8_type=DEFAULT_TARGET_SIZE_Q8_TYPE, + quantization_device="auto", + progress_callback=None, + lora_paths=None, + lora_strengths=None, + streamed=False, +): + if streamed: + return convert_safetensors_streamed( + path, + dst_path=dst_path, + interact=interact, + overwrite=overwrite, + quant_type_name=quant_type_name, + max_size_mb=max_size_mb, + target_size_q8_type=target_size_q8_type, + quantization_device=quantization_device, + progress_callback=progress_callback, + lora_paths=lora_paths, + lora_strengths=lora_strengths, + ) + state_dict = load_state_dict(path, progress_callback=progress_callback) + restored_int8_count = materialize_int8_source_weights(state_dict) + if restored_int8_count: + logging.info( + "Restored %d scaled INT8 source weight(s) to FP16 before conversion.", + restored_int8_count, + ) + source_metadata = load_safetensors_metadata(path) + lora_paths = lora_paths or [] + lora_strengths = lora_strengths or [] + if lora_strengths and len(lora_strengths) != len(lora_paths): + raise ValueError("Provide one --lora-strength for each --lora.") + if not lora_strengths: + lora_strengths = [1.0] * len(lora_paths) + if lora_paths: + device = resolve_quantization_device(quantization_device) + for lora_path, strength in zip(lora_paths, lora_strengths): + if not os.path.isfile(lora_path): + raise FileNotFoundError(f"LoRA does not exist: {lora_path}") + _, targets, _ = load_lora(lora_path) + fused_count = fuse_targets_into_state_dict(state_dict, targets, strength, device) + logging.info("Merged %d LoRA targets from %s.", fused_count, lora_path) + if device.type == "cuda": + torch.cuda.empty_cache() + return convert_state_dict( + state_dict, + dst_path=dst_path, + source_path=path, + source_metadata=source_metadata, + interact=interact, + overwrite=overwrite, + quant_type_name=quant_type_name, + max_size_mb=max_size_mb, + target_size_q8_type=target_size_q8_type, + quantization_device=quantization_device, + progress_callback=progress_callback, + ) + + +def _streamed_safetensors_layout(path): + if not path.endswith(".safetensors"): + raise ValueError("--streamed supports only .safetensors source checkpoints.") + dtype_map = { + "BOOL": torch.bool, + "U8": torch.uint8, + "I8": torch.int8, + "I16": torch.int16, + "I32": torch.int32, + "I64": torch.int64, + "F16": torch.float16, + "BF16": torch.bfloat16, + "F32": torch.float32, + "F64": torch.float64, + "F8_E4M3": getattr(torch, "float8_e4m3fn", None), + "F8_E4M3FN": getattr(torch, "float8_e4m3fn", None), + "F8_E5M2": getattr(torch, "float8_e5m2", None), + } + with safe_open(path, framework="pt", device="cpu") as checkpoint: + source_layout = OrderedDict() + for key in checkpoint.keys(): + tensor_slice = checkpoint.get_slice(key) + dtype_name = tensor_slice.get_dtype() + dtype = dtype_map.get(dtype_name) + if dtype is None: + raise ValueError( + f"Streamed conversion does not support safetensors dtype {dtype_name!r} " + f"for tensor {key!r}." + ) + source_layout[key] = torch.empty( + tuple(tensor_slice.get_shape()), dtype=dtype, device="meta" + ) + layout = strip_prefix(source_layout) + source_keys = {id(value): key for key, value in source_layout.items()} + return layout, {key: source_keys[id(value)] for key, value in layout.items()} + + +def convert_safetensors_streamed( + path, + dst_path=None, + interact=True, + overwrite=False, + quant_type_name=None, + max_size_mb=None, + target_size_q8_type=DEFAULT_TARGET_SIZE_Q8_TYPE, + quantization_device="auto", + progress_callback=None, + lora_paths=None, + lora_strengths=None, +): + """Convert one safetensors tensor at a time, flushing GGUF payload staging to disk.""" + state_dict, source_keys = _streamed_safetensors_layout(path) + source_metadata = load_safetensors_metadata(path) + lora_paths = lora_paths or [] + lora_strengths = lora_strengths or [] + if lora_strengths and len(lora_strengths) != len(lora_paths): + raise ValueError("Provide one --lora-strength for each --lora.") + if not lora_strengths: + lora_strengths = [1.0] * len(lora_paths) + + model_arch = detect_arch(state_dict) + logging.info(f"* Architecture detected from input: {model_arch.arch}") + validate_key_patterns(model_arch, state_dict) + if max_size_mb is not None and quant_type_name not in (None, TARGET_SIZE_QUANT_TYPE): + raise ValueError("--max-size-mb cannot be combined with --quant-type.") + + quantization_plan = None + if max_size_mb is not None: + quant_type_name = TARGET_SIZE_QUANT_TYPE + quantization_plan, maximum_size, selected_size = plan_target_size_quantization( + state_dict, model_arch, max_size_mb, target_size_q8_type=target_size_q8_type + ) + logging.info( + "TARGET_SIZE selected %.2f MiB from a %s baseline of %.2f MiB.", + selected_size / MEBIBYTE, + target_size_q8_type, + maximum_size / MEBIBYTE, + ) + + quant_type = None + if quantization_plan is not None: + ftype_name, ftype_gguf = TARGET_SIZE_QUANT_TYPE, None + elif quant_type_name is not None and quant_type_name in QUANT_TYPE_MAP: + quant_type, ftype_gguf = QUANT_TYPE_MAP[quant_type_name] + ftype_name = quant_type_name + else: + dtypes = [value.dtype for value in state_dict.values()] + main_dtype = max(set(dtypes), key=dtypes.count) + ftype_name = "BF16" if main_dtype == torch.bfloat16 else "F16" + ftype_gguf = ( + gguf.LlamaFileType.MOSTLY_BF16 + if main_dtype == torch.bfloat16 + else gguf.LlamaFileType.MOSTLY_F16 + ) + + if dst_path is None: + dst_path = f"{os.path.splitext(path)[0]}-{ftype_name}.gguf" + elif "{ftype}" in dst_path: + dst_path = dst_path.replace("{ftype}", ftype_name) + if os.path.isfile(dst_path) and not overwrite: + if interact: + input("Output exists enter to continue or ctrl+c to abort!") + else: + raise OSError("Output exists and overwriting is disabled!") + + resolved_loras = {} + if lora_paths: + device = resolve_quantization_device(quantization_device) + for lora_path, strength in zip(lora_paths, lora_strengths): + if not os.path.isfile(lora_path): + raise FileNotFoundError(f"LoRA does not exist: {lora_path}") + _, targets, _ = load_lora(lora_path) + resolved = resolve_fusion_targets(state_dict, targets) + fused_count = sum(len(entries) for entries in resolved.values()) + logging.info("Merged %d LoRA targets from %s.", fused_count, lora_path) + for key, entries in resolved.items(): + resolved_loras.setdefault(key, []).append((entries, strength)) + + writer = gguf.GGUFWriter(path=None, arch=model_arch.arch, use_temp_file=True) + writer.temp_file = tempfile.TemporaryFile(mode="w+b") + writer.add_quantization_version(gguf.GGML_QUANT_VERSION) + if ftype_gguf is not None: + writer.add_file_type(ftype_gguf) + if "config" in source_metadata: + writer.add_string("config", source_metadata["config"]) + + streamed_progress = None + if progress_callback is None: + streamed_progress = tqdm( + total=len(state_dict), desc="Quantizing", unit="tensor" + ) + try: + with safe_open(path, framework="pt", device="cpu") as checkpoint: + total = len(state_dict) + for index, (key, _) in enumerate(state_dict.items(), start=1): + data = checkpoint.get_tensor(source_keys[key]) + if key in resolved_loras: + source_scale = None + scale_key = f"{key}_scale" + if data.dtype in _FP8_DTYPES and scale_key in source_keys: + source_scale = checkpoint.get_tensor(source_keys[scale_key]) + for target_entries, strength in resolved_loras[key]: + data, _ = fuse_target_entries_into_tensor( + data, target_entries, strength, device, source_scale + ) + source_scale = None + fp8_scales = {} + if data.dtype in _FP8_DTYPES: + scale_key = f"{key}_scale" + if scale_key in source_keys: + fp8_scales[key] = checkpoint.get_tensor(source_keys[scale_key]).item() + handle_tensors( + writer, + OrderedDict(((key, data),)), + model_arch, + quant_type=quant_type, + quant_type_name=quant_type_name, + quantization_plan=quantization_plan, + quantization_device=quantization_device, + progress_callback=progress_callback, + fp8_scales=fp8_scales, + progress_offset=index - 1, + progress_total=total, + show_progress=False, + verbose=False, + ) + del data + # A GGUF header requires the complete tensor table, so the final + # file can only be assembled after conversion. Keep the payload + # staging file durable and growing throughout the conversion. + if writer.temp_file is not None: + writer.temp_file.flush() + if streamed_progress is not None: + streamed_progress.update() + writer.write_header_to_file(path=dst_path) + writer.write_kv_data_to_file() + writer.write_tensors_to_file(progress=True) + finally: + if streamed_progress is not None: + streamed_progress.close() + writer.close() + if writer.temp_file is not None and not writer.temp_file.closed: + writer.temp_file.close() + return dst_path, model_arch + + +def convert_state_dict( + state_dict, + dst_path, + source_path="", + source_metadata=None, + interact=False, + overwrite=False, + quant_type_name=None, + max_size_mb=None, + target_size_q8_type=DEFAULT_TARGET_SIZE_Q8_TYPE, + quantization_device="auto", + progress_callback=None, +): + """Convert an already loaded, prefix-normalized diffusion-model state dict.""" + source_metadata = source_metadata or {} + model_arch = detect_arch(state_dict) + logging.info(f"* Architecture detected from input: {model_arch.arch}") + validate_key_patterns(model_arch, state_dict) + + if max_size_mb is not None and quant_type_name not in (None, TARGET_SIZE_QUANT_TYPE): + raise ValueError("--max-size-mb cannot be combined with --quant-type.") + + quantization_plan = None + if max_size_mb is not None: + quant_type_name = TARGET_SIZE_QUANT_TYPE + quantization_plan, maximum_size, selected_size = plan_target_size_quantization( + state_dict, + model_arch, + max_size_mb, + target_size_q8_type=target_size_q8_type, + ) + logging.info( + "TARGET_SIZE selected %.2f MiB from a %s baseline of %.2f MiB.", + selected_size / MEBIBYTE, + target_size_q8_type, + maximum_size / MEBIBYTE, + ) + + # resolve quant type from name if provided + quant_type = None + if quantization_plan is not None: + ftype_name = TARGET_SIZE_QUANT_TYPE + ftype_gguf = None + elif quant_type_name is not None and quant_type_name in QUANT_TYPE_MAP: + quant_type, ftype_gguf = QUANT_TYPE_MAP[quant_type_name] + ftype_name = quant_type_name + else: + # detect & set dtype from source file + dtypes = [x.dtype for x in state_dict.values()] + dtypes = {x: dtypes.count(x) for x in set(dtypes)} + main_dtype = max(dtypes, key=dtypes.get) + + if main_dtype == torch.bfloat16: + ftype_name = "BF16" + ftype_gguf = gguf.LlamaFileType.MOSTLY_BF16 + # elif main_dtype == torch.float32: + # ftype_name = "F32" + # ftype_gguf = None + else: + ftype_name = "F16" + ftype_gguf = gguf.LlamaFileType.MOSTLY_F16 + + if dst_path is None: + dst_path = f"{os.path.splitext(source_path)[0]}-{ftype_name}.gguf" + elif "{ftype}" in dst_path: # lcpp logic + dst_path = dst_path.replace("{ftype}", ftype_name) + + if os.path.isfile(dst_path) and not overwrite: + if interact: + input("Output exists enter to continue or ctrl+c to abort!") + else: + raise OSError("Output exists and overwriting is disabled!") + + # handle actual file + writer = gguf.GGUFWriter(path=None, arch=model_arch.arch) + writer.add_quantization_version(gguf.GGML_QUANT_VERSION) + if ftype_gguf is not None: + writer.add_file_type(ftype_gguf) + if "config" in source_metadata: + writer.add_string("config", source_metadata["config"]) + + handle_tensors( + writer, + state_dict, + model_arch, + quant_type=quant_type, + quant_type_name=quant_type_name, + quantization_plan=quantization_plan, + quantization_device=quantization_device, + progress_callback=progress_callback, + ) + writer.write_header_to_file(path=dst_path) + writer.write_kv_data_to_file() + writer.write_tensors_to_file(progress=True) + writer.close() + + fix = f"./fix_5d_tensors_{model_arch.arch}.safetensors" + if os.path.isfile(fix): + logging.warning(f"\n### Warning! Fix file found at '{fix}'") + logging.warning(" you most likely need to run 'fix_5d_tensors.py' after quantization.") + + return dst_path, model_arch + +if __name__ == "__main__": + args = parse_args() + convert_file( + args.src, + args.dst, + quant_type_name=args.quant_type, + max_size_mb=args.max_size_mb, + target_size_q8_type=args.target_size_q8_type, + quantization_device=args.quantization_device, + lora_paths=args.lora, + lora_strengths=args.lora_strength, + streamed=args.streamed, + ) diff --git a/tools/convert_krea2_gguf.py b/tools/convert_krea2_gguf.py new file mode 100644 index 00000000..5a4afce2 --- /dev/null +++ b/tools/convert_krea2_gguf.py @@ -0,0 +1,144 @@ +""" +Batch GGUF converter for Krea-2 (Base and Turbo) models. + +Produces Q4_0, Q4_1, Q5_0, Q5_1, and Q8_0 GGUF files from a single +BF16/F16 safetensors source file. The converter auto-detects architecture +(ModelKrea2 or ModelIdeogram) from the tensor keys in the source file. + +Usage +----- +# Convert all five quant levels from a BF16 safetensors: +python tools/convert_krea2_gguf.py --src krea2_base_bf16.safetensors + +# Convert specific quant levels only: +python tools/convert_krea2_gguf.py --src krea2_turbo_bf16.safetensors --quant Q5_0 Q8_0 + +# Explicit output directory: +python tools/convert_krea2_gguf.py --src krea2_base_bf16.safetensors --outdir /path/to/output + +Notes +----- +- 1-D tensors (biases, norms, scales) are always kept in F32 regardless of + the requested quant type. +- Tensors with fewer than 1024 elements are kept in F32. +- No model weights are downloaded by this script. You must supply the source + safetensors file yourself. +- This script calls convert_file() from tools/convert.py; both files must be + on the Python path (they are when run from the repo root or from tools/). + +Approximate output sizes per quant level (Krea-2 ~24 B parameters): + Q4_0 ~ 14 GB Q4_1 ~ 15 GB Q5_0 ~ 17 GB Q5_1 ~ 18 GB Q8_0 ~ 26 GB +""" + +import os +import sys +import logging +import argparse +from tqdm import tqdm + +# Allow running directly from tools/ or from the repo root. +_HERE = os.path.dirname(os.path.abspath(__file__)) +_REPO = os.path.dirname(_HERE) +for _p in (_HERE, _REPO): + if _p not in sys.path: + sys.path.insert(0, _p) + +from convert import convert_file, QUANT_TYPE_MAP # noqa: E402 (tools/convert.py) + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + +DEFAULT_QUANTS = ["Q4_0", "Q8_CR", "Q8_0"] + + +def build_dst_path(src: str, quant: str, outdir: str | None) -> str: + """Derive an output path like -Q4_0.gguf next to the source.""" + stem = os.path.splitext(os.path.basename(src))[0] + filename = f"{stem}-{quant}.gguf" + directory = outdir if outdir else os.path.dirname(os.path.abspath(src)) + return os.path.join(directory, filename) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Batch-convert a Krea-2 safetensors file to multiple GGUF quant levels.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + parser.add_argument( + "--src", + required=True, + help="Path to the source BF16/F16 safetensors (or ckpt) file.", + ) + parser.add_argument( + "--outdir", + default=None, + help="Directory to write GGUF files into. Defaults to the same directory as --src.", + ) + parser.add_argument( + "--quant", + nargs="+", + choices=list(QUANT_TYPE_MAP.keys()), + default=DEFAULT_QUANTS, + metavar="QUANT", + help=( + "One or more quantization levels to produce. " + f"Choices: {list(QUANT_TYPE_MAP.keys())}. " + f"Default: {DEFAULT_QUANTS}" + ), + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="Overwrite existing output files without prompting.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + if not os.path.isfile(args.src): + logging.error(f"Source file not found: {args.src}") + sys.exit(1) + + if args.outdir: + os.makedirs(args.outdir, exist_ok=True) + + results: list[tuple[str, str, bool]] = [] # (quant, dst_path, success) + + for quant in tqdm(args.quant, desc="GGUF quant levels", unit="quant"): + dst = build_dst_path(args.src, quant, args.outdir) + logging.info(f"\n{'='*60}") + logging.info(f"Converting {os.path.basename(args.src)} → {os.path.basename(dst)}") + logging.info(f"Quant type : {quant}") + logging.info(f"Output : {dst}") + logging.info(f"{'='*60}") + + try: + out_path, model_arch = convert_file( + args.src, + dst_path=dst, + interact=False, + overwrite=args.overwrite, + quant_type_name=quant, + ) + logging.info(f"[OK] {quant} → {out_path} (arch={model_arch.arch})") + results.append((quant, out_path, True)) + except Exception as exc: + logging.error(f"[FAILED] {quant}: {exc}") + results.append((quant, dst, False)) + + # Summary + print(f"\n{'='*60}") + print("Conversion summary:") + for quant, path, ok in results: + status = "OK " if ok else "FAIL" + print(f" [{status}] {quant:6s} {path}") + print(f"{'='*60}") + + if not all(ok for _, _, ok in results): + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/tools/krea2_gguf_readme.md b/tools/krea2_gguf_readme.md new file mode 100644 index 00000000..5321d578 --- /dev/null +++ b/tools/krea2_gguf_readme.md @@ -0,0 +1,100 @@ +--- +license: other +license_name: krea-2-community-license +license_link: https://huggingface.co/krea/Krea-2-Turbo/blob/main/LICENSE +pipeline_tag: text-to-image +tags: + - image-generation + - diffusion + - flow-matching + - dit + - krea2 + - quantization + - comfyui +library_name: gguf +base_model: + - krea/Krea-2-Raw + - krea/Krea-2-Turbo +--- + +# Krea 2 GGUF + +Quantized GGUF diffusion transformer weights for [Krea 2](https://huggingface.co/krea/Krea-2-Turbo), +converted from the original BF16 releases for use with ComfyUI GGUF loader nodes. + +This repository provides GGUF files for two checkpoints of the Krea 2 model family: + +- `krea2_raw_bf16-*.gguf` — converted from [krea/Krea-2-Raw](https://huggingface.co/krea/Krea-2-Raw), the base release checkpoint. +- `krea2_turbo_bf16-*.gguf` — converted from [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo), the post-trained checkpoint with additional fine-tuning and distillation. + +Krea 2 is a 12-billion parameter Diffusion Transformer with a novel architecture featuring layerwise and refiner text-fusion blocks. It is not based on Flux or any prior open-weight architecture. + +These files are not a complete standalone Krea 2 package. Your workflow still needs the text encoder and VAE components. + +## ComfyUI Support + +Use these models with the ComfyUI nodes from [molbal/ComfyUI-GGUF](https://github.com/molbal/ComfyUI-GGUF). +Install that custom node repository into your ComfyUI `custom_nodes` folder, then restart ComfyUI. + +> **Important:** This repository requires [molbal/ComfyUI-GGUF](https://github.com/molbal/ComfyUI-GGUF), +> which is a fork of [city96/ComfyUI-GGUF](https://github.com/city96/ComfyUI-GGUF) with added support +> for the Krea 2 architecture. The original city96 plugin does **not** support these files. (as of 2026-06-24) + +Place the downloaded `.gguf` files in one of ComfyUI's diffusion model folders: + +``` +ComfyUI/models/diffusion_models/ +ComfyUI/models/unet/ +``` + +Load the file with `Unet Loader (GGUF)` in a Krea 2 workflow. Krea 2 uses a single transformer +(unlike Ideogram 4, there is no separate unconditional transformer component). + +## Files + +| Quant | Raw (base) | Turbo | Size | +|-------|-----------|-------------------|------| +| Q4_0 | krea2_raw_bf16-Q4_0.gguf | krea2_turbo_bf16-Q4_0.gguf | 7.74 GB | +| Q4_1 | krea2_raw_bf16-Q4_1.gguf | krea2_turbo_bf16-Q4_1.gguf | 8.47 GB | +| Q5_0 | krea2_raw_bf16-Q5_0.gguf | krea2_turbo_bf16-Q5_0.gguf | 9.20 GB | +| Q5_1 | krea2_raw_bf16-Q5_1.gguf | krea2_turbo_bf16-Q5_1.gguf | 9.93 GB | +| Q8_0 | krea2_raw_bf16-Q8_0.gguf | krea2_turbo_bf16-Q8_0.gguf | 13.56 GB | + +Choose either the Raw or Turbo variant depending on your workflow; they are not paired with each other. + +## Which Checkpoint to Use + +| Checkpoint | Steps | CFG | Notes | +|------------|-------|-----|-------| +| **Turbo** | 4–8 | 0.0 | Distilled; CFG-free. Fast, good for most use cases. | +| **Raw** | 20–30 | 3.0–7.0 | Full CFG; more controllable, higher inference cost. | + +The Turbo checkpoint has been post-trained with distillation and runs well at 8 steps with `CFG=1`. The Raw checkpoint behaves like a standard flow-matching DiT and benefits from more steps and positive CFG. + +## When GGUFs Are Worth the Tradeoff + +The BF16 source weights for Krea 2 are 26.6 GB each — far beyond what most consumer GPUs can hold entirely in VRAM. GGUFs make sense when: + +- **Limited VRAM:** Q4_0 at 7.74 GB fits entirely in an 8 GB GPU; Q5_1 at 9.93 GB targets 10–12 GB cards. Running BF16 on these GPUs would require heavy CPU offloading and become impractically slow. +- **CPU offload workflows:** If you are already offloading model layers to RAM, GGUF reduces the RAM footprint proportionally alongside VRAM, which is often the actual bottleneck. +- **Acceptable quality loss at Q5+:** At Q5_0 and above the visual output of Krea 2 is very close to BF16. Q4 levels show mild softening on fine detail but remain usable for most creative tasks. + +GGUFs are generally **not** worth it if you have a 24 GB+ GPU and want maximum fidelity — load the FP8 or BF16 source directly in that case. + +## Download + +Download the file you want from the Files tab, or use the Hugging Face CLI. For example: + +``` +huggingface-cli download molbal/krea2-gguf krea2_turbo_bf16-Q5_1.gguf --local-dir ComfyUI/models/diffusion_models +``` + +## Compatibility Notes + +These are non-K GGUF quantizations intended for PyTorch dequantization in ComfyUI. K-quants are not included because this ComfyUI loading path does not use fused quantized linear kernels. + +Krea 2 GGUF support requires ComfyUI to have the `krea2` architecture registered in its model detection system. If your ComfyUI installation does not recognise the checkpoint, update ComfyUI core and [molbal/ComfyUI-GGUF](https://github.com/molbal/ComfyUI-GGUF) to their latest versions. + +## License + +These files are derived from [krea/Krea-2-Raw](https://huggingface.co/krea/Krea-2-Raw) and [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo) and follow the [Krea 2 Community License](https://huggingface.co/krea/Krea-2-Turbo/blob/main/LICENSE). diff --git a/tools/lcpp.patch b/tools/lcpp.patch index 92396e17..9018f5ac 100644 --- a/tools/lcpp.patch +++ b/tools/lcpp.patch @@ -1,5 +1,5 @@ diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h -index de3c706f..0267c1fa 100644 +index de3c706fc..0267c1faa 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -223,7 +223,7 @@ @@ -20,7 +20,7 @@ index de3c706f..0267c1fa 100644 GGML_API void gguf_set_tensor_data(struct gguf_context * ctx, const char * name, const void * data, size_t size); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c -index b16c462f..6d1568f1 100644 +index b16c462fa..6d1568f1b 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -22960,6 +22960,14 @@ void gguf_add_tensor( @@ -39,10 +39,10 @@ index b16c462f..6d1568f1 100644 const int idx = gguf_find_tensor(ctx, name); if (idx < 0) { diff --git a/src/llama.cpp b/src/llama.cpp -index 24e1f1f0..25db4c69 100644 +index 24e1f1f01..23df89523 100644 --- a/src/llama.cpp +++ b/src/llama.cpp -@@ -205,6 +205,17 @@ enum llm_arch { +@@ -205,6 +205,18 @@ enum llm_arch { LLM_ARCH_GRANITE, LLM_ARCH_GRANITE_MOE, LLM_ARCH_CHAMELEON, @@ -57,10 +57,11 @@ index 24e1f1f0..25db4c69 100644 + LLM_ARCH_HIDREAM, + LLM_ARCH_COSMOS, + LLM_ARCH_LUMINA2, ++ LLM_ARCH_KREA2, LLM_ARCH_UNKNOWN, }; -@@ -258,6 +269,17 @@ static const std::map LLM_ARCH_NAMES = { +@@ -258,6 +270,18 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GRANITE, "granite" }, { LLM_ARCH_GRANITE_MOE, "granitemoe" }, { LLM_ARCH_CHAMELEON, "chameleon" }, @@ -75,10 +76,11 @@ index 24e1f1f0..25db4c69 100644 + { LLM_ARCH_HIDREAM, "hidream" }, + { LLM_ARCH_COSMOS, "cosmos" }, + { LLM_ARCH_LUMINA2, "lumina2" }, ++ { LLM_ARCH_KREA2, "krea2" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; -@@ -1531,6 +1553,17 @@ static const std::map> LLM_TENSOR_N +@@ -1531,6 +1555,18 @@ static const std::map> LLM_TENSOR_N { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, }, }, @@ -93,10 +95,11 @@ index 24e1f1f0..25db4c69 100644 + { LLM_ARCH_HIDREAM, {}}, + { LLM_ARCH_COSMOS, {}}, + { LLM_ARCH_LUMINA2, {}}, ++ { LLM_ARCH_KREA2, {}}, { LLM_ARCH_UNKNOWN, { -@@ -5403,6 +5436,25 @@ static void llm_load_hparams( +@@ -5403,6 +5439,26 @@ static void llm_load_hparams( // get general kv ml.get_key(LLM_KV_GENERAL_NAME, model.name, false); @@ -113,6 +116,7 @@ index 24e1f1f0..25db4c69 100644 + case LLM_ARCH_HIDREAM: + case LLM_ARCH_COSMOS: + case LLM_ARCH_LUMINA2: ++ case LLM_ARCH_KREA2: + model.ftype = ml.ftype; + return; + default: @@ -122,7 +126,7 @@ index 24e1f1f0..25db4c69 100644 // get hparams kv ml.get_key(LLM_KV_VOCAB_SIZE, hparams.n_vocab, false) || ml.get_arr_n(LLM_KV_TOKENIZER_LIST, hparams.n_vocab); -@@ -18016,6 +18068,134 @@ static void llama_tensor_dequantize_internal( +@@ -18016,6 +18072,135 @@ static void llama_tensor_dequantize_internal( workers.clear(); } @@ -158,6 +162,7 @@ index 24e1f1f0..25db4c69 100644 + (name.find(".v.weight") != std::string::npos) || + (name.find(".attn.w1v.weight") != std::string::npos) || + (name.find(".attn.w2v.weight") != std::string::npos) || ++ (name.find(".attn.wv.weight") != std::string::npos) || + (name.find("_attn.v_proj.weight") != std::string::npos) + ){ + if (ftype == LLAMA_FTYPE_MOSTLY_Q2_K) { @@ -257,7 +262,7 @@ index 24e1f1f0..25db4c69 100644 static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type new_type, const ggml_tensor * tensor, llama_ftype ftype) { const std::string name = ggml_get_name(tensor); -@@ -18513,7 +18693,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18513,7 +18698,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s if (llama_model_has_encoder(&model)) { n_attn_layer *= 3; } @@ -268,7 +273,7 @@ index 24e1f1f0..25db4c69 100644 } size_t total_size_org = 0; -@@ -18547,6 +18729,51 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18547,6 +18734,60 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s ctx_outs[i_split] = gguf_init_empty(); } gguf_add_tensor(ctx_outs[i_split], tensor); @@ -316,11 +321,20 @@ index 24e1f1f0..25db4c69 100644 + gguf_set_tensor_ndim(ctx_outs[i_split], tensor->name, n_dim); + LLAMA_LOG_INFO("\n%s: Correcting shape for Wan FLF2V: [key:%s]\n", __func__, tensor->name); + } ++ } ++ // Krea2's txtfusion.projector.weight has shape (1, 12) — leading dim of 1 gets truncated ++ if (model.arch == LLM_ARCH_KREA2) { ++ const std::string name = ggml_get_name(tensor); ++ if (name == "txtfusion.projector.weight" && tensor->ne[1] == 1) { ++ const int n_dim = 2; ++ gguf_set_tensor_ndim(ctx_outs[i_split], "txtfusion.projector.weight", n_dim); ++ LLAMA_LOG_INFO("\n%s: Correcting txtfusion.projector.weight shape for Krea2: [key:%s]\n", __func__, tensor->name); ++ } + } } // Set split info if needed -@@ -18647,6 +18874,110 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18647,6 +18888,119 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s // do not quantize relative position bias (T5) quantize &= name.find("attn_rel_b.weight") == std::string::npos; @@ -423,6 +437,15 @@ index 24e1f1f0..25db4c69 100644 + quantize &= name.find("context_refiner.") == std::string::npos; + quantize &= name.find("noise_refiner.") == std::string::npos; + } ++ if (model.arch == LLM_ARCH_KREA2) { ++ image_model = true; ++ quantize &= name.find("first.") == std::string::npos; ++ quantize &= name.find("last.") == std::string::npos; ++ quantize &= name.find("tproj.") == std::string::npos; ++ quantize &= name.find("tmlp.") == std::string::npos; ++ quantize &= name.find("txtmlp.") == std::string::npos; ++ quantize &= name.find("txtfusion.projector.") == std::string::npos; ++ } + // ignore 3D/4D tensors for image models as the code was never meant to handle these + if (image_model) { + quantize &= ggml_n_dims(tensor) == 2; @@ -431,7 +454,7 @@ index 24e1f1f0..25db4c69 100644 enum ggml_type new_type; void * new_data; size_t new_size; -@@ -18655,6 +18986,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18655,6 +19009,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s new_type = default_type; // get more optimal quantization type based on the tensor shape, layer, etc. @@ -441,11 +464,11 @@ index 24e1f1f0..25db4c69 100644 if (!params->pure && ggml_is_quantized(default_type)) { new_type = llama_tensor_get_type(qs, new_type, tensor, ftype); } -@@ -18664,6 +18998,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s +@@ -18664,6 +19021,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s if (params->output_tensor_type < GGML_TYPE_COUNT && strcmp(tensor->name, "output.weight") == 0) { new_type = params->output_tensor_type; } + } // If we've decided to quantize to the same type the tensor is already - // in then there's nothing to do. + // in then there's nothing to do. \ No newline at end of file diff --git a/tools/read_tensors.py b/tools/read_tensors.py index 1bdff028..08554e71 100644 --- a/tools/read_tensors.py +++ b/tools/read_tensors.py @@ -1,21 +1,185 @@ #!/usr/bin/python3 +""" +read_tensors.py - list tensors in a GGUF file, optionally comparing two files. + +Single-file mode: + python read_tensors.py model.gguf [--quantized-only] + + Lists every tensor with its quant type and shape. By default this + INCLUDES F32 tensors (unlike the original version), since those are + exactly the ones that can silently suffer a dimension collapse when + quantizing (see below). Pass --quantized-only to restore the old + behavior of skipping F32 tensors. Tensors with a suspicious shape + (see is_suspicious_shape) are flagged inline. + +Compare mode (recommended when you have both files): + python read_tensors.py before.gguf after.gguf + + Compares two GGUF files tensor-by-tensor (e.g. the pre-quantize + BF16/F16 GGUF from convert.py vs. the output of llama-quantize). + Flags any tensor whose number of dimensions changed between the two + files -- this is the exact, unambiguous signature of a GGML + ggml_n_dims() collapse (a tensor with a trailing `ne` dim of 1 gets + silently truncated to fewer dims when llama.cpp reloads it into a + ggml_tensor and writes it back out) -- e.g. SD3's pos_embed, + AuraFlow's positional_encoding/register_tokens, Wan's .modulation + tensors, or Krea2's txtfusion.projector.weight. This mode doesn't + rely on any heuristic, so prefer it whenever you have both files. + + Any tensor flagged here is a candidate for a gguf_set_tensor_ndim() + fix in lcpp.patch, inside the llama_model_quantize_internal() + tensor-writing loop. + +Batch/log-friendly compare mode: + python read_tensors.py --verify before.gguf after.gguf + + Same check as compare mode, but prints exactly one line -- "OK" or + "FAIL" plus a compact issue list -- instead of the full per-tensor + table. No --quantized-only equivalent needed here since the ndim + check inherently only concerns itself with what actually differs. + Skips the interactive "press enter to close" pause, and sets the + process exit code (0 = OK, 1 = FAIL) so it plays well in scripts. + Intended for looping over several quant types and collecting one + line per run into a shared log file, e.g.: + + for QUANT in Q4_0 Q5_0 Q8_0 Q2_K Q5_K_M Q6_K; do + ./llama-quantize model-BF16.gguf model-${QUANT}.gguf ${QUANT} + python read_tensors.py --verify model-BF16.gguf model-${QUANT}.gguf \\ + >> quant_test_results.log + done +""" import os import sys import gguf -def read_tensors(path): + +def shape_str(tensor): + return "x".join(str(d) for d in tensor.shape) + + +def is_suspicious_shape(shape): + """ + Flags shapes with a dimension of size 1 anywhere except shape[0]. + + gguf.GGUFReader reports tensor.shape in GGML `ne` order, i.e. + reversed relative to the original torch/numpy shape. For example + txtfusion.projector.weight was (1, 12) in the safetensors file, but + shows up here as (12, 1) -- the "1" moved to the end. llama.cpp's + C++ quantizer determines a tensor's dimensionality via ggml_n_dims(), + which scans the `ne` array from the highest index downward and drops + trailing 1s -- so a 1 anywhere in shape[1:] is at risk of being + silently collapsed once the tensor passes through llama-quantize, + even though the GGUF written by convert.py still has the correct + shape. A 1 at shape[0] is safe, since the scan stops there. + + Only usable as an early warning on a single file (typically the + *-BF16.gguf straight out of convert.py); prefer compare mode when + you have both the pre- and post-quantize files, since that checks + the actual outcome instead of guessing. + """ + dims = list(shape) + if len(dims) < 2: + return False + return any(d == 1 for d in dims[1:]) + + +def print_single(path, quantized_only): reader = gguf.GGUFReader(path) + n_flagged = 0 for tensor in reader.tensors: - if tensor.tensor_type == gguf.GGMLQuantizationType.F32: + if quantized_only and tensor.tensor_type == gguf.GGMLQuantizationType.F32: continue - print(f"{str(tensor.tensor_type):32}: {tensor.name}") - -try: - path = sys.argv[1] - assert os.path.isfile(path), "Invalid path" - print(f"input: {path}") -except Exception as e: - input(f"failed: {e}") -else: - read_tensors(path) - input() + suspicious = is_suspicious_shape(tensor.shape) + flag = "" + if suspicious: + flag = " <-- dim other than shape[0] is 1, check for ndim collapse risk" + n_flagged += 1 + print(f"{str(tensor.tensor_type):26} {shape_str(tensor):18}: {tensor.name}{flag}") + if n_flagged: + print(f"\n{n_flagged} tensor(s) flagged as at risk of a ggml_n_dims() collapse.") + + +def print_compare(path_a, path_b, verify_only=False): + tensors_a = {t.name: t for t in gguf.GGUFReader(path_a).tensors} + tensors_b = {t.name: t for t in gguf.GGUFReader(path_b).tensors} + + names = sorted(set(tensors_a) | set(tensors_b)) + ndim_mismatches = [] + missing = [] + + for name in names: + a = tensors_a.get(name) + b = tensors_b.get(name) + + if a is None: + missing.append(f"{name} (only in B)") + if not verify_only: + print(f"only in B ({path_b}): {name}") + continue + if b is None: + missing.append(f"{name} (only in A)") + if not verify_only: + print(f"only in A ({path_a}): {name}") + continue + + flag = "" + if len(a.shape) != len(b.shape): + flag = " <-- DIM COUNT CHANGED (ggml_n_dims collapse!)" + ndim_mismatches.append(f"{name} (dim {len(a.shape)}->{len(b.shape)})") + + if not verify_only: + print( + f"{name:55} A: {str(a.tensor_type):10} {shape_str(a):14} " + f"B: {str(b.tensor_type):10} {shape_str(b):14}{flag}" + ) + + issues = ndim_mismatches + missing + + if verify_only: + # One compact line per run -- meant to be piped/redirected into a + # shared log file across a batch of quantization types, e.g.: + # python read_tensors.py --verify base.gguf quant.gguf >> quant_test.log + status = "OK " if not issues else "FAIL" + summary = f"{status} {os.path.basename(path_b):40} ({len(names)} tensors, {len(issues)} issue(s))" + if issues: + summary += ": " + "; ".join(issues) + print(summary) + return not issues + + print() + if ndim_mismatches: + print(f"!!! {len(ndim_mismatches)} tensor(s) changed dimensionality between the two files:") + for name in ndim_mismatches: + print(f" - {name}") + print("These are prime suspects for needing gguf_set_tensor_ndim() in lcpp.patch.") + else: + print("No dimensionality changes detected between the two files.") + return not issues + + +if __name__ == "__main__": + raw_args = sys.argv[1:] + quantized_only = "--quantized-only" in raw_args + verify_only = "--verify" in raw_args + paths = [a for a in raw_args if not a.startswith("--")] + + try: + assert len(paths) in (1, 2), "Usage: read_tensors.py [file_to_compare.gguf] [--quantized-only] [--verify]" + assert not (verify_only and len(paths) != 2), "--verify requires two files to compare" + for p in paths: + assert os.path.isfile(p), f"Invalid path: {p}" + except Exception as e: + input(f"failed: {e}") + sys.exit(1) + else: + if len(paths) == 1: + print(f"input: {paths[0]}") + print_single(paths[0], quantized_only) + input() + else: + if verify_only: + ok = print_compare(paths[0], paths[1], verify_only=True) + sys.exit(0 if ok else 1) + print(f"comparing:\n A: {paths[0]}\n B: {paths[1]}\n") + print_compare(paths[0], paths[1]) + input() diff --git a/ui.html b/ui.html new file mode 100644 index 00000000..740d44ab --- /dev/null +++ b/ui.html @@ -0,0 +1,995 @@ + + + + + + + JSON Prompt Generation Tool + + + + + + + +
+
+
+
+
+ + 1024 +
+ +
+
+
+ + 1024 +
+ +
+ +
+
+ +
+
+
+
+

Prompt Canvas

+ 1024 x 1024 +
+
+
+
+
+
+
+
+
+ +
+
+
+

Prompt JSON

+
+ + +
+
+ +
+
+
+
+ + +
+
+ + + + + + diff --git a/web/docs/CLIPLoaderGGUF.md b/web/docs/CLIPLoaderGGUF.md new file mode 100644 index 00000000..940b4a13 --- /dev/null +++ b/web/docs/CLIPLoaderGGUF.md @@ -0,0 +1,18 @@ +# CLIP Loader (GGUF) + +Loads one text encoder from a GGUF model file and returns a ComfyUI `CLIP`. + +## Parameters + +- **clip_name**: A GGUF text encoder from `models/clip` or + `models/text_encoders`. Regular compatible checkpoint files are also listed. +- **type**: Select the model family expected by the workflow, such as + `stable_diffusion`, `sd3`, or `krea2`. + +## Usage + +Connect **CLIP** to `CLIP Text Encode` nodes. The selected **type** must match +the diffusion model; for example, Krea 2 workflows require the `krea2` type. + +GGUF text encoders are loaded with the package's GGML operations and may use +less VRAM than full-precision checkpoints. diff --git a/web/docs/CLIPLoaderGGUFDynamicVRAM.md b/web/docs/CLIPLoaderGGUFDynamicVRAM.md new file mode 100644 index 00000000..6bf6bbeb --- /dev/null +++ b/web/docs/CLIPLoaderGGUFDynamicVRAM.md @@ -0,0 +1,5 @@ +# CLIPLoader (Dynamic VRAM) + +Loads one text encoder through ComfyUI's DynamicVRAM path. + +Use this node when DynamicVRAM is enabled and the text encoder benefits from demand loading. It accepts GGUF and regular text-encoder files. diff --git a/web/docs/DualCLIPLoaderGGUF.md b/web/docs/DualCLIPLoaderGGUF.md new file mode 100644 index 00000000..c55803f2 --- /dev/null +++ b/web/docs/DualCLIPLoaderGGUF.md @@ -0,0 +1,18 @@ +# Dual CLIP Loader (GGUF) + +Loads two text encoder files and returns one combined ComfyUI `CLIP` object. + +## Parameters + +- **clip_name1**: First encoder file. +- **clip_name2**: Second encoder file. +- **type**: The model family required by the downstream diffusion model. + +Files can be GGUF or compatible regular text-encoder checkpoints. Select the +same pair and model type normally used by the equivalent built-in Dual CLIP +loader. + +## Usage + +Connect the output to the workflow's text-encoding nodes. Both encoders must +match the selected model family. diff --git a/web/docs/DualCLIPLoaderGGUFDynamicVRAM.md b/web/docs/DualCLIPLoaderGGUFDynamicVRAM.md new file mode 100644 index 00000000..db6baaa4 --- /dev/null +++ b/web/docs/DualCLIPLoaderGGUFDynamicVRAM.md @@ -0,0 +1,5 @@ +# DualCLIPLoader (Dynamic VRAM) + +Loads two text encoders through ComfyUI's DynamicVRAM path. + +Use this node when DynamicVRAM is enabled and your workflow requires two CLIP/text-encoder inputs. diff --git a/web/docs/QuadrupleCLIPLoaderGGUF.md b/web/docs/QuadrupleCLIPLoaderGGUF.md new file mode 100644 index 00000000..eb84fa68 --- /dev/null +++ b/web/docs/QuadrupleCLIPLoaderGGUF.md @@ -0,0 +1,18 @@ +# Quadruple CLIP Loader (GGUF) + +Loads four text encoder files and returns one combined ComfyUI `CLIP` object. + +## Parameters + +- **clip_name1** +- **clip_name2** +- **clip_name3** +- **clip_name4** + +Each input accepts a GGUF text encoder or a compatible regular checkpoint. +The default model type is `stable_diffusion`. + +## Usage + +Use this node only for workflows whose model family requires four text +encoders. Connect **CLIP** to the workflow's normal text-encoding node. diff --git a/web/docs/QuadrupleCLIPLoaderGGUFDynamicVRAM.md b/web/docs/QuadrupleCLIPLoaderGGUFDynamicVRAM.md new file mode 100644 index 00000000..1fd2a8c7 --- /dev/null +++ b/web/docs/QuadrupleCLIPLoaderGGUFDynamicVRAM.md @@ -0,0 +1,5 @@ +# QuadrupleCLIPLoader (Dynamic VRAM) + +Loads four text encoders through ComfyUI's DynamicVRAM path. + +Use this node when DynamicVRAM is enabled and your workflow requires four CLIP/text-encoder inputs. diff --git a/web/docs/TripleCLIPLoaderGGUF.md b/web/docs/TripleCLIPLoaderGGUF.md new file mode 100644 index 00000000..e229f3c4 --- /dev/null +++ b/web/docs/TripleCLIPLoaderGGUF.md @@ -0,0 +1,17 @@ +# Triple CLIP Loader (GGUF) + +Loads three text encoder files and returns one combined ComfyUI `CLIP` object. + +## Parameters + +- **clip_name1** +- **clip_name2** +- **clip_name3** + +Each input accepts a GGUF text encoder or a compatible regular checkpoint. +This node uses the `sd3` model type by default. + +## Usage + +Use this node where an SD3-style workflow requires three text encoders. Connect +the resulting **CLIP** output to the normal text-encoding node. diff --git a/web/docs/TripleCLIPLoaderGGUFDynamicVRAM.md b/web/docs/TripleCLIPLoaderGGUFDynamicVRAM.md new file mode 100644 index 00000000..c946c261 --- /dev/null +++ b/web/docs/TripleCLIPLoaderGGUFDynamicVRAM.md @@ -0,0 +1,5 @@ +# TripleCLIPLoader (Dynamic VRAM) + +Loads three text encoders through ComfyUI's DynamicVRAM path. + +Use this node when DynamicVRAM is enabled and your workflow requires three CLIP/text-encoder inputs. diff --git a/web/docs/UnetLoaderGGUF.md b/web/docs/UnetLoaderGGUF.md new file mode 100644 index 00000000..41e73f22 --- /dev/null +++ b/web/docs/UnetLoaderGGUF.md @@ -0,0 +1,20 @@ +# Unet Loader (GGUF) + +Loads a diffusion model stored in GGUF format and returns a ComfyUI `MODEL`. + +Place GGUF diffusion models in `ComfyUI/models/unet` or +`ComfyUI/models/diffusion_models`, then select the file with **unet_name**. + +## Q8_CR models + +`Q8_CR` GGUFs use ComfyUI's native INT8 ConvRot path. Eligible Linear weights +remain INT8 during inference; 1-D, small, and designated high-precision tensors +are retained in FP32, while convolution weights remain FP16. + +CUDA is optional. CUDA uses the optimized ComfyUI backend when available; +non-CUDA environments use the eager backend and are slower. + +## Usage + +Connect **MODEL** to a sampler such as `KSampler`. Use the normal model-family +workflow inputs, including the matching text encoder and VAE. \ No newline at end of file diff --git a/web/docs/UnetLoaderGGUFAdvanced.md b/web/docs/UnetLoaderGGUFAdvanced.md new file mode 100644 index 00000000..de2beb1b --- /dev/null +++ b/web/docs/UnetLoaderGGUFAdvanced.md @@ -0,0 +1,26 @@ +# Unet Loader (GGUF/Advanced) + +Loads a GGUF diffusion model with the same behavior as **Unet Loader (GGUF)**, +with additional controls for dequantization and LoRA patch handling. + +## Parameters + +- **unet_name**: GGUF diffusion model in `models/unet` or + `models/diffusion_models`. +- **dequant_dtype**: Dtype used when a conventional GGUF quant needs + dequantization. Use **default** unless troubleshooting a model. +- **patch_dtype**: Dtype used when applying weight patches such as LoRAs. Use + **default** unless a specific workflow requires another dtype. +- **patch_on_device**: Applies patches on the model load device. Leave this off + for the usual lower-VRAM behavior. + +## Q8_CR + +For `Q8_CR` files, the loader selects ComfyUI's native INT8 ConvRot path. +The advanced dtype controls do not convert native INT8 weights to floating +point. + +## Usage + +Connect **MODEL** to the same downstream nodes as the standard GGUF UNet loader. +Start with all advanced options set to **default**. diff --git a/web/docs/UnetLoaderGGUFDynamicVRAM.md b/web/docs/UnetLoaderGGUFDynamicVRAM.md new file mode 100644 index 00000000..042b97b4 --- /dev/null +++ b/web/docs/UnetLoaderGGUFDynamicVRAM.md @@ -0,0 +1,7 @@ +# Unet Loader (Dynamic VRAM) + +Loads a GGUF diffusion model through ComfyUI's DynamicVRAM path. + +DynamicVRAM keeps quantized GGUF weights in host memory and loads them on demand. It requires a ComfyUI build with DynamicVRAM enabled; use the regular GGUF loader when DynamicVRAM is unavailable. + +For Q8_CR models, supported linear layers retain ComfyUI's native INT8 ConvRot layout.