89 lines
4.0 KiB
Markdown
89 lines
4.0 KiB
Markdown
# ComfyUI Quantization Node Pack
|
|
|
|
A custom node pack for ComfyUI that provides tools for quantizing model weights to lower precision formats like FP16, BF16, or true FP8 types.
|
|
|
|
## Overview
|
|
|
|
This node pack provides tools for quantizing models directly within ComfyUI. It includes:
|
|
|
|
1. **Model To State Dict**: Extracts the state dictionary from a model object and attempts to normalize keys.
|
|
2. **Quantize Model to FP8 Format**: Converts model weights directly to `float8_e4m3fn` or `float8_e5m2` format (requires CUDA).
|
|
3. **Quantize Model Scaled**: Applies simulated FP8 scaling (per-tensor or per-channel) and then casts the model to `float16`, `bfloat16`, or keeps the original format.
|
|
4. **Save As SafeTensor**: Saves the processed state dictionary to a `.safetensors` file at a specified path.
|
|
|
|
## Installation
|
|
|
|
1. Clone or download this repository into your ComfyUI's `custom_nodes` directory.
|
|
* Example using git:
|
|
```bash
|
|
cd ComfyUI/custom_nodes
|
|
git clone [https://github.com/YourUsername/YourRepoName.git](https://github.com/YourUsername/YourRepoName.git) ComfyUI_QuantNodes
|
|
# Replace with your actual repo URL and desired folder name
|
|
```
|
|
* Alternatively, download the ZIP and extract it into `ComfyUI/custom_nodes/ComfyUI_QuantNodes`.
|
|
2. Install dependencies (optional, check if already installed):
|
|
```bash
|
|
pip install -r ComfyUI/custom_nodes/ComfyUI_QuantNodes/requirements.txt
|
|
```
|
|
3. Restart ComfyUI.
|
|
|
|
## Usage
|
|
|
|
### Model To State Dict
|
|
* **Category:** `Model Quantization/Utils`
|
|
* **Function:** Extracts state dict from a MODEL object, stripping common prefixes.
|
|
* **Inputs**:
|
|
* `model`: The input `MODEL` object.
|
|
* **Outputs**:
|
|
* `model_state_dict`: The extracted state dictionary.
|
|
|
|
### Quantize Model to FP8 Format
|
|
* **Category:** `Model Quantization/FP8 Direct`
|
|
* **Function:** Converts model weights directly to a specific FP8 format. Requires CUDA.
|
|
* **Inputs**:
|
|
* `model_state_dict`: The state dictionary to quantize.
|
|
* `fp8_format`: The target FP8 format (`float8_e5m2` or `float8_e4m3fn`).
|
|
* **Outputs**:
|
|
* `quantized_model_state_dict`: The state dictionary with FP8 tensors.
|
|
|
|
### Quantize Model Scaled
|
|
* **Category:** `Model Quantization`
|
|
* **Function:** Applies simulated FP8 value scaling and then casts to FP16, BF16, or keeps the original dtype. Useful for size reduction with good compatibility.
|
|
* **Inputs**:
|
|
* `model_state_dict`: The state dictionary to quantize.
|
|
* `scaling_strategy`: How to simulate scaling (`per_tensor` or `per_channel`).
|
|
* `processing_device`: Where to perform calculations (`Auto`, `CPU`, `GPU`).
|
|
* `output_dtype`: Final data type (`Original`, `float16`, `bfloat16`). Defaults to `float16`.
|
|
* **Outputs**:
|
|
* `quantized_model_state_dict`: The processed state dictionary.
|
|
|
|
### Save As SafeTensor
|
|
* **Category:** `Model Quantization/Save`
|
|
* **Function:** Saves the processed state dictionary to a `.safetensors` file.
|
|
* **Inputs**:
|
|
* `quantized_model_state_dict`: The state dictionary to save.
|
|
* `absolute_save_path`: The full path (including filename) where the model will be saved.
|
|
* **Outputs**: None (Output node).
|
|
|
|
## Example Workflow
|
|
|
|
1. Load a model using a standard loader (e.g., `Load Checkpoint`).
|
|
2. Connect the `MODEL` output to the `Model To State Dict` node.
|
|
3. Connect the `model_state_dict` output from `Model To State Dict` to `Quantize Model Scaled`.
|
|
4. In `Quantize Model Scaled`, select your desired `scaling_strategy` and set `output_dtype` to `float16` (for size reduction).
|
|
5. Connect the `quantized_model_state_dict` output from `Quantize Model Scaled` to the `Save Model as SafeTensor` node.
|
|
6. Specify the `absolute_save_path` in the `Save Model as SafeTensor` node.
|
|
7. Queue the prompt.
|
|
8. Restart ComfyUI or refresh loaders to find the saved model.
|
|
|
|
## Requirements
|
|
|
|
* PyTorch (usually included with ComfyUI)
|
|
* `safetensors`
|
|
* `tqdm`
|
|
* CUDA-enabled GPU is required for the `Quantize Model scaled` node.
|
|
|
|
## License
|
|
|
|
MIT (Or your chosen license)
|