Restore NF4 low VRAM workflow

CHANGELOG.md

@@ -18,6 +18,10 @@ All notable changes to this project will be documented in this file.
 - Generation nodes now return 3 outputs: `(image, rewritten_prompt, status)` instead of just `(image,)`
 - Status messages provide better feedback about generation settings
 
+### Fixed
+- **Low VRAM NF4 Loader**: Resolved validation errors on 24GB/32GB cards by implementing a custom device map strategy that forces NF4 layers to GPU while allowing other components to offload to CPU.
+- **Device Mapping**: Added logic to prevent `bitsandbytes` from seeing 4-bit layers on CPU, which was causing crashes in Low VRAM mode.
+
 ### Technical Details
 - `rewritten_prompt`: STRING - The final prompt used for generation (either original or LLM-rewritten)
 - `status`: STRING - Human-readable status message about the generation process
@@ -30,3 +34,11 @@ All notable changes to this project will be documented in this file.
 - Official HunyuanImage-3.0 prompt enhancement with LLM APIs
 - Large image generation with CPU offload
 - Professional resolution presets with megapixel indicators
+
+## [Low VRAM Fix] - 2024-11-19
+
+### Fixed Low VRAM NF4 Loader
+- Resolved validation errors on 24GB/32GB cards by implementing a custom device map strategy that forces NF4 layers to GPU while allowing other components to offload to CPU.
+
+### Enhanced Device Mapping
+- Added logic to prevent `bitsandbytes` from seeing 4-bit layers on CPU, which was causing crashes in Low VRAM mode.

Docs/INT8_LOADING_GUIDE.md

