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:

     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):

   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)