commit 1dad2ab24267ce89db5bfc3924bcea2b555eccc6 Author: Licon Studio Date: Wed Apr 15 15:14:15 2026 +0800 Add files via upload diff --git a/README.md b/README.md new file mode 100644 index 0000000..0d9390c --- /dev/null +++ b/README.md @@ -0,0 +1,112 @@ +# SNR-Quant + +ComfyUI nodes for adaptive FP8/BF16 mixed precision quantization. + +## Features + +- **Automatic model quantization** - Quantize any diffusion model to a target size +- **Adaptive threshold selection** - Automatically detects model "fragility" and adjusts quantization strategy +- **SNR-based layer selection** - Protects sensitive layers, quantizes robust ones +- **Outlier protection** - Layers with extreme weights are kept in BF16 +- **Progress bar** - Shows quantization progress in ComfyUI console +- **CSV report** - Generates detailed report showing which layers were quantized +- **Re-quantization safe** - Already quantized FP8 layers are preserved (no size increase on re-run) + +## Installation + +1. Clone or copy this folder to `ComfyUI/custom_nodes/SNR-quant` +2. Restart ComfyUI + +Dependencies (usually already installed): +- `torch>=2.1.0` (for FP8 support) +- `safetensors>=0.4.0` +- `numpy>=1.24.0` + +## Usage + +### SNR Quant (FP8/BF16) + +Quantize a model to fit within a target size. + +**Inputs:** +- `model_name` - Select from `ComfyUI/models/diffusion_models` +- `target_size_GB` - Target size in GB (e.g., 22.0) + +**Output:** +- Saved directly to `ComfyUI/output/diffusion_models/quant_{model_name}_{target}GB.safetensors` +- CSV report: `ComfyUI/output/diffusion_models/quant_{model_name}_{target}GB_report.csv` + +**Example:** +``` +model_name: wan2.2_fp16.safetensors +target_size_GB: 22.0 +→ output: quant_wan2.2_fp16_22.0GB.safetensors +→ report: quant_wan2.2_fp16_22.0GB_report.csv +``` + +### SNR Model Analyzer + +Analyze a model's SNR and outlier distribution. + +**Input:** +- `model_name` - Select from `ComfyUI/models/diffusion_models` + +**Outputs:** +- `report` - Text report with statistics and quantization recommendations + +## How It Works + +``` +Load model → Analyze layers → Detect fragility → Select layers → Quantize → Save + ↓ ↓ ↓ + SNR + Outlier fragile/moderate/ FP8 or BF16 + robust per layer +``` + +### Fragility Detection + +| Fragility | SNR Range | SNR Median | Strategy | +|-----------|-----------|------------|----------| +| Fragile | < 2 dB | < 32 dB | Protect more (90th percentile) | +| Moderate | - | - | Balanced (75th percentile) | +| Robust | > 5 dB | > 34 dB | Quantize more (50th percentile) | + +### Layer Selection + +1. Filter layers by outlier threshold (adaptive percentile) +2. Sort by SNR (descending - highest SNR = safest to quantize) +3. Select layers until target size is reached + +### Size Calculation + +- Baseline: All 2D layers as BF16 (2 bytes/param) + 1D layers as BF16 +- Minimum: All 2D layers as FP8 (1 byte/param) + 1D layers as BF16 +- If target < minimum: auto-adjust to minimum + +### Re-quantization Behavior + +When quantizing an already quantized model: +- Existing FP8 layers are **preserved** (not converted back to BF16) +- Only remaining BF16 layers are considered for further quantization +- File size will not increase on re-run + +## Output Format + +Quantized models are saved as `.safetensors` with mixed precision: +- FP8 layers: `torch.float8_e4m3fn` +- BF16 layers: `torch.bfloat16` +- 1D layers (bias, norms): always BF16 + +## CSV Report Format + +The generated report includes: +- `layer_name` - Full layer path +- `precision` - Final precision (FP8, BF16, or FP8 (existing)) +- `params_numel` - Number of parameters +- `SNR_dB` - Signal-to-Noise ratio (higher = safer to quantize) +- `outlier` - Outlier index (lower = safer to quantize) +- `precision_change` - Whether layer was quantized or preserved + +## License + +Apache-2.0 diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..38db322 --- /dev/null +++ b/__init__.py @@ -0,0 +1,4 @@ +# (c) SNR-Quant || Apache-2.0 +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..ba31816 --- /dev/null +++ b/nodes.py @@ -0,0 +1,616 @@ +""" +SNR-Quant - ComfyUI nodes for adaptive FP8/ BF16 quantization. +""" +import os +import torch +import numpy as np +from safetensors.torch import load_file, save_file +from typing import List, Dict, Tuple, Optional +from dataclasses import dataclass +from enum import Enum + +import folder_paths +from comfy.utils import ProgressBar + + +class ModelFragility(Enum): + """Model fragility levels.""" + FRAGILE = "fragile" + MODERATE = "moderate" + ROBUST = "robust" + + +@dataclass +class LayerInfo: + """Information about a single layer.""" + name: str + numel: int + snr: float + outlier: float + + @property + def bf16_size(self) -> int: + return self.numel * 2 + + @property + def fp8_size(self) -> int: + return self.numel + + +def compute_snr(weight: torch.Tensor) -> float: + """Compute Signal-to-Noise Ratio for FP8 E4M3 quantization.""" + f32_w = weight.float() + + if torch.isnan(f32_w).any() or torch.isinf(f32_w).any(): + return float('inf') + + max_abs = torch.max(torch.abs(f32_w)).item() + if max_abs == 0: + return float('inf') + + scale = max_abs / 448.0 + q = (f32_w / scale).clamp(-448, 448) + q = q.to(torch.float8_e4m3fn).to(f32_w.dtype) + dq = q * scale + + signal_pwr = torch.mean(f32_w ** 2).item() + noise_pwr = torch.mean((f32_w - dq) ** 2).item() + + if noise_pwr == 0: + return float('inf') + + return 10 * np.log10(signal_pwr / noise_pwr) + + +def compute_outlier_index(weight: torch.Tensor) -> float: + """Compute outlier index = max_abs / std.""" + f32_w = weight.float() + max_abs = torch.max(torch.abs(f32_w)).item() + std = torch.std(f32_w).item() + if std == 0: + return 0.0 + return max_abs / std + + +def analyze_model(sd: Dict[str, torch.Tensor]) -> Tuple[List[LayerInfo], int, int]: + """Analyze all 2D weight layers. + + Returns: + layers: List of LayerInfo for quantizable (non-FP8) layers + total_params: Total params in 2D layers + existing_fp8_count: Number of already quantized FP8 layers + """ + layers = [] + total_params = 0 + existing_fp8_count = 0 + + for name, tensor in sd.items(): + if tensor.dim() == 2: + numel = tensor.numel() + total_params += numel + + # Skip already quantized FP8 layers + if tensor.dtype == torch.float8_e4m3fn: + existing_fp8_count += 1 + continue + + snr = compute_snr(tensor) + outlier = compute_outlier_index(tensor) + layers.append(LayerInfo(name=name, numel=numel, snr=snr, outlier=outlier)) + + return layers, total_params, existing_fp8_count + + +def detect_fragility(layers: List[LayerInfo]) -> Tuple[ModelFragility, float, float]: + """Detect model fragility based on SNR distribution.""" + snrs = [l.snr for l in layers] + snr_range = max(snrs) - min(snrs) + snr_median = np.median(snrs) + + if snr_range < 2 and snr_median < 32: + return ModelFragility.FRAGILE, snr_range, snr_median + elif snr_range > 5 and snr_median > 34: + return ModelFragility.ROBUST, snr_range, snr_median + else: + return ModelFragility.MODERATE, snr_range, snr_median + + +class SNRQuantNode: + """Quantize model to target size using FP8/ BF16 mixed precision.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_name": (folder_paths.get_filename_list("diffusion_models"), { + "tooltip": "Model from ComfyUI/models/diffusion_models" + }), + "target_size_GB": ("FLOAT", { + "default": 22.0, + "min": 1.0, + "step": 0.1, + "tooltip": "Target model size in GB (e.g., 22.0)" + }), + } + } + + RETURN_TYPES = () + OUTPUT_NODE = True + FUNCTION = "quantize" + CATEGORY = "SNR-Quant" + + def quantize(self, model_name: str, target_size_GB: float): + # Get paths + model_path = folder_paths.get_full_path("diffusion_models", model_name) + output_dir = folder_paths.get_output_directory() + # Create diffusion_models folder in output if it doesn't exist + diffusion_models_dir = os.path.join(output_dir, "diffusion_models") + if not os.path.exists(diffusion_models_dir): + os.makedirs(diffusion_models_dir) + # Replace path separators with underscores for flat output filename + base_name = os.path.splitext(model_name)[0].replace('\\', '_').replace('/', '_') + # Include target size in filename to distinguish multiple quantizations + output_path = os.path.join(diffusion_models_dir, f"quant_{base_name}_{target_size_GB}GB.safetensors") + + print(f"\n{'='*60}") + print(f"[SNR-Quant] Auto Quant - FP8/BF16 Mixed Precision") + print(f"{'='*60}") + print(f"[SNR-Quant] Model: {model_name}") + print(f"[SNR-Quant] Target size: {target_size_GB:.2f} GB") + print(f"{'='*60}") + + # Load model to CPU + print(f"\n[Step 1/6] Loading model to CPU...") + sd = load_file(model_path, device="cpu") + print(f" ✓ Loaded {len(sd)} tensors") + + # Detect input dtype + sample = next(iter(sd.values())) + input_dtype = sample.dtype + print(f" ✓ Input dtype: {input_dtype}") + + # Analyze layers + print(f"\n[Step 2/6] Analyzing layers (computing SNR and outliers)...") + layers, total_params, existing_fp8_count = analyze_model(sd) + if existing_fp8_count > 0: + print(f" ✓ Found {len(layers)} quantizable 2D layers ({existing_fp8_count} FP8 layers skipped)") + else: + print(f" ✓ Found {len(layers)} 2D weight layers") + + # Detect fragility + fragility, snr_range, snr_median = detect_fragility(layers) + outliers = [l.outlier for l in layers] + outlier_median = np.median(outliers) + + print(f"\n[Step 3/6] Detecting model fragility...") + print(f" SNR range: {snr_range:.2f} dB") + print(f" SNR median: {snr_median:.2f} dB") + print(f" Outlier med: {outlier_median:.1f}") + print(f" → Fragility: {fragility.value.upper()}") + + # Set adaptive percentile + if fragility.value == "fragile": + outlier_percentile = 90 + print(f" → Using 90th percentile (conservative)") + elif fragility.value == "robust": + outlier_percentile = 50 + print(f" → Using 50th percentile (aggressive)") + else: + outlier_percentile = 75 + print(f" → Using 75th percentile (balanced)") + + # Calculate sizes + # layers only contains quantizable 2D weights (excludes existing FP8) + existing_fp8_numel = sum(t.numel() for t in sd.values() if t.dim() == 2 and t.dtype == torch.float8_e4m3fn) + existing_fp8_size = existing_fp8_numel # FP8 = 1 byte/param + + size_2d_bf16 = sum(l.numel * 2 for l in layers) # quantizable 2D layers as BF16 + size_2d_fp8 = sum(l.numel for l in layers) # quantizable 2D layers as FP8 + size_1d_bf16 = sum(t.numel() * 2 for t in sd.values() if t.dim() != 2) # 1D layers always BF16 + + # Current total size (existing FP8 + quantizable as BF16 + 1D) + current_total_size = existing_fp8_size + size_2d_bf16 + size_1d_bf16 + + # Minimum achievable (all quantizable as FP8) + min_size = existing_fp8_size + size_2d_fp8 + size_1d_bf16 + + target_bytes = target_size_GB * 1024**3 + need_save = current_total_size - target_bytes + + print(f"\n[Step 4/6] Calculating size targets...") + print(f" Current model size: {current_total_size / 1024**3:.2f} GB") + print(f" └─ Existing FP8: {existing_fp8_size / 1024**3:.2f} GB") + print(f" └─ Quantizable BF16: {size_2d_bf16 / 1024**3:.2f} GB") + if existing_fp8_count > 0: + print(f" Note: {existing_fp8_count} FP8 layers already quantized (skipped)") + print(f" Minimum achievable: {min_size / 1024**3:.2f} GB") + print(f" Target size: {target_bytes / 1024**3:.2f} GB") + print(f" Need to save: {need_save / 1024**3:.2f} GB") + + # Check if target is achievable + if target_bytes < min_size: + print(f"\n ⚠ WARNING: Target {target_size_GB:.2f} GB is below minimum!") + print(f" Adjusting to {min_size / 1024**3:.2f} GB (all quantizable layers as FP8)") + target_bytes = min_size + need_save = current_total_size - min_size + + if need_save <= 0: + print(f"\n[Step 5/6] Target >= Current size, no quantization needed") + print(f" Saving all layers as-is...") + + # Build layer info lookup + layer_info_map = {l.name: l for l in layers} + + quantized = {} + for name, tensor in sd.items(): + if tensor.dim() == 2 and tensor.dtype == torch.float8_e4m3fn: + # Keep existing FP8 layers as is + quantized[name] = tensor + elif tensor.dim() == 2: + quantized[name] = tensor.to(torch.bfloat16) + else: + quantized[name] = tensor.to(torch.bfloat16) + save_file(quantized, output_path) + + # Generate report - all layers BF16 + report_path = os.path.join(diffusion_models_dir, f"quant_{base_name}_{target_size_GB}GB_report.csv") + print(f"\n Generating quantization report: {report_path}") + + report_lines = ["layer_name,precision,params_numel,SNR_dB,outlier,precision_change"] + + # Add existing FP8 layers (already quantized, skipped) + for name, tensor in sd.items(): + if tensor.dim() == 2 and tensor.dtype == torch.float8_e4m3fn: + report_lines.append(f"{name},FP8 (existing),{tensor.numel()},N/A,N/A,already quantized") + + # Add other layers + for name, tensor in sd.items(): + if tensor.dim() == 2 and name in layer_info_map: + l = layer_info_map[name] + report_lines.append(f"{name},BF16,{tensor.numel()},{l.snr:.2f},{l.outlier:.2f},unchanged") + elif tensor.dim() != 2: + report_lines.append(f"{name},BF16,{tensor.numel()},N/A,N/A,unchanged") + + with open(report_path, 'w', encoding='utf-8') as f: + f.write('\n'.join(report_lines)) + + print(f" ✓ Report saved: {report_path}") + + del sd, quantized + print(f"\n[Step 6/6] Saving to: {output_path}") + print(f" ✓ Done! Output: quant_{base_name}.safetensors") + print(f"{'='*60}\n") + return () + + # Iterative selection + print(f"\n[Step 5/6] Selecting layers for FP8 quantization...") + final_fp8 = [] + iteration = 0 + + while outlier_percentile >= 50: + iteration += 1 + threshold = np.percentile(outliers, outlier_percentile) + quantizable = [l for l in layers if l.outlier < threshold] + quantizable_sorted = sorted(quantizable, key=lambda x: x.snr, reverse=True) + max_save = sum(l.numel for l in quantizable) + + if max_save >= need_save: + saved = 0 + selected = 0 + for l in quantizable_sorted: + if saved >= need_save: + break + final_fp8.append(l.name) + saved += l.numel + selected += 1 + print(f" Iteration {iteration}: {selected} layers selected at {outlier_percentile}th percentile") + print(f" Saved: {saved / 1024**3:.2f} GB / {need_save / 1024**3:.2f} GB needed") + break + else: + print(f" Iteration {iteration}: Not reachable at {outlier_percentile}th ({max_save / 1024**3:.2f} GB < {need_save / 1024**3:.2f} GB)") + outlier_percentile -= 5 + + if not final_fp8: + quantizable_sorted = sorted(layers, key=lambda x: x.snr, reverse=True) + final_fp8 = [l.name for l in quantizable_sorted] + print(f" Using ALL {len(final_fp8)} layers for maximum quantization") + + fp8_set = set(final_fp8) + bf16_count_layers = len(layers) - len(fp8_set) + + # Calculate final sizes + final_2d_fp8 = sum(l.numel for l in layers if l.name in fp8_set) + final_2d_bf16 = sum(l.numel * 2 for l in layers if l.name not in fp8_set) + final_size_gb = (existing_fp8_size + final_2d_fp8 + final_2d_bf16 + size_1d_bf16) / 1024**3 + + print(f"\n Final configuration:") + print(f" FP8 (existing): {existing_fp8_count} ({existing_fp8_size / 1024**3:.2f} GB)") + print(f" FP8 (new): {len(fp8_set)} ({final_2d_fp8 / 1024**3:.2f} GB)") + print(f" BF16 (2D): {bf16_count_layers} ({final_2d_bf16 / 1024**3:.2f} GB)") + print(f" └─ Est. size: {final_size_gb:.2f} GB") + + # Generate quantization report before quantizing + report_path = os.path.join(diffusion_models_dir, f"quant_{base_name}_{target_size_GB}GB_report.csv") + print(f"\n Generating quantization report: {report_path}") + + # Build layer info lookup + layer_info_map = {l.name: l for l in layers} + + report_lines = ["layer_name,precision,params_numel,SNR_dB,outlier,precision_change"] + + # Add existing FP8 layers (already quantized, skipped) + for name, tensor in sd.items(): + if tensor.dim() == 2 and tensor.dtype == torch.float8_e4m3fn: + report_lines.append(f"{name},FP8 (existing),{tensor.numel()},N/A,N/A,already quantized") + + # Add quantizable layers + for l in layers: + if l.name in fp8_set: + report_lines.append(f"{l.name},FP8,{l.numel},{l.snr:.2f},{l.outlier:.2f},BF16→FP8") + else: + report_lines.append(f"{l.name},BF16,{l.numel},{l.snr:.2f},{l.outlier:.2f},unchanged") + + # Add 1D layers + for name, tensor in sd.items(): + if tensor.dim() != 2: + report_lines.append(f"{name},BF16,{tensor.numel()},N/A,N/A,unchanged") + + with open(report_path, 'w', encoding='utf-8') as f: + f.write('\n'.join(report_lines)) + + print(f" ✓ Report saved: {report_path}") + + # Quantize with progress bar + print(f"\n[Step 6/6] Quantizing model...") + quantized = {} + pbar = ProgressBar(len(sd)) + fp8_count = 0 + bf16_count = 0 + + for i, (name, tensor) in enumerate(sd.items()): + if tensor.dim() == 2 and tensor.dtype == torch.float8_e4m3fn: + # Already FP8 - keep as is + quantized[name] = tensor + fp8_count += 1 + elif tensor.dim() == 2 and name in fp8_set: + # 2D layer selected for FP8 quantization + quantized[name] = tensor.to(torch.float8_e4m3fn) + fp8_count += 1 + else: + # BF16: all 1D layers + 2D layers not selected for FP8 + quantized[name] = tensor.to(torch.bfloat16) + bf16_count += 1 + pbar.update_absolute(i + 1) + if (i + 1) % 100 == 0 or i == len(sd) - 1: + print(f" Progress: {i + 1}/{len(sd)} (FP8: {fp8_count}, BF16: {bf16_count})") + + # Save + print(f"\n Saving to: {output_path}") + save_file(quantized, output_path) + + # Cleanup + del quantized, sd + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + print(f" ✓ Quantization complete!") + print(f" Output: quant_{base_name}.safetensors") + print(f"{'='*60}\n") + return () + + +class SNRAnalyzeNode: + """Analyze model and print SNR/outlier statistics.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_name": (folder_paths.get_filename_list("diffusion_models"), { + "tooltip": "Model from ComfyUI/models/diffusion_models" + }), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("report",) + FUNCTION = "analyze" + CATEGORY = "SNR-Quant" + + def analyze(self, model_name: str): + model_path = folder_paths.get_full_path("diffusion_models", model_name) + + print(f"\n{'='*60}") + print(f"[SNR-Quant] Model Analyzer") + print(f"{'='*60}") + print(f"[SNR-Quant] Model: {model_name}") + + print(f"\n[Step 1/3] Loading model...") + sd = load_file(model_path, device="cpu") + print(f" ✓ Loaded {len(sd)} tensors") + + print(f"\n[Step 2/3] Analyzing layers...") + layers, total_params, existing_fp8_count = analyze_model(sd) + if existing_fp8_count > 0: + print(f" ✓ Found {len(layers)} quantizable 2D layers ({existing_fp8_count} FP8 layers skipped)") + else: + print(f" ✓ Found {len(layers)} 2D weight layers") + + print(f"\n[Step 3/3] Computing statistics...") + fragility, snr_range, snr_median = detect_fragility(layers) + outliers = [l.outlier for l in layers] + outlier_median = np.median(outliers) + snr_min = min([l.snr for l in layers]) + snr_max = max([l.snr for l in layers]) + + print(f" SNR: min={snr_min:.2f}, max={snr_max:.2f}, median={snr_median:.2f} dB") + print(f" Outliers: min={min(outliers):.2f}, max={max(outliers):.2f}, median={outlier_median:.1f}") + print(f" Fragility: {fragility.value}") + + # Build report - calculate actual model sizes + size_2d_bf16 = sum(l.numel * 2 for l in layers) # quantizable 2D as BF16 + size_2d_fp8 = sum(l.numel for l in layers) # quantizable 2D as FP8 + size_1d_bf16 = sum(t.numel() * 2 for t in sd.values() if t.dim() != 2) # 1D layers BF16 + + # Calculate already quantized FP8 layers size + existing_fp8_numel = sum(t.numel() for t in sd.values() if t.dim() == 2 and t.dtype == torch.float8_e4m3fn) + existing_fp8_size = existing_fp8_numel # FP8 = 1 byte/param + + # Current total size (with existing FP8 + all quantizable as BF16) + current_total_size = existing_fp8_size + size_2d_bf16 + size_1d_bf16 + + # Minimum achievable (all quantizable as FP8) + min_size = existing_fp8_size + size_2d_fp8 + size_1d_bf16 + + # Analyze SNR and outlier distribution for recommendations + snrs = [l.snr for l in layers] + outliers = [l.outlier for l in layers] + + # Count layers by SNR quality + snr_excellent = sum(1 for s in snrs if s >= 35) + snr_good = sum(1 for s in snrs if 32 <= s < 35) + snr_moderate = sum(1 for s in snrs if 30 <= s < 32) + snr_poor = sum(1 for s in snrs if s < 30) + + # Count layers by outlier risk + outlier_low = sum(1 for o in outliers if o < 15) + outlier_medium = sum(1 for o in outliers if 15 <= o < 30) + outlier_high = sum(1 for o in outliers if o >= 30) + + # Calculate safety analysis + safe_to_quantize = sum(1 for l in layers if l.snr >= 32 and l.outlier < 20) + moderate_risk = sum(1 for l in layers if (30 <= l.snr < 32 or 20 <= l.outlier < 30)) + high_risk = sum(1 for l in layers if l.snr < 30 or l.outlier >= 30) + + safe_ratio = safe_to_quantize / len(layers) if len(layers) > 0 else 0 + moderate_ratio = moderate_risk / len(layers) if len(layers) > 0 else 0 + high_risk_ratio = high_risk / len(layers) if len(layers) > 0 else 0 + + # Weight by layer size for accurate size estimation + safe_numel = sum(l.numel for l in layers if l.snr >= 32 and l.outlier < 20) + moderate_numel = sum(l.numel for l in layers if (30 <= l.snr < 32 or 20 <= l.outlier < 30)) + high_risk_numel = sum(l.numel for l in layers if l.snr < 30 or l.outlier >= 30) + + # Size calculations (savings from quantizing) + conservative_fp8_size = safe_numel # Only safe layers as FP8 + balanced_fp8_size = safe_numel + moderate_numel * 0.6 # Safe + 60% of moderate + aggressive_fp8_size = safe_numel + moderate_numel * 0.9 + high_risk_numel * 0.3 # Most except very risky + + # Final target sizes = current size - savings + conservative_target = current_total_size - conservative_fp8_size + balanced_target = current_total_size - balanced_fp8_size + aggressive_target = current_total_size - aggressive_fp8_size + + # Ensure targets are within achievable range + conservative_target = max(conservative_target, min_size) + balanced_target = max(balanced_target, min_size) + aggressive_target = max(aggressive_target, min_size) + + # Build strategy description + if len(layers) == 0: + strategy_note = "No quantizable layers remaining - model is already fully quantized" + elif high_risk_ratio > 0.3: + strategy_note = "High-risk layers detected - recommend conservative quantization" + elif safe_ratio > 0.5: + strategy_note = "Many safe-to-quantize layers - can use aggressive strategy" + else: + strategy_note = "Mixed layer characteristics - balanced approach recommended" + + # Handle case where all layers are already quantized + if len(layers) == 0: + lines = [ + f"Model: {model_name}", + f"Fragility: N/A (no quantizable layers)", + f"2D Layers: 0 + {existing_fp8_count} FP8 (already quantized)", + f"1D Layers: {len(sd) - existing_fp8_count}", + f"", + f"This model is already fully quantized to FP8.", + f"Current size: {current_total_size / 1024**3:.2f} GB", + f"No further quantization recommended.", + ] + report = "\n".join(lines) + print(f"\n{report}") + print(f"\n{'='*60}\n") + del sd + return (report,) + + lines = [ + f"Model: {model_name}", + f"Fragility: {fragility.value}", + f"2D Layers: {len(layers)} + {existing_fp8_count} FP8 (skipped)", + f"1D Layers: {len(sd) - len(layers) - existing_fp8_count}", + f"", + f"Note: {existing_fp8_count} FP8 layers already quantized (skipped in analysis)" if existing_fp8_count > 0 else "", + f"", + f"SNR Distribution:", + f" Min: {snr_min:.2f} dB", + f" Max: {snr_max:.2f} dB", + f" Median: {snr_median:.2f} dB", + f" Range: {snr_range:.2f} dB", + f"", + f"SNR Quality Breakdown:", + f" Excellent (>=35 dB): {snr_excellent} layers ({snr_excellent/max(len(layers),1)*100:.0f}%)", + f" Good (32-35 dB): {snr_good} layers ({snr_good/max(len(layers),1)*100:.0f}%)", + f" Moderate (30-32 dB): {snr_moderate} layers ({snr_moderate/max(len(layers),1)*100:.0f}%)", + f" Poor (<30 dB): {snr_poor} layers ({snr_poor/max(len(layers),1)*100:.0f}%)", + f"", + f"Outlier Distribution:", + f" Min: {min(outliers):.2f}", + f" Max: {max(outliers):.2f}", + f" Median: {outlier_median:.1f}", + f"", + f"Outlier Risk Breakdown:", + f" Low (<15): {outlier_low} layers ({outlier_low/max(len(layers),1)*100:.0f}%)", + f" Medium (15-30): {outlier_medium} layers ({outlier_medium/max(len(layers),1)*100:.0f}%)", + f" High (>=30): {outlier_high} layers ({outlier_high/max(len(layers),1)*100:.0f}%)", + f"", + f"Quantization Safety Analysis:", + f" Safe to quantize (SNR>=32 AND outlier<20): {safe_to_quantize} layers ({safe_ratio*100:.0f}%)", + f" Moderate risk: {moderate_risk} layers ({moderate_ratio*100:.0f}%)", + f" High risk (protect recommended): {high_risk} layers ({high_risk_ratio*100:.0f}%)", + f"", + f"Size estimates:", + f" Current model size: {current_total_size / 1024**3:.2f} GB", + f" └─ Existing FP8: {existing_fp8_size / 1024**3:.2f} GB", + f" └─ Quantizable BF16: {size_2d_bf16 / 1024**3:.2f} GB", + f" Minimum achievable: {min_size / 1024**3:.2f} GB (all quantizable as FP8)", + f"", + f"Quantization Recommendations:", + f" Strategy: {strategy_note}", + f"", + f" Recommended target sizes (final model size):", + f" Conservative: {conservative_target / 1024**3:.2f} GB (save {(current_total_size - conservative_target) / 1024**3:.2f} GB)", + f" └─ FP8 layers: only safe layers (SNR>=32, outlier<20)", + f" └─ Expected quality: Near-lossless", + f" └─ Best for: Production, quality-critical work", + f"", + f" Balanced: {balanced_target / 1024**3:.2f} GB (save {(current_total_size - balanced_target) / 1024**3:.2f} GB)", + f" └─ FP8 layers: safe + partial moderate-risk layers", + f" └─ Expected quality: Minimal degradation", + f" └─ Best for: General use, good quality/size tradeoff", + f"", + f" Aggressive: {aggressive_target / 1024**3:.2f} GB (save {(current_total_size - aggressive_target) / 1024**3:.2f} GB)", + f" └─ FP8 layers: all except high-risk", + f" └─ Expected quality: Acceptable for most cases", + f" └─ Best for: VRAM-constrained environments", + ] + report = "\n".join(lines) + + print(f"\n{report}") + print(f"\n{'='*60}\n") + del sd + + return (report,) + + +NODE_CLASS_MAPPINGS = { + "SNRQuant": SNRQuantNode, + "SNRAnalyze": SNRAnalyzeNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SNRQuant": "SNR Quant (FP8/ BF16)", + "SNRAnalyze": "SNR Model Analyzer", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..de5014f --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +torch>=2.1.0 +safetensors>=0.4.0 +numpy>=1.24.0