@@ -0,0 +1,365 @@
+# INT8 Model Loading - Technical Notes
+
+## Session Date: November 25, 2025
+
+---
+
+## Key Learnings
+
+### 0. New Budget-Friendly Nodes
+
+To keep legacy nodes untouched, all GPU budget controls now live in dedicated nodes:
+
+- **Loader:** `Hunyuan 3 Loader (INT8 Budget)` exposes the GPU budget slider and records reserve metadata.
+- **Loader (NF4):** `Hunyuan 3 Loader (NF4 Low VRAM+)` mirrors the same behavior for NF4 checkpoints.
+- **Generators:** `Hunyuan 3 Generate (Large Budget)` and `Hunyuan 3 Generate (Low VRAM Budget)` understand the loader metadata, implement smarter "smart" mode checks, and emit MemoryTracker telemetry in the node status.
+
+Use these new nodes whenever you want budget sliders or telemetry output; existing nodes behave exactly as before this change.
+
+### 1. INT8 Model Variants
+
+There are **TWO** different INT8 quantization variants for HunyuanImage-3:
+
+#### **INT8-Full** (60GB on disk → 96GB loaded)
+- Everything quantized to INT8 on disk
+- During loading, critical layers are **dequantized** back to bfloat16
+- Final size: ~96GB (60GB INT8 weights + 36GB dequantized critical layers)
+- Metadata: `"quantization_method": "bitsandbytes_int8_full"`
+
+#### **INT8-Selective** (96GB on disk → 130GB loaded)
+- Only transformer weights quantized to INT8
+- VAE, attention, and critical layers already stored as bfloat16 on disk
+- Final size: ~130GB (96GB on disk + 34GB loading overhead)
+- Metadata: `"quantization_method": "bitsandbytes_int8"` (not "full")
+
+**CRITICAL**: Both variants end up with the same critical layers in bfloat16 after loading. The difference is whether we're dequantizing (Full) or preserving (Selective).
+
+---
+
+### 2. Bitsandbytes Configuration Issues
+
+#### **The `llm_int8_enable_fp32_cpu_offload` Conflict**
+
+**Problem**: 
+```python
+quant_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    llm_int8_enable_fp32_cpu_offload=True,  # ❌ CONFLICTS with device_map="auto"
+)
+```
+
+When using `device_map="auto"` with `max_memory` constraints, bitsandbytes validation fails with:
+```
+ValueError: Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM...
+```
+
+**Solution**:
+```python
+quant_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    llm_int8_skip_modules=skip_modules,
+    llm_int8_enable_fp32_cpu_offload=False,  # ✅ Let device_map handle offload
+)
+
+load_kwargs = dict(
+    device_map="auto",
+    max_memory={0: target_vram, "cpu": "100GiB"},  # We control offload
+    quantization_config=quant_config,
+)
+```
+
+**Why**: `llm_int8_enable_fp32_cpu_offload=True` is incompatible with `device_map="auto"`. We manage CPU offload ourselves via `max_memory` constraints.
+
+---
+
+### 3. Critical Layers That MUST Stay Full Precision
+
+The following layers cause severe corruption (colorful noise artifacts) if quantized to INT8:
+
+```python
+skip_modules = [
+    # VAE/Autoencoder (image encoding/decoding)
+    "vae", "model.vae", "vae.decoder", "vae.encoder",
+    "autoencoder", "model.autoencoder", "autoencoder.decoder", "autoencoder.encoder",
+    
+    # Vision components
+    "vision_model", "model.vision_model",
+    "vision_aligner", "model.vision_aligner",
+    
+    # Embeddings and core layers
+    "timestep_emb", "model.timestep_emb",
+    "patch_embed", "model.patch_embed",
+    "time_embed", "model.time_embed",
+    "time_embed_2", "model.time_embed_2",
+    "final_layer", "model.final_layer",
+    
+    # Language model components
+    "model.wte", "wte",
+    "model.ln_f", "ln_f",
+    "lm_head", "model.lm_head",
+    
+    # Attention layers
+    "attn.q_proj", "attn.k_proj", "attn.v_proj", "attn.o_proj",
+    "attn.qkv_proj", "attn.out_proj",
+    "self_attn", "cross_attn",
+]
+```
+
+**Important**: These are **pattern matches** - any layer name containing these strings will be skipped.
+
+---
+
+### 4. Memory Management Strategy
+
+#### **Target: 55% GPU Usage**
+
+For a 96GB GPU:
+- **Target**: ~52GB on GPU (55% of total)
+- **Offloaded**: ~44GB to CPU (45% of model)
+- **Reserved**: ~43GB free VRAM for inference pulsing
+
+```python
+if torch.cuda.is_available():
+    free, total = torch.cuda.mem_get_info(0)
+    target_model_vram = int(total * 0.55)  # Use 55% for model
+    max_memory = {
+        0: target_model_vram,
+        "cpu": "100GiB"
+    }
+```
+
+#### **Why 55%?**
+- 96GB model on 96GB GPU leaves no headroom
+- Inference causes temporary VRAM spikes (up to 12GB per megapixel)
+- Without headroom: 99% VRAM usage → stalls/OOM
+- With 55% target: ~43GB free → smooth inference
+
+#### **GPU Budget Slider (Loaders & Generators)**
+
+- The budget slider now lives in `Hunyuan 3 Loader (NF4 Low VRAM+)` and `Hunyuan 3 Loader (INT8 Budget)`.
+- Set this to the VRAM you are comfortable dedicating to weights; anything beyond is auto-offloaded to CPU RAM.
+- Suggested starting points:
+    - 24GB GPUs → **18‑20GB** target
+    - 32GB GPUs → **26‑28GB** target
+    - 96GB GPUs → **52GB** (default)
+- Large/offload generators (`Hunyuan 3 Generate (Large Budget)` / `Low VRAM Budget`) also include an optional per-run override so you can experiment with different budgets without reloading the model.
+
+---
+
+### 5. Smart Mode Detection (Fixed)
+
+**Old "Stupid" Logic**: Tried to detect model type and predict requirements
+**New "Actually Smart" Logic**: Just check current VRAM state
+
+```python
+if offload_mode == "smart" and torch.cuda.is_available():
+    megapixels = (width * height) / 1_000_000
+    free_bytes, total_bytes = torch.cuda.mem_get_info(0)
+    free_gb = free_bytes / 1024**3
+    allocated_gb = (total_bytes - free_bytes) / 1024**3
+    percent_used = (allocated_gb / total_gb) * 100
+    
+    loader_reserve_gb = getattr(model, '_loader_reserve_gb', 0)
+    inference_requirement = (megapixels * 12.0) + 5.0 + loader_reserve_gb
+    
+    # Three simple rules:
+    if free_gb < 25.0:  # Less than 25GB free (allows 5090/4090 to skip)
+        should_offload = True
+    elif percent_used > 70.0:  # Using more than 70% of GPU
+        should_offload = True
+    elif free_gb < inference_requirement:  # Not enough for inference
+        should_offload = True
+```
+
+**Key Insight**: Model is already loaded - don't try to predict model size, just check if there's enough room for inference.
+
+---
+
+### 6. Terminology Clarification
+
+#### **"Keep in bfloat16" vs "Dequantize to bfloat16"**
+
+**Wrong terminology** (what we kept saying):
+
+- "Keep VAE in bfloat16"
+
+**Correct terminology**:
+
+- **INT8-Full**: "Dequantize/expand VAE from INT8 → bfloat16"
+- **INT8-Selective**: "Preserve VAE as bfloat16 (already full precision)"
+
+The `llm_int8_skip_modules` parameter does **different things** depending on what's on disk:
+
+1. If layer is INT8 on disk → causes **dequantization** to bfloat16
+2. If layer is bfloat16 on disk → **prevents quantization**, keeps as-is
+
+Same parameter, different behavior based on checkpoint content.
+
+---
+
+### 7. Performance Expectations
+
+**Target**: INT8 should be **2-3x faster** than BF16
+
+- BF16: ~6 minutes for 1MP image
+- INT8 Expected: ~2-3 minutes
+- INT8 Reality (before fixes): 12 minutes (2x SLOWER!)
+
+**Why it was slow**:
+
+- Loading 81GB to GPU (85% usage)
+- During inference: 99% VRAM → constant swapping/stalls
+- OOM after a few minutes
+
+**After fixes** (expected):
+
+- Loading ~52GB to GPU (55% usage)
+- During inference: ~70-80% VRAM (room to pulse)
+- Should hit 2-3x speedup target
+
+---
+
+### 8. Debug Verification Added
+
+Added logging to verify which layers are actually in which dtype:
+
+```python
+logger.info("Verifying layer dtypes after load:")
+dtype_summary = {}
+for name, param in model.named_parameters():
+    dtype = str(param.dtype)
+    if any(key in name.lower() for key in ['vae', 'autoencoder', 'vision', 'embed', 'final', 'lm_head', 'wte', 'ln_f']):
+        dtype_summary[dtype].append(name)
+
+for dtype, layers in dtype_summary.items():
+    logger.info(f"  {dtype}: {len(layers)} layers")
+    logger.info(f"    Examples: {', '.join(layers[:3])}")
+```
+
+This shows if skip_modules patterns are actually matching the real layer names.
+
+---
+
+### 9. Common Errors & Solutions
+
+#### **Error**: "Cannot copy out of meta tensor"
+
+**Cause**: Model already has meta tensors, trying to apply cpu_offload() again
+**Solution**: Check for `has_meta` flag before applying cpu_offload
+
+#### **Error**: "Some modules are dispatched on the CPU..."
+
+**Cause**: `llm_int8_enable_fp32_cpu_offload=True` with `device_map="auto"`
+**Solution**: Set to `False`, manage offload via max_memory
+
+#### **Error**: VAE corruption (colorful noise)
+
+**Cause**: VAE layers still quantized to INT8
+**Solution**: Verify skip_modules patterns match actual layer names, check dtype verification log
+
+#### **Error**: 99% VRAM usage, OOM
+
+**Cause**: Model loaded too close to GPU limit (81GB on 96GB GPU)
+**Solution**: Use max_memory to target 55% GPU usage
+
+---
+
+### 10. Telemetry & Status Output
+
+- All generate nodes (standard, large/offload, low VRAM) now track **peak GPU VRAM** and **system RAM** for every run.
+- The values are logged and appended to the node's status string, e.g. `... | Peak VRAM 27.8GB | Peak RAM 42.1GB`.
+- Use those stats to tune the GPU budget slider—keep peaks ~10GB under your physical limit to leave room for inference spikes.
+
+---
+
+## Optimal Device Map (Manual Configuration)
+
+If you want to manually control layer placement instead of using `device_map="auto"`, here's a recommended optimal mapping for 96GB GPU:
+
+```python
+optimal_device_map = {
+    # Critical layers on GPU (must be fast)
+    "vision_model": 0,
+    "vision_aligner": 0,
+    "timestep_emb": 0,
+    "patch_embed": 0,
+    "time_embed": 0,
+    "final_layer": 0,
+    "time_embed_2": 0,
+    "model.wte": 0,
+    "model.ln_f": 0,
+    "lm_head": 0,
+    
+    # First 6 transformer layers on GPU
+    "model.layers.0": 0,
+    "model.layers.1": 0,
+    "model.layers.2": 0,
+    "model.layers.3": 0,
+    "model.layers.4": 0,
+    "model.layers.5": 0,
+    
+    # Remaining layers offloaded to CPU
+    "model.layers.6": "cpu",
+    "model.layers.7": "cpu",
+    "model.layers.8": "cpu",
+    "model.layers.9": "cpu",
+    "model.layers.10": "cpu",
+    "model.layers.11": "cpu",
+    "model.layers.12": "cpu",
+    "model.layers.13": "cpu",
+    "model.layers.14": "cpu",
+    "model.layers.15": "cpu",
+    "model.layers.16": "cpu",
+    "model.layers.17": "cpu",
+    "model.layers.18": "cpu",
+    "model.layers.19": "cpu",
+    "model.layers.20": "cpu",
+    "model.layers.21": "cpu",
+    "model.layers.22": "cpu",
+    "model.layers.23": "cpu",
+    "model.layers.25": "cpu",
+    "model.layers.26": "cpu",
+    "model.layers.27": "cpu",
+    "model.layers.28": "cpu",
+    "model.layers.29": "cpu",
+    "model.layers.30": "cpu",
+    "model.layers.31": "cpu",
+    "model.layers.32": "cpu",
+}
+```
+
+**Note**: This keeps critical components and first 6 transformer layers on GPU, offloading the bulk of transformer layers to CPU. This provides good balance between speed and VRAM usage.
+
+---
+
+## Code Changes Summary
+
+### File: `hunyuan_quantized_nodes.py`
+
+1. **Lines 773-831**: Unified skip_modules list for both variants
+2. **Lines 833-837**: Disabled `llm_int8_enable_fp32_cpu_offload` (set to False)
+3. **Lines 852-875**: INT8-Full uses 55% GPU target with max_memory
+4. **Lines 876-898**: INT8-Selective uses 55% GPU target with max_memory
+5. **Lines 912-928**: Added dtype verification logging
+6. **Lines 1623-1650**: Simplified smart mode detection (removed model size prediction)
+
+---
+
+## Testing Checklist
+
+- [ ] INT8-Full loads to ~52GB (not 81GB)
+- [ ] INT8-Selective loads without bitsandbytes validation error
+- [ ] No VAE corruption (colorful noise) in generated images
+- [ ] VRAM stays below 85% during inference (not 99%)
+- [ ] Performance is 2-3x faster than BF16 (not slower)
+- [ ] Smart mode correctly detects when offload is needed
+- [ ] Dtype verification log shows VAE/critical layers in bfloat16
+
+---
+
+## Notes for Future
+
+- Both INT8 variants have similar final memory footprint (~96-130GB)
+- The 55% GPU target is optimized for 96GB GPU - may need adjustment for other GPU sizes
+- Smart mode threshold is 25GB free (works for 5090/4090, adjust if needed)
+- Skip_modules uses pattern matching - if corruption occurs, check actual layer names with dtype verification log

