diff --git a/README.md b/README.md new file mode 100644 index 0000000..6f23b85 --- /dev/null +++ b/README.md @@ -0,0 +1,88 @@ +# 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 to FP8 Format` node. + +## License + +MIT (Or your chosen license) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..d3ae638 --- /dev/null +++ b/__init__.py @@ -0,0 +1,39 @@ +# __init__.py +# This file is necessary to make Python treat the directory as a package. +# It's also where ComfyUI looks for node mappings. + +# Import node classes from nodes.py +# Renamed QuantizeScaled to QuantizeModel +from .nodes import ModelToStateDict, QuantizeFP8Format, QuantizeModel, SaveAsSafeTensor + +# A dictionary that ComfyUI uses to map node class names to node objects +NODE_CLASS_MAPPINGS = { + # Node Utils + "ModelToStateDict": ModelToStateDict, + # Direct FP8 Conversion + "QuantizeFP8Format": QuantizeFP8Format, + # Scaled Quantization + Casting (FP16/BF16) + "QuantizeModel": QuantizeModel, # Renamed class QuantizeScaled -> QuantizeModel + # Saving Node + "SaveAsSafeTensor": SaveAsSafeTensor, +} + +# A dictionary that ComfyUI uses to map node class names to display names +NODE_DISPLAY_NAME_MAPPINGS = { + "ModelToStateDict": "Model To State Dict", + "QuantizeFP8Format": "Quantize Model to FP8 Format", + "QuantizeModel": "Quantize Model Scaled", # Display name for the renamed node + "SaveAsSafeTensor": "Save Model as SafeTensor", +} + +# Optional: Print a message to the console when the extension is loaded +print("----------------------------------------------------") +print("--- ComfyUI Quantization Node Pack Loaded ---") +print("--- Renamed QuantizeScaled to QuantizeModel ---") +print("--- Removed FP8 output from QuantizeModel ---") +print("--- Using User-Provided Scaling Logic ---") # Added note about scaling +print("--- Developed by [Lum3on] ---") # Remember to change this! +print("--- Version 0.5.0 ---") # Incremented version +print("----------------------------------------------------") + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/__pycache__/__init__.cpython-312.pyc b/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..07cdf27 Binary files /dev/null and b/__pycache__/__init__.cpython-312.pyc differ diff --git a/__pycache__/nodes.cpython-312.pyc b/__pycache__/nodes.cpython-312.pyc new file mode 100644 index 0000000..964ea03 Binary files /dev/null and b/__pycache__/nodes.cpython-312.pyc differ diff --git a/examples/workflow_quantize.json b/examples/workflow_quantize.json new file mode 100644 index 0000000..977cc1f --- /dev/null +++ b/examples/workflow_quantize.json @@ -0,0 +1,358 @@ +{ + "id": "0c605522-962a-4c82-9089-e97d4235261f", + "revision": 0, + "last_node_id": 14, + "last_link_id": 14, + "nodes": [ + { + "id": 5, + "type": "UNETLoader", + "pos": [ + 141.08615112304688, + 169.95050048828125 + ], + "size": [ + 287.13525390625, + 83.07095336914062 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 4 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.32", + "Node name for S&R": "UNETLoader", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [ + "flux\\flux1-fill-dev.safetensors", + "default" + ] + }, + { + "id": 11, + "type": "QuantizeFP8Format", + "pos": [ + 466.7684020996094, + 179.52532958984375 + ], + "size": [ + 323.7391357421875, + 59.94718933105469 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "model_state_dict", + "type": "MODEL_STATE_DICT", + "link": 11 + } + ], + "outputs": [ + { + "name": "quantized_model_state_dict", + "type": "MODEL_STATE_DICT", + "links": [ + 12 + ] + } + ], + "properties": { + "Node name for S&R": "QuantizeFP8Format", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [ + "float8_e4m3fn" + ] + }, + { + "id": 2, + "type": "ModelToStateDict", + "pos": [ + 142.7964630126953, + 307.8086853027344 + ], + "size": [ + 291.43902587890625, + 79.07095336914062 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 4 + } + ], + "outputs": [ + { + "name": "model_state_dict", + "type": "MODEL_STATE_DICT", + "slot_index": 0, + "links": [ + 11, + 13 + ] + } + ], + "properties": { + "Node name for S&R": "ModelToStateDict", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [] + }, + { + "id": 13, + "type": "QuantizeModel", + "pos": [ + 119.18730163574219, + 500.8162536621094 + ], + "size": [ + 345.158203125, + 106 + ], + "flags": {}, + "order": 4, + "mode": 4, + "inputs": [ + { + "name": "model_state_dict", + "type": "MODEL_STATE_DICT", + "link": 13 + } + ], + "outputs": [ + { + "name": "quantized_model_state_dict", + "type": "MODEL_STATE_DICT", + "links": [ + 14 + ] + } + ], + "properties": { + "Node name for S&R": "QuantizeModel", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [ + "per_tensor", + "Auto", + "float16" + ] + }, + { + "id": 4, + "type": "SaveAsSafeTensor", + "pos": [ + 471.4664611816406, + 311.7754211425781 + ], + "size": [ + 318.212890625, + 69.432373046875 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "quantized_model_state_dict", + "type": "MODEL_STATE_DICT", + "link": 12 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "SaveAsSafeTensor", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [ + "C:\\Users\\RAIIN Studios\\Documents\\protable\\ComfyUI\\models\\fluxfill-fp8e4m3fn.safetensors" + ] + }, + { + "id": 14, + "type": "SaveAsSafeTensor", + "pos": [ + 490.1119689941406, + 509.20977783203125 + ], + "size": [ + 363.431396484375, + 83.39691162109375 + ], + "flags": {}, + "order": 6, + "mode": 4, + "inputs": [ + { + "name": "quantized_model_state_dict", + "type": "MODEL_STATE_DICT", + "link": 14 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "SaveAsSafeTensor", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [ + "C:\\Users\\RAIIN Studios\\Documents\\protable\\ComfyUI\\models\\fluxfill-fp8e4m3fn.safetensors" + ] + }, + { + "id": 9, + "type": "Note", + "pos": [ + 911.549560546875, + 467.0907897949219 + ], + "size": [ + 371.62591552734375, + 149.91725158691406 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Available Scaling Strategies:\n\n- per_tensor : Uses a single scale factor for the entire tensor (fast, less precise).\n\n- per_channel : Computes a separate scale for each output channel (more accurate)." + ], + "color": "#432", + "bgcolor": "#653" + } + ], + "links": [ + [ + 4, + 5, + 0, + 2, + 0, + "MODEL" + ], + [ + 11, + 2, + 0, + 11, + 0, + "MODEL_STATE_DICT" + ], + [ + 12, + 11, + 0, + 4, + 0, + "MODEL_STATE_DICT" + ], + [ + 13, + 2, + 0, + 13, + 0, + "MODEL_STATE_DICT" + ], + [ + 14, + 13, + 0, + 14, + 0, + "MODEL_STATE_DICT" + ] + ], + "groups": [ + { + "id": 1, + "title": "fp8", + "bounding": [ + 131.08615112304688, + 96.35050201416016, + 680.042724609375, + 311.2386169433594 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + }, + { + "id": 2, + "title": "fp/bf16", + "bounding": [ + 109.18731689453125, + 427.21624755859375, + 775.532958984375, + 198.77328491210938 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + } + ], + "config": {}, + "extra": { + "frontendVersion": "1.18.9", + "ue_links": [], + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..75323c6 --- /dev/null +++ b/nodes.py @@ -0,0 +1,371 @@ +# nodes.py +# This file contains the implementation of your custom nodes. + +import torch +from safetensors.torch import save_file +from tqdm import tqdm +import os + +# --- Helper Classes (Using original simpler scaling logic as provided by user) --- + +class FP8Quantizer: + """A class to apply FP8 quantization to a state_dict.""" + + def __init__(self, quant_dtype: str = "float8_e5m2"): + if not hasattr(torch, quant_dtype): + raise ValueError(f"Unsupported quant_dtype: {quant_dtype}. PyTorch does not have this attribute.") + self.quant_dtype = quant_dtype + self.scale_factors = {} # Not used in current quantize_weights but kept for potential future use + + def quantize_weights(self, weight: torch.Tensor, layer_name: str) -> torch.Tensor: + """Quantizes a weight tensor to the specified FP8 format using simple scaling.""" + if not weight.is_floating_point(): + return weight # Only quantize floating point tensors + + original_device = weight.device + # Ensure FP8 conversion happens on CUDA if possible, as FP8 types require it + can_use_cuda = torch.cuda.is_available() + target_device = torch.device("cuda") if can_use_cuda else torch.device("cpu") + + # Warn if trying FP8 on CPU without CUDA + if not can_use_cuda and "float8" in self.quant_dtype: + print(f"[FP8Quantizer] Warning: CUDA not available. True {self.quant_dtype} conversion requires CUDA. Attempting on CPU, but results may be unexpected or errors may occur.") + target_device = torch.device("cpu") # Try on CPU despite warning + + weight_on_target = weight.to(target_device) + + max_val = torch.max(torch.abs(weight_on_target)) + if max_val == 0: + # For zero tensor, just cast to target dtype on the target device + target_torch_dtype = getattr(torch, self.quant_dtype) + return torch.zeros_like(weight_on_target, dtype=target_torch_dtype) + else: + # Using the simple scaling from user's provided script + scale = max_val / 127.0 + # Clamp scale to avoid division by zero if max_val is extremely small + scale = torch.max(scale, torch.tensor(1e-12, device=target_device, dtype=weight_on_target.dtype)) + + # Quantize: scale, round, unscale (simulated int8 range mapping) + quantized_weight_simulated = torch.round(weight_on_target / scale * 127.0) / 127.0 * scale + + # Final cast to the target FP8 dtype + target_torch_dtype = getattr(torch, self.quant_dtype) + quantized_weight = quantized_weight_simulated.to(dtype=target_torch_dtype) + + # Return on the device where conversion happened (target_device, ideally CUDA) + return quantized_weight + + def apply_quantization(self, state_dict: dict) -> dict: + """Applies direct FP8 quantization to all applicable weights.""" + quantized_state_dict = {} + eligible_tensors = {name: param for name, param in state_dict.items() if isinstance(param, torch.Tensor) and param.is_floating_point()} + progress_bar = tqdm(eligible_tensors.items(), desc=f"Quantizing to {self.quant_dtype}", unit="tensor", leave=False) + + for name, param in progress_bar: + # quantize_weights handles device logic now + quantized_state_dict[name] = self.quantize_weights(param.clone(), name) + + for name, param in state_dict.items(): + if name not in quantized_state_dict: + quantized_state_dict[name] = param + return quantized_state_dict + +class FP8ScaledQuantizer: + """ + Simulated FP8 quantizer using 8-bit scaled float approximation based on user provided script. + Operations are performed on the device of the input tensors. + (Used internally by QuantizeModel node for the value simulation step). + """ + def __init__(self, scaling_strategy: str = "per_tensor"): + self.scaling_strategy = scaling_strategy + self.scale_factors = {} # Stores the calculated scales (Python floats or lists) + + def _quantize_fp8_simulated(self, tensor: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Simulate quantization by scaling, clamping to 8-bit range, and dequantizing.""" + # Ensure scale is a tensor on the correct device and dtype + scale = scale.to(device=tensor.device, dtype=tensor.dtype) + # Prevent division by zero + scale = torch.where(scale == 0, torch.tensor(1e-9, device=tensor.device, dtype=tensor.dtype), scale) + + # Perform simulation: scale, round, clamp, unscale + quantized_intermediate = tensor / scale * 127.0 + quantized = torch.round(quantized_intermediate).clamp_(-127.0, 127.0) + dequantized = quantized / 127.0 * scale + return dequantized + + def quantize_weights(self, weight: torch.Tensor, layer_name: str) -> torch.Tensor: + """Applies the simulated quantization based on the chosen strategy.""" + if not isinstance(weight, torch.Tensor) or not weight.is_floating_point(): + return weight # Skip non-float tensors + + current_device = weight.device + + if self.scaling_strategy == "per_tensor": + scale_val = torch.max(torch.abs(weight)) + # Ensure scale_val is a tensor for consistent handling in _quantize_fp8_simulated + scale = scale_val if scale_val != 0 else torch.tensor(1.0, device=current_device, dtype=weight.dtype) + self.scale_factors[layer_name] = scale.item() # Store the scale value + quantized_weight = self._quantize_fp8_simulated(weight, scale) + + elif self.scaling_strategy == "per_channel": + if weight.ndim < 2: # Fallback to per-tensor for 1D tensors + scale_val = torch.max(torch.abs(weight)) + scale = scale_val if scale_val != 0 else torch.tensor(1.0, device=current_device, dtype=weight.dtype) + self.scale_factors[layer_name] = scale.item() + quantized_weight = self._quantize_fp8_simulated(weight, scale) + else: + # Assume channel dimension is 0 for typical Conv layers + # For Linear layers (e.g., [out_features, in_features]), dim 0 is also common. + # If weights are [in, out], dim 1 might be needed. Defaulting to dim 0. + channel_dim = 0 + dims_to_reduce = [d for d in range(weight.ndim) if d != channel_dim] + if not dims_to_reduce: # Handle edge case if channel_dim is the only dim somehow + scale_val = torch.max(torch.abs(weight)) + scale = scale_val if scale_val != 0 else torch.tensor(1.0, device=current_device, dtype=weight.dtype) + self.scale_factors[layer_name] = scale.item() + else: + scale = torch.amax(torch.abs(weight), dim=dims_to_reduce, keepdim=True) + # Store scales as a list of floats + self.scale_factors[layer_name] = scale.squeeze().tolist() + + quantized_weight = self._quantize_fp8_simulated(weight, scale) + else: + raise ValueError(f"Unknown scaling strategy: {self.scaling_strategy}") + + # The output tensor retains the original dtype but has modified values + return quantized_weight + + def apply_quantization(self, state_dict: dict) -> dict: + """Applies simulated FP8 quantization to all applicable weights.""" + quantized_state_dict = {} + # Process only floating point tensors + eligible_tensors = {name: param for name, param in state_dict.items() if isinstance(param, torch.Tensor) and param.is_floating_point()} + progress_bar = tqdm(eligible_tensors.items(), desc=f"Applying scaled ({self.scaling_strategy}) quantization", unit="tensor", leave=False) + + for name, param in progress_bar: + # Pass the clone to avoid modifying the original dict if errors occur mid-way + quantized_state_dict[name] = self.quantize_weights(param.clone(), name) + + # Add back non-floating point tensors and non-tensor data + for name, param in state_dict.items(): + if name not in quantized_state_dict: + quantized_state_dict[name] = param + return quantized_state_dict + +# --- ComfyUI Nodes --- + +class ModelToStateDict: + @classmethod + def INPUT_TYPES(s): return {"required": {"model": ("MODEL",)}} + RETURN_TYPES = ("MODEL_STATE_DICT",); RETURN_NAMES = ("model_state_dict",) + FUNCTION = "get_state_dict"; CATEGORY = "Model Quantization/Utils" + def get_state_dict(self, model): + print("[ModelToStateDict] Attempting to extract state_dict...") + if not hasattr(model, 'model') or not hasattr(model.model, 'state_dict'): + print("[ModelToStateDict] Error: Invalid MODEL structure."); return ({},) + try: + original_state_dict = model.model.state_dict() + print(f"[ModelToStateDict] Original keys sample: {list(original_state_dict.keys())[:5]}") + state_dict_to_return = original_state_dict; prefixes_to_try = ["diffusion_model.", "model."]; prefix_found = False + for prefix in prefixes_to_try: + num_keys = len(original_state_dict); + if num_keys == 0: break + matches = sum(1 for k in original_state_dict if k.startswith(prefix)) + if matches > 0 and (matches / num_keys > 0.5 or matches == num_keys): + print(f"[ModelToStateDict] Stripping prefix '{prefix}'...") + state_dict_to_return = {k[len(prefix):] if k.startswith(prefix) else k: v for k, v in original_state_dict.items()} + print(f"[ModelToStateDict] New keys sample: {list(state_dict_to_return.keys())[:5]}"); prefix_found = True; break + if not prefix_found: print("[ModelToStateDict] No common prefixes stripped.") + dtypes = {}; total = 0 + for k, v in state_dict_to_return.items(): + if isinstance(v, torch.Tensor): total += 1; dt = str(v.dtype); dtypes[dt] = dtypes.get(dt, 0) + 1 + print(f"[ModelToStateDict] DEBUG: Output Tensors: {total}, Dtypes: {dtypes}") + return (state_dict_to_return,) + except Exception as e: print(f"[ModelToStateDict] Error: {e}"); return ({},) + +class QuantizeFP8Format: # Direct FP8 conversion node + @classmethod + def INPUT_TYPES(s): return { "required": { "model_state_dict": ("MODEL_STATE_DICT",), "fp8_format": (["float8_e4m3fn", "float8_e5m2"], {"default": "float8_e5m2"}), } } + RETURN_TYPES = ("MODEL_STATE_DICT",); RETURN_NAMES = ("quantized_model_state_dict",) + FUNCTION = "quantize_model"; CATEGORY = "Model Quantization/FP8 Direct" + def quantize_model(self, model_state_dict: dict, fp8_format: str): + print(f"[QuantizeFP8Format] To {fp8_format}. Keys(sample): {list(model_state_dict.keys())[:3]}") + if not isinstance(model_state_dict, dict) or not model_state_dict: print("[QuantizeFP8Format] Invalid input."); return ({},) + try: + quantizer = FP8Quantizer(quant_dtype=fp8_format) # Uses helper class with simple scaling + quantized_state_dict = quantizer.apply_quantization(model_state_dict) + found = False; + for n, p in quantized_state_dict.items(): + if isinstance(p, torch.Tensor) and "float8" in str(p.dtype): print(f"[QuantizeFP8Format] Sample '{n}' dtype: {p.dtype}, dev: {p.device}"); found=True; break + if not found: print(f"[QuantizeFP8Format] No tensor converted to {fp8_format}.") + print("[QuantizeFP8Format] Complete."); return (quantized_state_dict,) + except Exception as e: print(f"[QuantizeFP8Format] Error: {e}"); return (model_state_dict,) + +class QuantizeModel: # <<< RENAMED CLASS from QuantizeScaled + """ + Applies simulated FP8 scaling (per-tensor/per-channel) and casts + to a specified output dtype (float16, bfloat16, or Original). + """ + # Removed "FP8" from the list as requested + OUTPUT_DTYPES_LIST = ["Original", "float16", "bfloat16"] + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model_state_dict": ("MODEL_STATE_DICT",), + "scaling_strategy": (["per_tensor", "per_channel"], {"default": "per_tensor"}), + "processing_device": (["Auto", "CPU", "GPU"], {"default": "Auto"}), + # Default to float16 for size reduction, user can choose others + "output_dtype": (s.OUTPUT_DTYPES_LIST, {"default": "float16"}), + } + } + + RETURN_TYPES = ("MODEL_STATE_DICT",) + RETURN_NAMES = ("quantized_model_state_dict",) + FUNCTION = "quantize_model_scaled" # Internal function name can stay + CATEGORY = "Model Quantization" # More general category + + def quantize_model_scaled(self, model_state_dict: dict, scaling_strategy: str, processing_device: str, output_dtype: str): + # Log using the new class name + print(f"[QuantizeModel] Strategy: {scaling_strategy}, Device: {processing_device}, Output Dtype: {output_dtype}") + + if not isinstance(model_state_dict, dict) or not model_state_dict: + print("[QuantizeModel] Error: Input model_state_dict is invalid."); + return (model_state_dict if isinstance(model_state_dict, dict) else {},) + + # Determine processing device + current_processing_device_str = "cpu" + if processing_device == "Auto": + first_tensor_device = next((p.device for p in model_state_dict.values() if isinstance(p, torch.Tensor)), torch.device("cpu")) + current_processing_device_str = str(first_tensor_device) + elif processing_device == "CPU": current_processing_device_str = "cpu" + elif processing_device == "GPU": + if torch.cuda.is_available(): current_processing_device_str = "cuda" + else: print("[QuantizeModel] Warning: GPU selected, CUDA unavailable. Defaulting to CPU."); current_processing_device_str = "cpu" + current_processing_device = torch.device(current_processing_device_str) + print(f"[QuantizeModel] Value scaling simulation target device: {current_processing_device}") + + # Move input state_dict to the processing device + state_dict_on_processing_device = {} + for name, param in model_state_dict.items(): + if isinstance(param, torch.Tensor): + state_dict_on_processing_device[name] = param.to(current_processing_device) + else: state_dict_on_processing_device[name] = param + + scaled_state_dict = {}; final_state_dict = {} + + try: + # Perform FP8 value simulation using FP8ScaledQuantizer helper (simple scaling version) + quantizer = FP8ScaledQuantizer(scaling_strategy=scaling_strategy) + scaled_state_dict = quantizer.apply_quantization(state_dict_on_processing_device) + print(f"[QuantizeModel] FP8 value scaling simulation performed on {current_processing_device}.") + + # Cast to final output_dtype (Original, float16, bfloat16) + if output_dtype == "Original": + print("[QuantizeModel] Output Dtype: Original. No further dtype casting.") + final_state_dict = scaled_state_dict + else: + # output_dtype is guaranteed to be 'float16' or 'bfloat16' + try: + target_torch_dtype = getattr(torch, output_dtype) + print(f"[QuantizeModel] Casting output to {output_dtype} ({target_torch_dtype})...") + for name, param in scaled_state_dict.items(): + if isinstance(param, torch.Tensor) and param.is_floating_point(): + final_state_dict[name] = param.to(dtype=target_torch_dtype) + else: + final_state_dict[name] = param # Pass non-float tensors or non-tensors + print(f"[QuantizeModel] Casting to {output_dtype} complete.") + except AttributeError: # Should not happen with the restricted list + print(f"[QuantizeModel] Error: Invalid torch dtype '{output_dtype}'. Using scaled tensors without final casting.") + final_state_dict = scaled_state_dict + except Exception as e_cast: + print(f"[QuantizeModel] Error during casting loop to {output_dtype}: {e_cast}. Using scaled tensors without final casting for affected tensors.") + for name_done, param_done in final_state_dict.items(): pass + for name_rem, param_rem in scaled_state_dict.items(): + if name_rem not in final_state_dict: final_state_dict[name_rem] = param_rem + + # Verification log + for name, param in final_state_dict.items(): + if isinstance(param, torch.Tensor) and param.is_floating_point(): + print(f"[QuantizeModel] Sample output tensor '{name}' final dtype: {param.dtype}, device: {param.device}") + break + print(f"[QuantizeModel] Processing complete.") + return (final_state_dict,) + except Exception as e: + print(f"[QuantizeModel] Major error during processing: {e}") + return (model_state_dict,) + + +class SaveAsSafeTensor: # No changes needed + @classmethod + def INPUT_TYPES(s): return { "required": { "quantized_model_state_dict": ("MODEL_STATE_DICT",), "absolute_save_path": ("STRING", {"default": "C:/temp/quantized_model.safetensors", "multiline": False}), } } + RETURN_TYPES = () ; OUTPUT_NODE = True ; FUNCTION = "save_model"; CATEGORY = "Model Quantization/Save" + def save_model(self, quantized_model_state_dict: dict, absolute_save_path: str): + print(f"[SaveAsSafeTensor] Saving to: {absolute_save_path}") + if not isinstance(quantized_model_state_dict, dict) or not quantized_model_state_dict: print("[SaveAsSafeTensor] Error: Input invalid."); return {"ui": {"text": ["Error: Input invalid."]}} + if not absolute_save_path: print("[SaveAsSafeTensor] Error: Path empty."); return {"ui": {"text": ["Error: Path empty."]}} + if not absolute_save_path.lower().endswith(".safetensors"): absolute_save_path += ".safetensors"; print(f"[SaveAsSafeTensor] Appended .safetensors") + try: + output_dir = os.path.dirname(absolute_save_path); + if output_dir and not os.path.exists(output_dir): os.makedirs(output_dir, exist_ok=True); print(f"[SaveAsSafeTensor] Created dir.") + cpu_state_dict = {}; dtype_counts = {}; total_tensors = 0 + for k, v in quantized_model_state_dict.items(): + if isinstance(v, torch.Tensor): + total_tensors += 1; tensor_to_save = v.cpu() if v.device.type != 'cpu' else v; cpu_state_dict[k] = tensor_to_save + dt_str = str(tensor_to_save.dtype); dtype_counts[dt_str] = dtype_counts.get(dt_str, 0) + 1 + else: cpu_state_dict[k] = v + print(f"[SaveAsSafeTensor] DEBUG: Tensors: {total_tensors}, Dtypes: {dtype_counts}") + save_file(cpu_state_dict, absolute_save_path) + print(f"[SaveAsSafeTensor] Saved successfully."); return {"ui": {"text": [f"Saved: {absolute_save_path}"]}} + except Exception as e: print(f"[SaveAsSafeTensor] Error saving: {e}"); return {"ui": {"text": [f"Error: {e}"]}} + +# --- Main (for testing outside ComfyUI, not strictly necessary for the plugin) --- +# (Test block remains the same as previous version, it already tests QuantizeModel) +if __name__ == '__main__': + print("--- Testing Quantization Nodes (Renamed QuantizeModel) ---") + + class MockCoreModel(torch.nn.Module): + def __init__(self): super().__init__(); self.layer1 = torch.nn.Linear(10,10).float(); self.layer2=torch.nn.Linear(10,10).float() + def forward(self,x): return self.layer2(self.layer1(x)) + def state_dict(self, *args, **kwargs): return {k:v.clone() for k,v in super().state_dict(*args,**kwargs).items()} + class MockModelWithPrefix(torch.nn.Module): + def __init__(self): super().__init__(); self.model = MockCoreModel() # Use "model." prefix + def forward(self,x): return self.model(x) + class MockModelPatcher: + def __init__(self): self.model = MockModelWithPrefix(); [p.data.normal_().float() for p in self.model.parameters()] + + mock_comfy_model = MockModelPatcher() + node_to_sd = ModelToStateDict() + base_sd_tuple = node_to_sd.get_state_dict(mock_comfy_model) + base_sd = base_sd_tuple[0] if base_sd_tuple else {} + if not base_sd or 'layer1.weight' not in base_sd: print("ModelToStateDict failed."); exit() # Check unprefixed key + print(f"Base SD 'layer1.weight' dtype: {base_sd['layer1.weight'].dtype}, device: {base_sd['layer1.weight'].device}") + + # Test the renamed node + node_quantize = QuantizeModel() + print("\n--- Test QuantizeModel ---") + + # Test Case 1: Output float16 + print("Testing QuantizeModel: Output Dtype = float16, Device = CPU") + result_fp16_tuple = node_quantize.quantize_model_scaled(base_sd, "per_tensor", "CPU", "float16") + result_fp16 = result_fp16_tuple[0] if result_fp16_tuple else {} + if result_fp16 and 'layer1.weight' in result_fp16: + tensor = result_fp16['layer1.weight'] + print(f" Output 'layer1.weight' dtype: {tensor.dtype} (Expected float16), device: {tensor.device}") + assert tensor.dtype == torch.float16 + assert tensor.device.type == 'cpu' + else: print(" Test Case 1 Failed.") + + # Test Case 2: Output Original + print("\nTesting QuantizeModel: Output Dtype = Original, Device = CPU") + result_orig_tuple = node_quantize.quantize_model_scaled(base_sd, "per_tensor", "CPU", "Original") + result_orig = result_orig_tuple[0] if result_orig_tuple else {} + if result_orig and 'layer1.weight' in result_orig: + tensor = result_orig['layer1.weight'] + print(f" Output 'layer1.weight' dtype: {tensor.dtype} (Expected {base_sd['layer1.weight'].dtype}), device: {tensor.device}") + assert tensor.dtype == base_sd['layer1.weight'].dtype # Should match original + assert tensor.device.type == 'cpu' + else: print(" Test Case 2 Failed.") + + print("\n--- Testing Complete ---") \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..8a3cd47 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +torch>=2.0.0 +safetensors>=0.3.1 +tqdm>=4.65.0