Eric_Hunyuan3.code-workspace

@@ -11,7 +11,12 @@
 		},
 		{
 			"path": "../../models/HunyuanImage-3-NF4"
+		},
+		{
+			"path": "../../models/HunyuanImage-3-INT8"
 		}
 	],
-	"settings": {}
+	"settings": {
+		"powershell.cwd": "Eric_Hunyuan3"
+	}
 }

README.md

@@ -6,15 +6,21 @@
 
 Professional ComfyUI custom nodes for [Tencent HunyuanImage-3.0](https://github.com/Tencent-Hunyuan/HunyuanImage-3.0), the powerful 80B parameter native multimodal image generation model.
 
+> **Latest**
+>
+> - ✅ **NF4 Low VRAM Loader + Low VRAM Budget generator are verified on 24‑32 GB cards** thanks to the custom device-map strategy that pins every quantized layer on GPU.
+> - ✅ **INT8 Budget loader works end-to-end** and a pre-quantized checkpoint is available on Hugging Face: [EricRollei/Hunyuan_Image_3_Int8](https://huggingface.co/EricRollei/Hunyuan_Image_3_Int8) for anyone who wants INT8 quality without running the quantizer locally.
+> - 📸 Added a reference workflow (below) showing the recommended node pair for Low VRAM setups.
+
 > **🙏 Acknowledgment**: This project integrates the HunyuanImage-3.0 model developed by **Tencent Hunyuan Team** and uses their official system prompts. The model and original code are licensed under [Apache 2.0](https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/LICENSE). This integration code is separately licensed under CC BY-NC 4.0 for non-commercial use.
 
 ## 📋 TODO / Known Issues
 
-- [ ] **INT8 Loader**: Currently non-functional due to bitsandbytes validation issues with CPU offload. Use NF4 or BF16 loaders instead.
-- [ ] Add example workflows to repository
-- [ ] Add screenshots/documentation for each node
-- [ ] Test and document multi-GPU setup
-- [ ] Optimize Low VRAM node for 24GB cards
+- [x] **NF4 Low VRAM Loader**: Custom device map keeps NF4 layers on GPU so 24‑32 GB cards can use the Low VRAM Budget workflow without bitsandbytes errors.
+- [x] Provide example workflow/screenshot for Low VRAM users (see below).
+- [ ] Add screenshots/documentation for every node (in progress).
+- [ ] Test and document multi-GPU setup.
+- [ ] Continue long-run stability testing on the INT8 Budget loader with CPU offload edge cases.
 
 ## 🎯 Features
 
@@ -88,6 +94,7 @@ This version is compressed to 8-bit, offering near-original quality with reduced
 - **Performance**: 
   - **Quality**: ~98% of full precision (better than NF4).
   - **Speed**: Faster inference than NF4 (less dequantization overhead) but requires more memory transfer if offloading.
+  - **Practical Requirement**: For the selective INT8 checkpoints shipped here, you realistically need a 96GB-class GPU (RTX 6000 Pro Blackwell, H100, etc.). Forcing CPU offload on smaller cards makes each step take minutes because the quantized tensors continually stream over PCIe.
 
 ### Quick Install
 
@@ -122,7 +129,17 @@ cd ../../models
 huggingface-cli download EricRollei/HunyuanImage-3-NF4-ComfyUI --local-dir HunyuanImage-3-NF4
 ```
 
-**Option C: Quantize Yourself (from Full Model)**
+**Option C: Download Pre-Quantized INT8 Model (~85GB)**
+If you want maximum fidelity without running the INT8 quantizer locally, grab the ready-to-go checkpoint:
+[EricRollei/Hunyuan_Image_3_Int8](https://huggingface.co/EricRollei/Hunyuan_Image_3_Int8)
+
+```bash
+# Download to ComfyUI/models/
+cd ../../models
+huggingface-cli download EricRollei/Hunyuan_Image_3_Int8 --local-dir HunyuanImage-3-INT8
+```
+
+**Option D: Quantize Yourself (from Full Model)**
 If you prefer to quantize it yourself:
 ```bash
 # First download full model (Option A), then quantize
@@ -136,20 +153,49 @@ python hunyuan_quantize_nf4.py \
 
 ## 🚀 Usage
 
+### 📸 Low VRAM Workflow Reference
+
+![Low VRAM Workflow](Docs/LowVRAM-Workflow-Image.png)
+
+Pair **Hunyuan 3 Loader (NF4 Low VRAM+)** with **Hunyuan 3 Generate (Low VRAM Budget)** for 24‑32 GB cards. The loader’s GPU budget slider keeps 18–20 GB free for inference while the generator adds smart telemetry and optional prompt rewriting. The screenshot above shows a full working chain (loader → rewriter → Low VRAM Budget generator → save/display) producing a 1600×1600 render on a 24 GB test rig without triggering bitsandbytes validation.
+
 ### Node Overview
 
 | Node Name | Purpose | VRAM Required | Speed |
 |-----------|---------|---------------|-------|
 | **Hunyuan 3 Loader (NF4)** | Load quantized model | ~45GB | Fast load |
+| **Hunyuan 3 Loader (NF4 Low VRAM+)** | NF4 loader with GPU budget slider for 24‑32GB GPUs | 18‑28GB (configurable) | Balanced |
+| **Hunyuan 3 Loader (INT8 Budget)** | INT8 loader with GPU budget & telemetry metadata | 20‑55GB (configurable) | Balanced |
 | **Hunyuan 3 Loader (Full BF16)** | Load full precision model | ~80GB | Moderate |
 | **Hunyuan 3 Loader (Full BF16 GPU)** | Single GPU with memory control | ~75GB+ | Moderate |
 | **Hunyuan 3 Loader (Multi-GPU BF16)** | Distribute across GPUs | 80GB total | Fast |
 | **Hunyuan 3 Loader (88GB GPU Optimized)** | **DEPRECATED** - Use Full BF16 Loader | - | - |
 | **Hunyuan 3 Generate** | Standard generation (<2MP) | Varies | **Fast** ⚡ |
+| **Hunyuan 3 Generate (Telemetry)** | Adds RAM/VRAM stats to status output | Varies | Fast |
 | **Hunyuan 3 Generate (Large/Offload)** | Large images (2-8MP+) | Varies | Moderate |
+| **Hunyuan 3 Generate (Large Budget)** | Large/offload with GPU budget override + telemetry | Varies | Moderate |
+| **Hunyuan 3 Generate (Low VRAM)** | Quantized-friendly large generation | Varies | Moderate |
+| **Hunyuan 3 Generate (Low VRAM Budget)** | Low VRAM mode with smart heuristics + telemetry | Varies | Moderate |
 | **Hunyuan 3 Unload** | Free VRAM | - | Instant |
 | **Hunyuan 3 GPU Info** | Diagnostic/GPU detection | - | Instant |
 
+> 🆕 **Budget/Telemetry Nodes**: New loader/generator variants keep the legacy nodes untouched while exposing GPU budget sliders, smarter "smart" mode logic, and MemoryTracker telemetry in the node status. Mix and match them only if you need the extra controls; existing workflows can stay on the classic nodes.
+
+### INT8 Budget Workflow Controls
+
+When using the **Hunyuan 3 Loader (INT8 Budget)** + **Hunyuan 3 Generate (Large Budget)** path (recommended for 96 GB RTX 6000 Pro / similar cards), keep these controls in mind:
+
+- `reserve_memory_gb` (loader): Amount of VRAM the loader leaves free for inference overhead. Set to ~20 GB for 1.5–2 MP work; increase when targeting 4K+ so Smart mode has headroom without forcing CPU offload.
+- `gpu_memory_target_gb` (loader): How much VRAM the INT8 weights are allowed to occupy. 80 GB is a good sweet spot on 96 GB GPUs. Lower values keep more VRAM free but may require CPU spillover during load.
+- `offload_mode` (Large Budget generator):
+  - `smart` (default) reads the loader’s reserve/budget metadata, estimates per-megapixel cost (~15 GB/MP + 6 GB), and only enables accelerate CPU offload when VRAM is critically low.
+  - `disabled` forces fully-on-GPU sampling (fastest) – use this for ≤2 MP jobs when you already fit in VRAM.
+  - `always` moves the model into accelerate offload before sampling. Only select this for extreme resolutions where you expect to exceed available VRAM; otherwise it slows INT8 inference to CPU speeds.
+- `keep_model_loaded` (Large Budget generator): When `False`, the node calls the same cache-clearing logic as **Hunyuan 3 Unload** after a run, freeing the INT8 weights from VRAM/RAM. Leave it `True` for iterative workflows to avoid reload times.
+- `gpu_budget_gb` override (Large Budget generator): Temporarily overrides the loader’s `_loader_gpu_budget_gb` value for Smart-mode math without remapping tensors. Set it if you want Smart mode to assume a different baseline (e.g., 65 GB) without reloading the model; leave at `-1` to use the loader’s recorded budget.
+
+These generator settings never remap the model by themselves—they only influence whether accelerate offload engages and whether we clear the cached model afterward. Adjust the loader slider first (20 GB reserve / 80 GB target recommended), then use the generator controls to decide when to keep models in memory or trigger CPU offload for very large renders.
+
 ### Node Compatibility Guide
 
 **⚠️ IMPORTANT**: Match the Loader to the correct Generate node for stability.

hunyuan_shared.py

@@ -25,6 +25,7 @@ Note: HunyuanImage-3.0 model is subject to Tencent's Apache 2.0 license.
 import logging
 from typing import Optional
 
+import psutil
 import torch
 
 logger = logging.getLogger(__name__)
@@ -33,6 +34,8 @@ _SDPA_PATCHED = False
 _ORIGINAL_SDPA = None
 _GELU_PATCHED = False
 _ORIGINAL_GELU = None
+_SILU_PATCHED = False
+_ORIGINAL_SILU = None
 _SCATTER_PATCHED = False
 _ORIGINAL_SCATTER = None
 _ORIGINAL_SCATTER_INPLACE = None
@@ -41,6 +44,41 @@ _ORIGINAL_INDEX_COPY = None
 _ORIGINAL_INDEX_COPY_INPLACE = None
 _DYNAMIC_CACHE_PATCHED = False
 _ORIGINAL_DYNAMIC_CACHE_UPDATE = None
+_DTYPE_HOOKS_INSTALLED = False
+
+
+def install_dtype_harmonization_hooks(model) -> None:
+    """Install forward hooks to harmonize dtypes across quantized and non-quantized layers."""
+    global _DTYPE_HOOKS_INSTALLED
+    if _DTYPE_HOOKS_INSTALLED:
+        return
+    
+    def _harmonize_dtype_hook(module, args, kwargs, output):
+        """Ensure output tensors are in float32/bfloat16, not int8."""
+        if isinstance(output, torch.Tensor):
+            # If output is int8 from a quantized layer, convert to bfloat16
+            if output.dtype == torch.int8:
+                output = output.to(dtype=torch.bfloat16)
+            # If output is uint8, convert to float32
+            elif output.dtype == torch.uint8:
+                output = output.to(dtype=torch.float32)
+        elif isinstance(output, (tuple, list)):
+            # Handle tuple/list outputs
+            output = type(output)(
+                _harmonize_dtype_hook(module, args, kwargs, item) if isinstance(item, torch.Tensor) else item
+                for item in output
+            )
+        return output
+    
+    # Register hooks on all modules
+    hook_count = 0
+    for name, module in model.named_modules():
+        if len(list(module.children())) == 0:  # Only leaf modules
+            module.register_forward_hook(_harmonize_dtype_hook, with_kwargs=True)
+            hook_count += 1
+    
+    _DTYPE_HOOKS_INSTALLED = True
+    logger.info(f"Installed dtype harmonization hooks on {hook_count} modules")
 
 
 def patch_scaled_dot_product_attention() -> None:
@@ -69,20 +107,49 @@ def patch_scaled_dot_product_attention() -> None:
 
 
 def patch_gelu_uint8() -> None:
-    """Promote stray uint8 activations to float before GELU."""
+    """Promote stray uint8/int8 activations to float before GELU."""
     global _GELU_PATCHED, _ORIGINAL_GELU
     if _GELU_PATCHED:
         return
     _ORIGINAL_GELU = torch.nn.functional.gelu
 
     def _patched_gelu(input, approximate='none'):
-        if isinstance(input, torch.Tensor) and input.dtype == torch.uint8:
-            input = input.to(dtype=torch.float32)
+        if isinstance(input, torch.Tensor):
+            # Handle quantized types (uint8, int8, etc.)
+            if input.dtype in (torch.uint8, torch.int8):
+                input = input.to(dtype=torch.float32)
+            # Handle bitsandbytes quantized tensors with quant_state
+            elif hasattr(input, 'quant_state'):
+                # This is a quantized tensor - it should have been dequantized before GELU
+                # Force dequantization by converting to float
+                input = input.to(dtype=torch.float32)
         return _ORIGINAL_GELU(input, approximate=approximate)
 
     torch.nn.functional.gelu = _patched_gelu  # type: ignore[assignment]
     _GELU_PATCHED = True
-    logger.info("Patched GELU to promote uint8 activations to float32")
+    logger.info("Patched GELU to promote uint8/int8/quantized activations to float32")
+
+
+def patch_silu_uint8() -> None:
+    """Promote stray uint8/int8 activations to float before SiLU."""
+    global _SILU_PATCHED, _ORIGINAL_SILU
+    if _SILU_PATCHED:
+        return
+    _ORIGINAL_SILU = torch.nn.functional.silu
+
+    def _patched_silu(input, inplace=False):
+        if isinstance(input, torch.Tensor):
+            # Handle quantized types (uint8, int8, etc.)
+            if input.dtype in (torch.uint8, torch.int8):
+                input = input.to(dtype=torch.float32)
+            # Handle bitsandbytes quantized tensors with quant_state
+            elif hasattr(input, 'quant_state'):
+                input = input.to(dtype=torch.float32)
+        return _ORIGINAL_SILU(input, inplace=inplace)
+
+    torch.nn.functional.silu = _patched_silu  # type: ignore[assignment]
+    _SILU_PATCHED = True
+    logger.info("Patched SiLU to promote uint8/int8/quantized activations to float32")
 
 
 def patch_scatter_dtype() -> None:
@@ -226,6 +293,7 @@ def ensure_model_on_device(model, device: torch.device, skip_quantized_params: b
 
     patch_scaled_dot_product_attention()
     patch_gelu_uint8()
+    patch_silu_uint8()
     patch_scatter_dtype()
     patch_index_copy_dtype()
     patch_dynamic_cache_dtype()
@@ -345,6 +413,19 @@ class HunyuanModelCache:
                 cls._cached_model = None
                 cls._cached_path = None
                 
+                # Remove accelerate hooks if present
+                hook = getattr(model_to_clear, '_cpu_offload_hook', None)
+                if hook is not None:
+                    try:
+                        hook.remove()
+                    except Exception as hook_err:
+                        logger.warning("Failed to remove cpu_offload hook: %s", hook_err)
+                if hasattr(model_to_clear, '_forced_device'):
+                    try:
+                        delattr(model_to_clear, '_forced_device')
+                    except Exception:
+                        pass
+
                 # Now clean up the model
                 if hasattr(model_to_clear, 'cpu'):
                     model_to_clear.cpu()
@@ -413,6 +494,96 @@ class HunyuanImage3Unload:
         return (cleared, float(time.time()))
 
 
+class MemoryTracker:
+    """Track system RAM and GPU VRAM usage during a generation run."""
+
+    def __init__(self, device_index: int = 0):
+        self.process = psutil.Process()
+        self.device_index = device_index if torch.cuda.is_available() else None
+        self.ram_start_bytes = self.process.memory_info().rss
+        self._finished_stats = None
+
+        if self.device_index is not None:
+            torch.cuda.reset_peak_memory_stats(self.device_index)
+            free_bytes, total_bytes = torch.cuda.mem_get_info(self.device_index)
+            self.vram_start_bytes = total_bytes - free_bytes
+            self.vram_total_bytes = total_bytes
+        else:
+            self.vram_start_bytes = None
+            self.vram_total_bytes = None
+
+    @staticmethod
+    def _bytes_to_gb(value: Optional[int]) -> Optional[float]:
+        if value is None:
+            return None
+        return value / 1024**3
+
+    @staticmethod
+    def _extract_peak_ram(mem_info) -> int:
+        candidates = []
+        for attr in ("peak_wset", "peak_rss", "peak_pagefile", "peak_pagefaults", "rss"):
+            value = getattr(mem_info, attr, None)
+            if isinstance(value, (int, float)) and value > 0:
+                candidates.append(int(value))
+        return max(candidates) if candidates else mem_info.rss
+
+    def finish(self) -> dict:
+        if self._finished_stats is not None:
+            return self._finished_stats
+
+        mem_end = self.process.memory_info()
+        ram_end_bytes = mem_end.rss
+        ram_peak_bytes = max(self._extract_peak_ram(mem_end), self.ram_start_bytes, ram_end_bytes)
+
+        vram_end_bytes = None
+        vram_peak_bytes = None
+        if self.device_index is not None:
+            free_bytes, total_bytes = torch.cuda.mem_get_info(self.device_index)
+            vram_end_bytes = total_bytes - free_bytes
+            vram_peak_bytes = torch.cuda.max_memory_reserved(self.device_index)
+
+        stats = {
+            "ram_start_gb": self._bytes_to_gb(self.ram_start_bytes),
+            "ram_end_gb": self._bytes_to_gb(ram_end_bytes),
+            "ram_peak_gb": self._bytes_to_gb(ram_peak_bytes),
+            "vram_start_gb": self._bytes_to_gb(self.vram_start_bytes),
+            "vram_end_gb": self._bytes_to_gb(vram_end_bytes),
+            "vram_peak_gb": self._bytes_to_gb(vram_peak_bytes),
+            "vram_total_gb": self._bytes_to_gb(self.vram_total_bytes),
+        }
+
+        self._finished_stats = stats
+        return stats
+
+    def start_snapshot(self) -> dict:
+        return {
+            "ram_start_gb": self._bytes_to_gb(self.ram_start_bytes),
+            "vram_start_gb": self._bytes_to_gb(self.vram_start_bytes),
+            "vram_total_gb": self._bytes_to_gb(self.vram_total_bytes),
+        }
+
+
+def format_memory_stats(stats: dict) -> str:
+    """Create a short human-readable summary of memory usage."""
+    parts = []
+    peak_vram = stats.get("vram_peak_gb")
+    if peak_vram is not None:
+        parts.append(f"Peak VRAM {peak_vram:.1f}GB")
+
+    vram_start = stats.get("vram_start_gb")
+    vram_end = stats.get("vram_end_gb")
+    if vram_start is not None and vram_end is not None:
+        parts.append(f"VRAM Start {vram_start:.1f}GB → End {vram_end:.1f}GB")
+
+    peak_ram = stats.get("ram_peak_gb")
+    if peak_ram is not None:
+        parts.append(f"Peak RAM {peak_ram:.1f}GB")
+
+    if not parts:
+        return ""
+    return " | ".join(parts)
+
+
 def patch_hunyuan_generate_image(model):
     """
     Monkey-patch the generate_image method on the model instance to correctly handle

requirements.txt

@@ -6,3 +6,4 @@ safetensors>=0.4.5
 pillow>=11.0.0
 numpy>=1.26.0
 requests>=2.32.0
+psutil>=5.9.0