diff --git a/.clinerules b/.clinerules deleted file mode 100644 index f04aadf..0000000 --- a/.clinerules +++ /dev/null @@ -1,256 +0,0 @@ -# ComfyUI-MultiGPU Development Rules - -## Project Context -This is ComfyUI-MultiGPU: a production-grade multi-device AI inference platform that transforms ComfyUI from single-GPU to universal multi-device support. The project enables previously impossible AI workflows across diverse hardware configurations. - -**Current Version**: v2.5.0 Release Candidate -**Status**: PRODUCTION READY -**Stability**: 9/10 - Verified working in production -**Community**: 300+ commits, 90+ resolved issues, active ecosystem - -## Memory Bank System -**CRITICAL**: Always read ALL files in the `memory-bank/` folder at the start of every session. The Memory Bank contains complete project context: - -### Core Documentation (Read These First) -1. `memory-bank/projectbrief.md` - Project identity, mission, evolution timeline -2. `memory-bank/productContext.md` - Problem space, user goals, success metrics -3. `memory-bank/activeContext.md` - Current work focus and priorities (UPDATED 2025-09-30) -4. `memory-bank/progress.md` - Production status, roadmap, lessons learned (UPDATED 2025-09-30) - -### Technical Deep Dive -5. `memory-bank/systemPatterns.md` - Architecture patterns and design decisions (UPDATED 2025-09-30) -6. `memory-bank/techContext.md` - Technology stack and development environment -7. `memory-bank/performance-benchmarks.md` - Quantified performance across hardware configurations -8. `memory-bank/comfyui-lineage.md` - Integration analysis with ComfyUI core - -## Development Philosophy -- **Extend, Don't Replace**: Build upon ComfyUI's existing patterns -- **Fail Loudly**: Immediate detection of API changes prevents silent failures -- **User Agency**: Let users specify device placement explicitly -- **Production Quality**: Stability and reliability over experimental features -- **Community First**: Solutions should benefit the entire ComfyUI ecosystem -- **Clean Code**: Remove debug artifacts, comprehensive production logging only - -## Key Technical Patterns -- **City96's Dynamic Class Override**: Elegant inheritance pattern for node creation (50 lines vs 400+) -- **Load-Patch-Distribute Pipeline**: Quality-preserving LoRA application workflow -- **ComfyCore Alignment**: Work WITH existing ComfyUI patterns, not against them -- **Multi-Device Native**: Treat all devices as equal citizens -- **Selective Unload**: Per-model granular control over memory management - -## Production Status (v2.5.0) - -### Core Features ✅ -- **DisTorch2 Distributed Loading**: Universal SafeTensor support with CLIP head preservation -- **Selective Unload System**: Verified working - keeps models with `keep_loaded=True`, ejects others -- **Multi-Device VRAM Management**: Clears allocator caches across all devices -- **Manager Parity**: Mirrors ComfyUI-Manager "Free model and node cache" behavior -- **Universal Device Support**: CUDA, CPU, MPS, XPU, NPU, MLU, DirectML, CoreX - -### Recent Achievements (2025-09-30) -- **Code Refactoring** (-219 lines total): - - DisTorch2 allocation consolidation (-179 lines): Unified UNET and CLIP allocation functions - - Production cleanup (-40 lines): Removed diagnostic instrumentation wrapper -- **Verified Working**: Selective unload tested in production with comprehensive logging -- **Clean Architecture**: Single responsibility modules, clear dependency direction - -### Performance Validation -- **NVLink**: 5-7% slowdown (near-native) -- **PCIe 4.0 x16**: 40-50% slowdown (excellent) -- **PCIe 3.0 x16**: 70-80% slowdown (good) -- **PCIe 4.0 x8**: 80-100% slowdown (acceptable) -- **PCIe 3.0 x8**: 150-200% slowdown (workable) -- **PCIe 3.0 x4**: 300-400% slowdown (last resort) - -### Ecosystem Integration -- 10+ custom node integrations with automatic detection -- Dynamic node creation for compatible loaders -- Fail-loudly compatibility with ComfyCore API - -## Module Architecture Rules - -### Module Boundary Principles -- **Single Responsibility**: Each module should have ONE clear purpose -- **Dependency Direction**: Dependencies flow UPWARD only - violations create circular imports -- **Import Hierarchy**: Base modules NEVER import from Feature/UI modules - -### Module Hierarchy (Dependency Order) -1. **`device_utils.py`** - BASE LAYER - - Device detection and enumeration - - VRAM cache management (`soft_empty_cache_multigpu`) - - Pure hardware abstraction - NO model tracking - -2. **`model_management_mgpu.py`** - CORE LAYER - - Model lifecycle tracking - - Memory logging infrastructure - - Cleanup orchestration (`force_full_system_cleanup`, `trigger_executor_cache_reset`) - - Patched `mm.unload_all_models` (selective ejection) - -3. **`distorch_2.py`, `distorch.py`** - FEATURE LAYER - - DisTorch distribution algorithms - - Allocation analysis and device assignment - - Per-model flag setting (`_mgpu_unload_distorch_model`) - - Imports from CORE/BASE only - -4. **`nodes.py`, `checkpoint_multigpu.py`** - UI LAYER - - Device-aware user interfaces - - Node implementations and definitions - - Imports from any lower level - -5. **`__init__.py`** - ASSEMBLY LAYER - - Final integration and patch registration - - Node mapping and registration - - Imports from all lower levels - -### Import Flow Architecture -``` - __init__.py ← Assembly - ↑ - UI Layer ← nodes.py, checkpoint_multigpu.py - ↑ - Feature Layer ← distorch_2.py, distorch.py - ↑ - Core Layer ← model_management_mgpu.py - ↑ - Base Layer ← device_utils.py -``` - -### Mandatory Architecture Checks -**BEFORE adding ANY import statement:** -1. **Check Direction**: Does this create upward dependency? (FORBIDDEN) -2. **Check Purpose**: Does the function belong in this module per Single Responsibility? -3. **Check Cycles**: Run `python -c "import sys; sys.path.append('.'); import "` to detect circular imports - -### Function Placement Rules -- **device_utils.py**: ONLY device detection, VRAM cache management -- **model_management_mgpu.py**: Model tracking, memory logging, cleanup utilities -- **Feature modules**: Import from CORE/BASE only, never each other -- **UI modules**: Import from any lower level, implement user interfaces only - -### Violation Detection -If import fails with "circular import" or "cannot import name": -1. STOP immediately - do not work around -2. Identify which module boundary was violated -3. Move misplaced function to correct architectural layer -4. Update ALL imports consistently - -## Memory Management System (Verified Working) - -### Selective Unload Pipeline -**Load Phase**: -```python -# DisTorch2 wrapper sets per-model flag based on keep_loaded parameter -if hasattr(out[0], 'model') and hasattr(out[0].model, '_mgpu_keep_loaded'): - keep_loaded = out[0].model._mgpu_keep_loaded - out[0].model._mgpu_unload_distorch_model = (not keep_loaded) -``` - -**Unload Phase** (patched `mm.unload_all_models`): -```python -# Categorize models by flag -models_to_unload = [flagged models] -kept_models = [unflagged models] - -if kept_models: - # Selective: eject flagged, retain others with GC anchors - for lm in models_to_unload: - lm.model_unload(unpatch_weights=True) - mm.current_loaded_models = kept_models -else: - # Standard cleanup when no models to keep - _mgpu_original_unload_all_models() -``` - -**Verified Working** (Production Logs 2025-09-30): -``` -[CATEGORIZE_SUMMARY] kept_models: 2, models_to_unload: 1, total: 3 -[SELECTIVE_UNLOAD] Proceeding with selective unload: retaining 2, unloading 1 -[REMAINING_MODEL] 0: AutoencodingEngine -[REMAINING_MODEL] 1: FluxClipModel_ -``` - -### Manager Parity -`force_full_system_cleanup()` mirrors ComfyUI-Manager "Free model and node cache": -- Sets `unload_models=True`, `free_memory=True` on PromptQueue -- Triggers patched `mm.unload_all_models` for selective ejection -- Triggers `PromptExecutor.reset()` for CPU memory management - -## Code Quality Standards - -### Production Requirements -- **No Debug Cruft**: Remove all diagnostic-only code before release -- **Comprehensive Logging**: Production-grade telemetry at major operations -- **Clean Modules**: Single responsibility, clear boundaries -- **Fail Loudly**: Surface API changes immediately, no defensive masking - -### Logging Conventions -```python -# Model Management logs -logger.mgpu_mm_log("[OPERATION] Description with context") - -# Memory state logging -multigpu_memory_log("identifier", "tag") - -# Debug logging (use sparingly) -logger.debug("[Component] Detailed diagnostic information") -``` - -### Code Style -- Self-documenting code over excessive comments -- Clear function/variable names conveying intent -- Minimal comments for non-obvious constraints only -- Structured logging for production debugging - -## Development Workflow - -### Before Making Changes -1. Read relevant Memory Bank files -2. Understand module architecture and dependencies -3. Check if change violates architectural boundaries -4. Consider impact on existing patterns - -### When Adding Features -1. Determine correct module placement (BASE/CORE/FEATURE/UI) -2. Verify no circular dependencies created -3. Add comprehensive logging at key operations -4. Test with production workflows -5. Update Memory Bank documentation - -### When Refactoring -1. Eliminate code duplication (DRY principle) -2. Remove debug artifacts and diagnostic code -3. Maintain or improve architectural clarity -4. Verify no functionality regressions -5. Document pattern changes in systemPatterns.md - -## Testing Philosophy - -### Manual Validation -- Test across hardware configurations (NVLink, PCIe variants, CPU) -- Verify selective unload with keep_loaded combinations -- Check memory usage patterns (VRAM + CPU) -- Validate quality parity with single-GPU baselines - -### Community Testing -- Active users provide hardware configuration validation -- Integration testing with custom node ecosystem -- Performance feedback across diverse setups - -## Next Steps (v2.5.0 Release) - -### Immediate -- [ ] Final testing pass across hardware configurations -- [ ] GitHub release notes and changelog -- [ ] Community announcement - -### Short-term -- [ ] Issue triage and community feedback integration -- [ ] New model format support (Mochi, community requests) -- [ ] Documentation refresh and tutorials - -### Long-term -- [ ] Model parallelism research -- [ ] Streaming inference for video -- [ ] Multi-node orchestration - -When working on this project, always reference the Memory Bank for context and maintain the established patterns and philosophy. The codebase is production-ready - focus on stability, community needs, and quality over experimental features. diff --git a/RELEASE_NOTES_v2.5.0.md b/RELEASE_NOTES_v2.5.0.md deleted file mode 100644 index d7a5a55..0000000 --- a/RELEASE_NOTES_v2.5.0.md +++ /dev/null @@ -1,350 +0,0 @@ -# ComfyUI-MultiGPU v2.5.0 Release Notes - -## Overview - -Version 2.5.0 marks a significant maturity milestone for ComfyUI-MultiGPU, delivering **production-grade stability** through comprehensive code refactoring, verified selective model unloading, and enhanced architectural clarity. This release removes 219 lines of code while adding powerful new capabilities. - -**Status**: Production Ready (9/10 Stability Rating) -**Total Changes**: +8,094 additions / -1,645 deletions across 23 files -**Code Quality**: Significant improvement through refactoring and cleanup - ---- - -## 🎯 Major Features - -### ✅ Selective Model Unloading (Verified Working) - -The flagship feature of v2.5.0 enables **granular control over model memory management** through a per-model `keep_loaded` parameter. - -**What It Does**: -- Keep specific models loaded in VRAM while unloading others -- Prevents expensive reload cycles for frequently-used models -- Reduces workflow iteration time by 50-80% in multi-model scenarios -- Works with **any** DisTorch2-enabled loader - -**How It Works**: -```python -# Example: Keep VAE and CLIP loaded, allow UNet to be unloaded -UNet Loader (DisTorch2): keep_loaded=False # Can be unloaded -CLIP Loader (DisTorch2): keep_loaded=True # Stays in VRAM -VAE Loader (DisTorch2): keep_loaded=True # Stays in VRAM -``` - -**Verification**: Confirmed working in production with comprehensive logging: -``` -[CATEGORIZE_SUMMARY] kept_models: 2, models_to_unload: 1, total: 3 -[SELECTIVE_UNLOAD] Proceeding with selective unload: retaining 2, unloading 1 -[REMAINING_MODEL] 0: AutoencodingEngine -[REMAINING_MODEL] 1: FluxClipModel_ -``` - -**Technical Implementation**: -- Per-model `_mgpu_unload_distorch_model` flag system -- Patched `mm.unload_all_models` with selective categorization -- GC anchor protection prevents premature garbage collection -- Manager parity with ComfyUI-Manager's "Free model and node cache" - ---- - -### 🏗️ Major Code Refactoring (-219 Lines) - -Significant architectural improvements through consolidation and cleanup. - -#### DisTorch2 Allocation Consolidation (-179 lines) - -**Before**: Separate functions with 85% code duplication -- `analyze_safetensor_loading()` for standard models -- `analyze_safetensor_loading_clip()` for CLIP models - -**After**: Single unified function with CLIP-specific handling -- `analyze_safetensor_loading(model_patcher, allocations, is_clip=False)` -- Helper function `_extract_clip_head_blocks()` for CLIP head preservation -- ~10% performance improvement over DisTorch V1 - -**Benefits**: -- Single source of truth for allocation logic -- Easier to maintain and extend -- Eliminates duplicate bug fixes -- Clearer code flow - -#### Production Cleanup (-40 lines) - -Removed all diagnostic and debug artifacts: -- Deleted `_mgpu_instrumented_soft_empty_cache()` wrapper -- Removed temporary diagnostic logging -- Clean separation: `device_utils.py` = functional, `model_management_mgpu.py` = lifecycle -- Only production-grade logging remains - ---- - -### 📚 Comprehensive Documentation System - -**Memory Bank** (7,739 new lines): -- `projectbrief.md` - Project identity and evolution timeline -- `productContext.md` - Problem space and user goals -- `activeContext.md` - Current work focus and priorities -- `progress.md` - Production status and roadmap -- `systemPatterns.md` - Architecture patterns and design decisions -- `techContext.md` - Technology stack and environment -- `performance-benchmarks.md` - Quantified performance data -- `comfyui-lineage.md` - ComfyUI core integration analysis - -**Code Quality**: -- All functions now PEP 257 compliant with single-line docstrings -- Comprehensive inline documentation -- Clear module boundaries and responsibilities - ---- - -### 🔧 Architecture Improvements - -#### Clean Module Boundaries - -**New File**: `wrappers.py` (+520 lines) -- Consolidated all node wrapper generation functions -- Clear separation from initialization logic -- Single location for override patterns - -**Improved Separation**: -- `device_utils.py` - Hardware detection and VRAM management -- `model_management_mgpu.py` - Model lifecycle tracking and cleanup -- `distorch_2.py` - Distribution algorithms -- `wrappers.py` - Node creation patterns -- `__init__.py` - Assembly and registration - -#### Single Responsibility Principle - -Each module now has ONE clear purpose: -- No circular dependencies -- Clear import hierarchy (Base → Core → Feature → UI → Assembly) -- Easier testing and maintenance - ---- - -## 🚀 Performance Validation - -### Hardware Performance Tiers (Verified) - -| Connection Type | Slowdown | Rating | Use Case | -|----------------|----------|---------|----------| -| **NVLink** | 5-7% | Excellent | Professional multi-GPU systems | -| **PCIe 4.0 x16** | 40-50% | Excellent | Modern consumer builds | -| **PCIe 3.0 x16** | 70-80% | Good | Standard desktop systems | -| **PCIe 4.0 x8** | 80-100% | Acceptable | Budget/compact builds | -| **PCIe 3.0 x8** | 150-200% | Workable | Older systems, still functional | -| **PCIe 3.0 x4** | 300-400% | Last Resort | Better than OOM errors | - -### Model Validation ✅ - -Tested and verified with: -- **FLUX** (1.dev, schnell, GGUF variants) -- **WAN Video** (1.3B, 2.0, 2.2) -- **HunyuanVideo** (text-to-video) -- **QWEN VL** (image understanding) -- **Florence2** (vision tasks) -- **SDXL, SD1.5** (classic models) - -**Quality Guarantee**: Bit-exact parity with single-GPU inference (zero precision loss) - ---- - -## 🔌 Integration Support - -### Verified Custom Node Integrations - -- ✅ **ComfyUI-GGUF** - Quantized model support -- ✅ **ComfyUI-WanVideoWrapper** - Video generation -- ✅ **ComfyUI-Florence2** - Vision tasks -- ✅ **ComfyUI-HunyuanVideoWrapper** - HunyuanVideo support -- ✅ **ComfyUI-LTXVideo** - LTXV models -- ✅ **ComfyUI-MMAudio** - Audio synthesis -- ✅ **PuLID_ComfyUI** - Identity preservation -- ✅ **ComfyUI_bitsandbytes_NF4** - NF4 quantization -- ✅ **x-flux-comfyui** - Flux ControlNet - -**Total**: 10+ integrations with automatic MultiGPU node generation - ---- - -## 🛠️ Technical Details - -### DisTorch2 Allocation Modes - -Three flexible ways to specify memory distribution: - -1. **Bytes Mode** (Explicit) - ``` - cuda:0,6gb;cuda:1,4gb;cpu,* - ``` - Direct byte allocation with wildcard support - -2. **Ratio Mode** (Percentage) - ``` - cuda:0,60%;cuda:1,30%;cpu,10% - ``` - Proportional model splitting - -3. **Fraction Mode** (Automatic) - ``` - compute_device=cuda:0, virtual_vram_gb=4.0, donor_device=cpu - ``` - Automatic calculation based on VRAM constraints - -### CLIP Head Preservation - -DisTorch2 now intelligently handles CLIP models: -- Automatically detects head layers (embeddings, positional encodings) -- Keeps heads on compute device for optimal performance -- Distributes remaining layers across donor devices -- Zero configuration required - -### Universal Device Support - -Supports all PyTorch accelerator types: -- **CUDA** (NVIDIA GPUs) -- **XPU** (Intel GPUs) -- **NPU** (Huawei Ascend) -- **MLU** (Cambricon) -- **MPS** (Apple Metal) -- **DirectML** (Windows DirectML) -- **CoreX** (Specialized accelerators) -- **CPU** (Always available) - ---- - -## 📊 What Users Are Saying - -> "Previously impossible workflows now run reliably on my 2x3090 setup" - -> "The selective unload feature saves me hours of iteration time" - -> "Finally can use my 8GB card alongside my 24GB card effectively" - ---- - -## 🔍 Under the Hood - -### Code Quality Metrics - -- **Lines Removed**: 219 (eliminating redundancy and debug code) -- **Documentation Added**: 7,739 lines (memory bank system) -- **Functions Documented**: 67 (100% PEP 257 compliance) -- **Module Refactoring**: 5 major files reorganized -- **Test Coverage**: Validated across 6 hardware configurations - -### Logging Infrastructure - -Production-grade telemetry at every major operation: -- Memory snapshots with timestamp alignment -- Device-specific cache management tracking -- Model lifecycle event logging -- Selective unload categorization details - -### Fail-Loudly Philosophy - -Rather than masking issues, v2.5.0 surfaces them immediately: -- API changes detected instantly -- Clear error messages with context -- Comprehensive diagnostic logging -- Community can identify and report issues quickly - ---- - -## 🚦 Migration from v2.4.x - -### Breaking Changes - -**None** - v2.5.0 is fully backward compatible. - -### New Features Available - -To use selective unloading, add `keep_loaded` parameter to DisTorch2 loaders: -```python -# Old (still works) -UNETLoader (DisTorch2) - -# New (recommended) -UNETLoader (DisTorch2): keep_loaded=True # Stays in VRAM -``` - -### Recommended Actions - -1. **Update workflows** to use selective unload where beneficial -2. **Review allocation strategies** with new CLIP head preservation -3. **Enable logging** during testing to verify behavior -4. **Report issues** on GitHub with comprehensive logs - ---- - -## 🎓 Learning Resources - -### Example Workflows - -20+ JSON examples in `/examples`: -- `distorch2/` - DisTorch2 allocation patterns -- `multiGPU/` - Standard MultiGPU workflows -- `gguf/` - Quantized model examples -- Model-specific examples (Florence2, HunyuanVideo, WanVideo, etc.) - -### Documentation - -- **README.md** - Architecture overview and quick start -- **Memory Bank** - Comprehensive technical documentation -- **Performance Benchmarks** - Hardware selection guide -- **.clinerules** - Development patterns and practices - ---- - -## 🙏 Acknowledgments - -### Community Contributions - -- **City96** - Dynamic class override pattern (foundation of architecture) -- **ComfyUI Core Team** - Extensible architecture enabling multi-device support -- **Custom Node Developers** - Integration partnerships and testing -- **Community Testers** - Hardware validation across diverse configurations - -### Special Thanks - -To the 300+ commits and 90+ resolved issues that shaped this release. - ---- - -## 📅 What's Next - -### Immediate (v2.5.1) -- Issue triage and community feedback -- Minor bug fixes -- Integration expansion - -### Short-term (v2.6.0) -- Allocation string generator with validation -- Hardware profiler tools -- Enhanced documentation and tutorials - -### Long-term (v3.0.0) -- Model parallelism experiments -- Streaming inference for video -- Multi-node orchestration -- Pipeline parallelism - ---- - -## 📞 Support & Community - -- **GitHub Issues**: Bug reports and feature requests -- **Discussions**: Architecture questions and optimization tips -- **Pull Requests**: Contributions welcome! - ---- - -## ⚖️ License - -MIT License - See LICENSE file for details - ---- - -**Version**: 2.5.0 -**Release Date**: September 30, 2025 -**Stability Rating**: 9/10 (Production Ready) -**Recommended**: Yes - Significant quality improvements over 2.4.x diff --git a/memory-bank/activeContext.md b/memory-bank/activeContext.md deleted file mode 100644 index a6388c3..0000000 --- a/memory-bank/activeContext.md +++ /dev/null @@ -1,194 +0,0 @@ -# Active Context: Production Ready v2.5.0 (Updated 2025-09-30) - -## Current Project State - -**Status**: PRODUCTION READY - v2.5.0 Release Candidate -**Stability**: 300+ commits, 90+ resolved issues, active community -**Performance**: Validated across 6 hardware configurations -**Code Quality**: Clean, refactored, comprehensive logging - -## Recent Session Achievements (2025-09-30) - -### ✅ DisTorch2 Allocation Refactoring (-179 lines) -**Problem**: 85% code duplication between UNET and CLIP allocation functions -**Solution**: Consolidated into unified `analyze_safetensor_loading(model_patcher, allocations, is_clip=False)` -- CLIP-specific head preservation via helper function `_extract_clip_head_blocks()` -- Single source of truth for allocation logic -- Easier maintenance and debugging -- **Verified working**: Logs show "Preserving 2 head layer(s) (72.49 MB)" - -### ✅ Production Cleanup (-40 lines) -**Removed**: Diagnostic instrumentation from model_management_mgpu.py -- Deleted `_mgpu_instrumented_soft_empty_cache()` wrapper (debug artifact) -- Retained production telemetry and functional patches -- Clear separation: device_utils.py = functional, model_management = lifecycle - -### ✅ Selective Unload VERIFIED WORKING -**Test Results** (from production logs): -``` -[CATEGORIZE_SUMMARY] kept_models: 2, models_to_unload: 1, total: 3 -[SELECTIVE_UNLOAD] Proceeding with selective unload: retaining 2, unloading 1 -[UNLOAD_EXECUTE] Unloading model: Flux -[REMAINING_MODEL] 0: AutoencodingEngine -[REMAINING_MODEL] 1: FluxClipModel_ -``` - -**Key Components Working**: -- Per-model `_mgpu_unload_distorch_model` flag setting (working) -- Selective unload logic in patched `mm.unload_all_models` (working) -- GC anchor system preventing premature collection (working) -- Multi-device cache clearing (working) - -## Architecture Status - -### Core Files - Production Ready -1. **__init__.py** (284 lines) - Clean initialization and node registration -2. **device_utils.py** (420 lines) - Universal device support + comprehensive memory patch -3. **distorch_2.py** (refactored) - Unified allocation with CLIP support -4. **model_management_mgpu.py** (cleaned) - Selective unload with diagnostics -5. **checkpoint_multigpu.py** (252 lines) - Advanced checkpoint loaders -6. **wrappers.py** - Dynamic node creation via City96 pattern - -### Memory Management Pipeline (Verified Working) - -**Load Phase**: -1. DisTorch2 wrapper detects `keep_loaded` parameter -2. Sets `_mgpu_unload_distorch_model = (not keep_loaded)` on ModelPatcher -3. Stores allocation in safetensor_allocation_store - -**Execution Phase**: -4. Models load with distributed blocks across devices -5. CLIP head preservation works (verified in logs) -6. Quality-preserving LoRA application on compute device - -**Unload Phase** (End of workflow): -7. `force_full_system_cleanup()` sets `unload_models=True`, `free_memory=True` -8. Patched `mm.unload_all_models()` categorizes models: - - `_mgpu_unload_distorch_model=True` → models_to_unload - - `_mgpu_unload_distorch_model=False` → kept_models (with GC anchors) -9. Selectively unloads flagged models -10. Rebuilds `mm.current_loaded_models` with kept models only -11. Multi-device cache clearing via `soft_empty_cache_multigpu()` - -## Current Development Priorities - -### 1) v2.5.0 Release Preparation (IMMEDIATE) -- [x] Refactor DisTorch2 allocation functions -- [x] Remove diagnostic code -- [x] Verify selective unload working -- [ ] Update memory bank documentation -- [ ] Final testing pass -- [ ] GitHub release notes - -### 2) Ecosystem Expansion (HIGH PRIORITY) -Active Integrations: -- ✅ ComfyUI-GGUF: DisTorch-enabled GGUF nodes -- ✅ WanVideoWrapper: MultiGPU video generation -- ✅ Florence2: Vision model support -- ✅ HunyuanVideoWrapper: Native VAE support -- ✅ LTXVideo: Video generation -- ✅ MMAudio: Audio synthesis -- ✅ PuLID: Identity preservation - -Next Targets: -- Mochi video models -- Community-requested integrations - -### 3) Documentation & UX (MEDIUM PRIORITY) -- 20+ example JSON workflows -- Clear error messages and guidance -- Hardware-specific recommendations -- Configuration validation - -### 4) Advanced Features (LOW PRIORITY - Research) -- Model parallelism experiments -- Memory compression techniques -- Quality metrics and parity validation -- Pipeline parallelism - -## Technical Design Principles - -### Memory Management Philosophy -1. **Conservative by default** - Explicit user control -2. **Quality preservation** - Patch LoRAs before distributing -3. **Transparency** - Comprehensive structured logging -4. **Fail-loudly** - Immediate detection of API changes - -### Integration Strategy -1. **Inheritance-based override** (City96 pattern) -2. **Minimal patch surface**: - - `mm.get_torch_device` / `mm.text_encoder_device` - Device selection - - `mm.soft_empty_cache` - Multi-device cache + CPU reset - - `mm.unload_all_models` - Selective ejection -3. **Single source of truth** - device_utils.py for device management - -### Hardware Support Tiers -- **Tier 1**: CUDA (primary validation) -- **Tier 2**: CPU, MPS (secondary validation) -- **Tier 3**: XPU, NPU, MLU, DirectML, CoreX (community validation) - -## Performance Characteristics (Validated) - -### Hardware Configurations -1. **NVLink (RTX 3090 x2)**: 5-7% slowdown vs native -2. **PCIe 4.0 x16**: 40-50% slowdown (excellent) -3. **PCIe 3.0 x16**: 70-80% slowdown (good) -4. **PCIe 4.0 x8**: 80-100% slowdown (acceptable) -5. **PCIe 3.0 x8**: 150-200% slowdown (workable) -6. **PCIe 3.0 x4**: 300-400% slowdown (last resort) - -### Model Validation -- ✅ FLUX (1.dev, schnell, GGUF variants) -- ✅ WAN Video (1.3B, 2.0, 2.2) -- ✅ QWEN VL (image understanding) -- ✅ HunyuanVideo (text-to-video) -- ✅ Florence2 (vision tasks) - -## Known Limitations & Workarounds - -1. **DirectML Performance**: Slower than native CUDA, but functional -2. **CPU Offload Overhead**: PCIe bandwidth bottleneck in extreme offload scenarios -3. **Quality**: Maintains bit-exact parity with single-GPU (validated) -4. **Memory Pressure**: Adaptive thresholds prevent OOM, may trigger premature unloads - -## Next Steps - -### Immediate (This Week) -- [ ] Commit memory bank updates -- [ ] Archive resolved issue docs -- [ ] Final v2.5.0 testing -- [ ] GitHub release with changelog - -### Short-term (2-4 Weeks) -- [ ] Triage GitHub issues -- [ ] Community feedback integration -- [ ] Performance dashboard updates - -### Medium-term (2-3 Months) -- [ ] New model format support -- [ ] Tutorial series refresh -- [ ] Quality measurement automation - -### Long-term (6-12 Months) -- [ ] Model parallelism research -- [ ] Streaming inference for video -- [ ] Multi-node orchestration - -## Development Environment - -- **IDE**: VSCode with Python language support -- **Version Control**: Git with conventional commits -- **Testing**: Manual validation + community testing -- **Primary Hardware**: Multi-GPU configurations (CUDA focus) -- **Limitation**: Limited access to cutting-edge GPUs (RTX 5090, etc.) - -## Summary - -The project has reached production maturity with v2.5.0. Key achievements: -- Selective unload working correctly (verified in logs) -- Clean refactored codebase (-219 lines of cruft) -- Comprehensive logging for production debugging -- Universal device support -- Quality-preserving distributed inference - -The architecture is stable, performant, and ready for release. diff --git a/memory-bank/code-references.md b/memory-bank/code-references.md deleted file mode 100644 index ba1b8ab..0000000 --- a/memory-bank/code-references.md +++ /dev/null @@ -1,231 +0,0 @@ -# Code References (Definitive): ComfyUI Manager “Free model and node cache” - -Purpose -- Provide an end-to-end, fully verified lineage of the ComfyUI Manager “Free model and node cache” button through to the exact consumption of flags in ComfyUI core, with exact file paths and code excerpts captured from the current snapshot in this workspace. -- Document MultiGPU patch integration points that participate in the free/unload flow, including selective unload behavior and current caveats. - -End‑to‑End Flow (Current Snapshot) -1) UI Button (Manager) → 2) JS helper free_models(...) → 3) POST /free (Comfy core) → 4) main.py prompt_worker thread polls flags and performs: - - unload_models: comfy.model_management.unload_all_models() - - free_memory: PromptExecutor.reset() - - Additionally triggers GC and comfy.model_management.soft_empty_cache() - -A) Frontend UI trigger (ComfyUI Manager) -- File: ../ComfyUI-Manager/js/comfyui-manager.js -- Location: app.registerExtension({ name: "Comfy.ManagerMenu", ... }) → setup() → ComfyButtonGroup -```js -new(await import("../../scripts/ui/components/button.js")).ComfyButton({ - icon: "vacuum-outline", - action: () => { - free_models(); - }, - tooltip: "Unload Models" -}).element, -new(await import("../../scripts/ui/components/button.js")).ComfyButton({ - icon: "vacuum", - action: () => { - free_models(true); - }, - tooltip: "Free model and node cache" -}).element, -``` -Semantics: -- “Unload Models” → free_models() (models only) -- “Free model and node cache” → free_models(true) (models + execution cache) - -B) Frontend request construction (ComfyUI Manager) -- File: ../ComfyUI-Manager/js/common.js -- Function: export async function free_models(free_execution_cache) -```js -export async function free_models(free_execution_cache) { - try { - let mode = ""; - if (free_execution_cache) { - mode = '{"unload_models": true, "free_memory": true}'; - } else { - mode = '{"unload_models": true}'; - } - - console.log(`[ManagerFreePath] POST /free payload: ${mode}`); - let res = await api.fetchApi(`/free`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: mode - }); - console.log(`[ManagerFreePath] /free status: ${res.status}`); - - if (res.status == 200) { - if (free_execution_cache) { - showToast("'Models' and 'Execution Cache' have been cleared.", 3000); - } else { - showToast("Models' have been unloaded.", 3000); - } - } else { - showToast('Unloading of models failed. Installed ComfyUI may be an outdated version.', 5000); - } - } catch (error) { - console.error('[ManagerFreePath] /free error:', error); - showToast('An error occurred while trying to unload models.', 5000); - } -} -``` -Semantics: -- free_models(true) → POST /free with {"unload_models": true, "free_memory": true} -- free_models() → POST /free with {"unload_models": true} - -C) Core server endpoint (flags are set on the queue) -- File: ../../server.py -- Route: @routes.post("/free") -```py -@routes.post("/free") -async def post_free(request): - json_data = await request.json() - unload_models = json_data.get("unload_models", False) - free_memory = json_data.get("free_memory", False) - if unload_models: - self.prompt_queue.set_flag("unload_models", unload_models) - if free_memory: - self.prompt_queue.set_flag("free_memory", free_memory) - return web.Response(status=200) -``` -Semantics: -- The HTTP endpoint itself does not unload/reset; instead it sets flags on PromptServer.prompt_queue for the background worker to consume. - -D) Flag consumption and execution (definitive mechanism) -- File: ../../main.py -- Function: prompt_worker(q, server_instance) -- Excerpt (poll and handle flags, then clean up): -```py - flags = q.get_flags() - free_memory = flags.get("free_memory", False) - - if flags.get("unload_models", free_memory): - comfy.model_management.unload_all_models() - need_gc = True - last_gc_collect = 0 - - if free_memory: - e.reset() - need_gc = True - last_gc_collect = 0 - - if need_gc: - current_time = time.perf_counter() - if (current_time - last_gc_collect) > gc_collect_interval: - gc.collect() - comfy.model_management.soft_empty_cache() - last_gc_collect = current_time - need_gc = False - hook_breaker_ac10a0.restore_functions() -``` -Context: -- e is a PromptExecutor (created earlier in prompt_worker): `e = execution.PromptExecutor(server_instance, ...)` -- The worker thread is started in start_comfyui(): -```py -threading.Thread(target=prompt_worker, daemon=True, args=(prompt_server.prompt_queue, prompt_server,)).start() -``` - -Interpretation (What the Manager button actually does) -- “Free model and node cache” sets unload_models: true and free_memory: true via POST /free. -- The background prompt_worker then: - - Calls comfy.model_management.unload_all_models() - - Calls e.reset() on the PromptExecutor to drop execution caches - - Performs gc.collect() and comfy.model_management.soft_empty_cache() -- This matches the “benchmark button” behavior required for CPU memory reclamation (models fully unloaded + executor reset + allocator/cache cleanup). - -Implications for MultiGPU P1 (force_full_system_cleanup) -- To 100% replicate the benchmark button behavior from within MultiGPU code paths: - - Call comfy.model_management.unload_all_models() - - Trigger PromptExecutor.reset() on the active executor - - Follow up with gc.collect() and comfy.model_management.soft_empty_cache() -- Or, trigger the core behavior indirectly by POST /free with both flags set, relying on ComfyUI’s running prompt worker. - -Verification Status -- All file paths and snippets above were extracted from this workspace: - - Manager JS files under ../ComfyUI-Manager/js/ - - ComfyUI server and main under ../../server.py and ../../main.py -- Consumption site conclusively identified in ../../main.py prompt_worker via q.get_flags → unload_all_models + PromptExecutor.reset - ---- - -MultiGPU Integration Points (This Repository) - -Overview -- In addition to the core /free flow, MultiGPU patches (in this repository) alter both the unload and soft-empty behaviors to enable selective ejection of DisTorch-managed models and multi-device cache clearing. - -1) Per-model transient flag (DisTorch2 nodes) -- File: memory-bank reference → implemented in code at: ./distorch_2.py -- Where: - - In DisTorch2 wrappers (UNET/CLIP/VAE) inside `override(...)`, after calling the original node: - - `out[0].model._mgpu_unload_distorch_model = (not keep_loaded)` -- Purpose: - - Mark models for ejection only when the user disables “keep_loaded”. - - This supplants the previously planned global sentinel; the implemented design is purely per-model. - -2) Selective unloading (patched unload_all_models) -- File: ./model_management_mgpu.py -- Patch site notes: - - At import time, we patch `mm.unload_all_models` with `_mgpu_patched_unload_all_models`. - - Behavior: - - Iterate `mm.current_loaded_models` into: - - `models_to_unload`: those with `_mgpu_unload_distorch_model == True` - - `kept_models`: the rest - - If any are flagged, unload only `models_to_unload` and rebuild `mm.current_loaded_models = kept_models`. - - If none are flagged (all kept), current code delegates to original `unload_all_models()` (known caveat; see below). -- Known caveat (to be fixed next): - - The “all kept” branch currently delegates to the original unload, which unloads everything. Target behavior is strict no-op when no models are flagged. - -3) Multi-device VRAM cache and CPU reset (patched soft_empty_cache) -- File: ./__init__.py -- Patch site notes: - - `mm.soft_empty_cache` → `soft_empty_cache_distorch2_patched` - - Behavior: - - Detect DisTorch2 active state; clear allocator caches on ALL devices via `soft_empty_cache_multigpu()` from `device_utils.py` - - Adaptive CPU memory reset with optional force to emulate Manager “free_memory”. - - This ensures cache clearing covers all devices in MultiGPU environments beyond a single `mm.get_torch_device()`. - -4) Manager parity helper -- File: ./model_management_mgpu.py -- Function: `force_full_system_cleanup(reason="manual", force=True)` - - Sets both flags (`unload_models=True`, `free_memory=True`) on PromptQueue, identical to Manager’s “Free model and node cache”. - - Useful for testing and ensuring parity from MultiGPU paths. - -Behavioral Summary -- End-to-end Manager parity: - - Manager “Free model and node cache” → POST /free sets flags → Comfy’s prompt_worker calls our patched `unload_all_models` (selective) → `PromptExecutor.reset()` → our patched `soft_empty_cache` (multi-device) → GC. -- Selectiveness guarantee (intended): - - Only DisTorch2 models flagged with `_mgpu_unload_distorch_model=True` are ejected. - - Unflagged models (keep_loaded=True) remain in `mm.current_loaded_models` after the entire flow. -- Current discrepancy: - - When no models are flagged, our patch currently delegates to the original unload (unloads everything). Target fix is to convert this branch to a strict no-op. - -Validation & Logging Hooks -- Memory snapshots: - - Use `multigpu_memory_log(identifier, tag)` in `model_management_mgpu.py` for timestamped CPU/VRAM snapshot lines. -- VRAM cache clearing: - - `soft_empty_cache_multigpu()` logs per-device clearing events (pre/post) in `device_utils.py`. -- Unload path tracing: - - `_mgpu_patched_unload_all_models` logs the counts of kept/unloaded models and updates to `mm.current_loaded_models`. - -Practical Test Recipes -1) Minimal retention test -- Load A(keep=false), B(keep=true), C(keep=true) -- POST /free payload: {"unload_models": true, "free_memory": true} -- Expected: - - Only A is ejected; B and C remain in `mm.current_loaded_models` post-flow. - - CPU RAM drops; VRAM caches clear on all devices. - -2) All-kept test -- Load D(keep=true), E(keep=true) -- POST /free payload: {"unload_models": true, "free_memory": true} -- Expected target behavior: - - No models are ejected (strict no-op in unload step), allocator/cache cleaning only. -- Current behavior (caveat): - - Delegates to original unload → all models may be ejected. This is the next change to reinstate strict no-op. - -References (paths in this repo) -- Per-model flagging: ./distorch_2.py -- Selective unload patch: ./model_management_mgpu.py -- Patched soft empty: ./__init__.py (soft_empty_cache_distorch2_patched) -- Multi-device cache clear: ./device_utils.py -- Manager parity helper: ./model_management_mgpu.py (force_full_system_cleanup) diff --git a/memory-bank/comfy_core.py b/memory-bank/comfy_core.py deleted file mode 100644 index d20add8..0000000 --- a/memory-bank/comfy_core.py +++ /dev/null @@ -1,4042 +0,0 @@ -################## Comfy Core model_management.py #################### - -""" - This file is part of ComfyUI. - Copyright (C) 2024 Comfy - - This program is free software: you can redistribute it and/or modify - it under the terms of the GNU General Public License as published by - the Free Software Foundation, either version 3 of the License, or - (at your option) any later version. - - This program is distributed in the hope that it will be useful, - but WITHOUT ANY WARRANTY; without even the implied warranty of - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - GNU General Public License for more details. - - You should have received a copy of the GNU General Public License - along with this program. If not, see . -""" - -import psutil -import logging -from enum import Enum -from comfy.cli_args import args, PerformanceFeature -import torch -import sys -import importlib -import platform -import weakref -import gc - -class VRAMState(Enum): - DISABLED = 0 #No vram present: no need to move models to vram - NO_VRAM = 1 #Very low vram: enable all the options to save vram - LOW_VRAM = 2 - NORMAL_VRAM = 3 - HIGH_VRAM = 4 - SHARED = 5 #No dedicated vram: memory shared between CPU and GPU but models still need to be moved between both. - -class CPUState(Enum): - GPU = 0 - CPU = 1 - MPS = 2 - -# Determine VRAM State -vram_state = VRAMState.NORMAL_VRAM -set_vram_to = VRAMState.NORMAL_VRAM -cpu_state = CPUState.GPU - -total_vram = 0 - -def get_supported_float8_types(): - float8_types = [] - try: - float8_types.append(torch.float8_e4m3fn) - except: - pass - try: - float8_types.append(torch.float8_e4m3fnuz) - except: - pass - try: - float8_types.append(torch.float8_e5m2) - except: - pass - try: - float8_types.append(torch.float8_e5m2fnuz) - except: - pass - try: - float8_types.append(torch.float8_e8m0fnu) - except: - pass - return float8_types - -FLOAT8_TYPES = get_supported_float8_types() - -xpu_available = False -torch_version = "" -try: - torch_version = torch.version.__version__ - temp = torch_version.split(".") - torch_version_numeric = (int(temp[0]), int(temp[1])) -except: - pass - -lowvram_available = True -if args.deterministic: - logging.info("Using deterministic algorithms for pytorch") - torch.use_deterministic_algorithms(True, warn_only=True) - -directml_enabled = False -if args.directml is not None: - import torch_directml - directml_enabled = True - device_index = args.directml - if device_index < 0: - directml_device = torch_directml.device() - else: - directml_device = torch_directml.device(device_index) - logging.info("Using directml with device: {}".format(torch_directml.device_name(device_index))) - # torch_directml.disable_tiled_resources(True) - lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default. - -try: - import intel_extension_for_pytorch as ipex # noqa: F401 -except: - pass - -try: - _ = torch.xpu.device_count() - xpu_available = torch.xpu.is_available() -except: - xpu_available = False - -try: - if torch.backends.mps.is_available(): - cpu_state = CPUState.MPS - import torch.mps -except: - pass - -try: - import torch_npu # noqa: F401 - _ = torch.npu.device_count() - npu_available = torch.npu.is_available() -except: - npu_available = False - -try: - import torch_mlu # noqa: F401 - _ = torch.mlu.device_count() - mlu_available = torch.mlu.is_available() -except: - mlu_available = False - -try: - ixuca_available = hasattr(torch, "corex") -except: - ixuca_available = False - -if args.cpu: - cpu_state = CPUState.CPU - -def is_intel_xpu(): - global cpu_state - global xpu_available - if cpu_state == CPUState.GPU: - if xpu_available: - return True - return False - -def is_ascend_npu(): - global npu_available - if npu_available: - return True - return False - -def is_mlu(): - global mlu_available - if mlu_available: - return True - return False - -def is_ixuca(): - global ixuca_available - if ixuca_available: - return True - return False - -def get_torch_device(): - global directml_enabled - global cpu_state - if directml_enabled: - global directml_device - return directml_device - if cpu_state == CPUState.MPS: - return torch.device("mps") - if cpu_state == CPUState.CPU: - return torch.device("cpu") - else: - if is_intel_xpu(): - return torch.device("xpu", torch.xpu.current_device()) - elif is_ascend_npu(): - return torch.device("npu", torch.npu.current_device()) - elif is_mlu(): - return torch.device("mlu", torch.mlu.current_device()) - else: - return torch.device(torch.cuda.current_device()) - -def get_total_memory(dev=None, torch_total_too=False): - global directml_enabled - if dev is None: - dev = get_torch_device() - - if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'): - mem_total = psutil.virtual_memory().total - mem_total_torch = mem_total - else: - if directml_enabled: - mem_total = 1024 * 1024 * 1024 #TODO - mem_total_torch = mem_total - elif is_intel_xpu(): - stats = torch.xpu.memory_stats(dev) - mem_reserved = stats['reserved_bytes.all.current'] - mem_total_xpu = torch.xpu.get_device_properties(dev).total_memory - mem_total_torch = mem_reserved - mem_total = mem_total_xpu - elif is_ascend_npu(): - stats = torch.npu.memory_stats(dev) - mem_reserved = stats['reserved_bytes.all.current'] - _, mem_total_npu = torch.npu.mem_get_info(dev) - mem_total_torch = mem_reserved - mem_total = mem_total_npu - elif is_mlu(): - stats = torch.mlu.memory_stats(dev) - mem_reserved = stats['reserved_bytes.all.current'] - _, mem_total_mlu = torch.mlu.mem_get_info(dev) - mem_total_torch = mem_reserved - mem_total = mem_total_mlu - else: - stats = torch.cuda.memory_stats(dev) - mem_reserved = stats['reserved_bytes.all.current'] - _, mem_total_cuda = torch.cuda.mem_get_info(dev) - mem_total_torch = mem_reserved - mem_total = mem_total_cuda - - if torch_total_too: - return (mem_total, mem_total_torch) - else: - return mem_total - -def mac_version(): - try: - return tuple(int(n) for n in platform.mac_ver()[0].split(".")) - except: - return None - -total_vram = get_total_memory(get_torch_device()) / (1024 * 1024) -total_ram = psutil.virtual_memory().total / (1024 * 1024) -logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram)) - -try: - logging.info("pytorch version: {}".format(torch_version)) - mac_ver = mac_version() - if mac_ver is not None: - logging.info("Mac Version {}".format(mac_ver)) -except: - pass - -try: - OOM_EXCEPTION = torch.cuda.OutOfMemoryError -except: - OOM_EXCEPTION = Exception - -XFORMERS_VERSION = "" -XFORMERS_ENABLED_VAE = True -if args.disable_xformers: - XFORMERS_IS_AVAILABLE = False -else: - try: - import xformers - import xformers.ops - XFORMERS_IS_AVAILABLE = True - try: - XFORMERS_IS_AVAILABLE = xformers._has_cpp_library - except: - pass - try: - XFORMERS_VERSION = xformers.version.__version__ - logging.info("xformers version: {}".format(XFORMERS_VERSION)) - if XFORMERS_VERSION.startswith("0.0.18"): - logging.warning("\nWARNING: This version of xformers has a major bug where you will get black images when generating high resolution images.") - logging.warning("Please downgrade or upgrade xformers to a different version.\n") - XFORMERS_ENABLED_VAE = False - except: - pass - except: - XFORMERS_IS_AVAILABLE = False - -def is_nvidia(): - global cpu_state - if cpu_state == CPUState.GPU: - if torch.version.cuda: - return True - return False - -def is_amd(): - global cpu_state - if cpu_state == CPUState.GPU: - if torch.version.hip: - return True - return False - -def amd_min_version(device=None, min_rdna_version=0): - if not is_amd(): - return False - - if is_device_cpu(device): - return False - - arch = torch.cuda.get_device_properties(device).gcnArchName - if arch.startswith('gfx') and len(arch) == 7: - try: - cmp_rdna_version = int(arch[4]) + 2 - except: - cmp_rdna_version = 0 - if cmp_rdna_version >= min_rdna_version: - return True - - return False - -MIN_WEIGHT_MEMORY_RATIO = 0.4 -if is_nvidia(): - MIN_WEIGHT_MEMORY_RATIO = 0.0 - -ENABLE_PYTORCH_ATTENTION = False -if args.use_pytorch_cross_attention: - ENABLE_PYTORCH_ATTENTION = True - XFORMERS_IS_AVAILABLE = False - -try: - if is_nvidia(): - if torch_version_numeric[0] >= 2: - if ENABLE_PYTORCH_ATTENTION == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False: - ENABLE_PYTORCH_ATTENTION = True - if is_intel_xpu() or is_ascend_npu() or is_mlu() or is_ixuca(): - if args.use_split_cross_attention == False and args.use_quad_cross_attention == False: - ENABLE_PYTORCH_ATTENTION = True -except: - pass - - -SUPPORT_FP8_OPS = args.supports_fp8_compute -try: - if is_amd(): - try: - rocm_version = tuple(map(int, str(torch.version.hip).split(".")[:2])) - except: - rocm_version = (6, -1) - arch = torch.cuda.get_device_properties(get_torch_device()).gcnArchName - logging.info("AMD arch: {}".format(arch)) - logging.info("ROCm version: {}".format(rocm_version)) - if args.use_split_cross_attention == False and args.use_quad_cross_attention == False: - if importlib.util.find_spec('triton') is not None: # AMD efficient attention implementation depends on triton. TODO: better way of detecting if it's compiled in or not. - if torch_version_numeric >= (2, 7): # works on 2.6 but doesn't actually seem to improve much - if any((a in arch) for a in ["gfx90a", "gfx942", "gfx1100", "gfx1101", "gfx1151"]): # TODO: more arches, TODO: gfx950 - ENABLE_PYTORCH_ATTENTION = True -# if torch_version_numeric >= (2, 8): -# if any((a in arch) for a in ["gfx1201"]): -# ENABLE_PYTORCH_ATTENTION = True - if torch_version_numeric >= (2, 7) and rocm_version >= (6, 4): - if any((a in arch) for a in ["gfx1200", "gfx1201", "gfx942", "gfx950"]): # TODO: more arches - SUPPORT_FP8_OPS = True - -except: - pass - - -if ENABLE_PYTORCH_ATTENTION: - torch.backends.cuda.enable_math_sdp(True) - torch.backends.cuda.enable_flash_sdp(True) - torch.backends.cuda.enable_mem_efficient_sdp(True) - - -PRIORITIZE_FP16 = False # TODO: remove and replace with something that shows exactly which dtype is faster than the other -try: - if (is_nvidia() or is_amd()) and PerformanceFeature.Fp16Accumulation in args.fast: - torch.backends.cuda.matmul.allow_fp16_accumulation = True - PRIORITIZE_FP16 = True # TODO: limit to cards where it actually boosts performance - logging.info("Enabled fp16 accumulation.") -except: - pass - -try: - if torch_version_numeric >= (2, 5): - torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(True) -except: - logging.warning("Warning, could not set allow_fp16_bf16_reduction_math_sdp") - -if args.lowvram: - set_vram_to = VRAMState.LOW_VRAM - lowvram_available = True -elif args.novram: - set_vram_to = VRAMState.NO_VRAM -elif args.highvram or args.gpu_only: - vram_state = VRAMState.HIGH_VRAM - -FORCE_FP32 = False -if args.force_fp32: - logging.info("Forcing FP32, if this improves things please report it.") - FORCE_FP32 = True - -if lowvram_available: - if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM): - vram_state = set_vram_to - - -if cpu_state != CPUState.GPU: - vram_state = VRAMState.DISABLED - -if cpu_state == CPUState.MPS: - vram_state = VRAMState.SHARED - -logging.info(f"Set vram state to: {vram_state.name}") - -DISABLE_SMART_MEMORY = args.disable_smart_memory - -if DISABLE_SMART_MEMORY: - logging.info("Disabling smart memory management") - -def get_torch_device_name(device): - if hasattr(device, 'type'): - if device.type == "cuda": - try: - allocator_backend = torch.cuda.get_allocator_backend() - except: - allocator_backend = "" - return "{} {} : {}".format(device, torch.cuda.get_device_name(device), allocator_backend) - elif device.type == "xpu": - return "{} {}".format(device, torch.xpu.get_device_name(device)) - else: - return "{}".format(device.type) - elif is_intel_xpu(): - return "{} {}".format(device, torch.xpu.get_device_name(device)) - elif is_ascend_npu(): - return "{} {}".format(device, torch.npu.get_device_name(device)) - elif is_mlu(): - return "{} {}".format(device, torch.mlu.get_device_name(device)) - else: - return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device)) - -try: - logging.info("Device: {}".format(get_torch_device_name(get_torch_device()))) -except: - logging.warning("Could not pick default device.") - - -current_loaded_models = [] - -def module_size(module): - module_mem = 0 - sd = module.state_dict() - for k in sd: - t = sd[k] - module_mem += t.nelement() * t.element_size() - return module_mem - -class LoadedModel: - def __init__(self, model): - self._set_model(model) - self.device = model.load_device - self.real_model = None - self.currently_used = True - self.model_finalizer = None - self._patcher_finalizer = None - - def _set_model(self, model): - self._model = weakref.ref(model) - if model.parent is not None: - self._parent_model = weakref.ref(model.parent) - self._patcher_finalizer = weakref.finalize(model, self._switch_parent) - - def _switch_parent(self): - model = self._parent_model() - if model is not None: - self._set_model(model) - - @property - def model(self): - return self._model() - - def model_memory(self): - return self.model.model_size() - - def model_loaded_memory(self): - return self.model.loaded_size() - - def model_offloaded_memory(self): - return self.model.model_size() - self.model.loaded_size() - - def model_memory_required(self, device): - if device == self.model.current_loaded_device(): - return self.model_offloaded_memory() - else: - return self.model_memory() - - def model_load(self, lowvram_model_memory=0, force_patch_weights=False): - self.model.model_patches_to(self.device) - self.model.model_patches_to(self.model.model_dtype()) - - # if self.model.loaded_size() > 0: - use_more_vram = lowvram_model_memory - if use_more_vram == 0: - use_more_vram = 1e32 - self.model_use_more_vram(use_more_vram, force_patch_weights=force_patch_weights) - real_model = self.model.model - - if is_intel_xpu() and not args.disable_ipex_optimize and 'ipex' in globals() and real_model is not None: - with torch.no_grad(): - real_model = ipex.optimize(real_model.eval(), inplace=True, graph_mode=True, concat_linear=True) - - self.real_model = weakref.ref(real_model) - self.model_finalizer = weakref.finalize(real_model, cleanup_models) - return real_model - - def should_reload_model(self, force_patch_weights=False): - if force_patch_weights and self.model.lowvram_patch_counter() > 0: - return True - return False - - def model_unload(self, memory_to_free=None, unpatch_weights=True): - if memory_to_free is not None: - if memory_to_free < self.model.loaded_size(): - freed = self.model.partially_unload(self.model.offload_device, memory_to_free) - if freed >= memory_to_free: - return False - self.model.detach(unpatch_weights) - self.model_finalizer.detach() - self.model_finalizer = None - self.real_model = None - return True - - def model_use_more_vram(self, extra_memory, force_patch_weights=False): - return self.model.partially_load(self.device, extra_memory, force_patch_weights=force_patch_weights) - - def __eq__(self, other): - return self.model is other.model - - def __del__(self): - if self._patcher_finalizer is not None: - self._patcher_finalizer.detach() - - def is_dead(self): - return self.real_model() is not None and self.model is None - - -def use_more_memory(extra_memory, loaded_models, device): - for m in loaded_models: - if m.device == device: - extra_memory -= m.model_use_more_vram(extra_memory) - if extra_memory <= 0: - break - -def offloaded_memory(loaded_models, device): - offloaded_mem = 0 - for m in loaded_models: - if m.device == device: - offloaded_mem += m.model_offloaded_memory() - return offloaded_mem - -WINDOWS = any(platform.win32_ver()) - -EXTRA_RESERVED_VRAM = 400 * 1024 * 1024 -if WINDOWS: - EXTRA_RESERVED_VRAM = 600 * 1024 * 1024 #Windows is higher because of the shared vram issue - if total_vram > (15 * 1024): # more extra reserved vram on 16GB+ cards - EXTRA_RESERVED_VRAM += 100 * 1024 * 1024 - -if args.reserve_vram is not None: - EXTRA_RESERVED_VRAM = args.reserve_vram * 1024 * 1024 * 1024 - logging.debug("Reserving {}MB vram for other applications.".format(EXTRA_RESERVED_VRAM / (1024 * 1024))) - -def extra_reserved_memory(): - return EXTRA_RESERVED_VRAM - -def minimum_inference_memory(): - return (1024 * 1024 * 1024) * 0.8 + extra_reserved_memory() - -def free_memory(memory_required, device, keep_loaded=[]): - cleanup_models_gc() - unloaded_model = [] - can_unload = [] - unloaded_models = [] - - for i in range(len(current_loaded_models) -1, -1, -1): - shift_model = current_loaded_models[i] - if shift_model.device == device: - if shift_model not in keep_loaded and not shift_model.is_dead(): - can_unload.append((-shift_model.model_offloaded_memory(), sys.getrefcount(shift_model.model), shift_model.model_memory(), i)) - shift_model.currently_used = False - - for x in sorted(can_unload): - i = x[-1] - memory_to_free = None - if not DISABLE_SMART_MEMORY: - free_mem = get_free_memory(device) - if free_mem > memory_required: - break - memory_to_free = memory_required - free_mem - logging.debug(f"Unloading {current_loaded_models[i].model.model.__class__.__name__}") - if current_loaded_models[i].model_unload(memory_to_free): - unloaded_model.append(i) - - for i in sorted(unloaded_model, reverse=True): - unloaded_models.append(current_loaded_models.pop(i)) - - if len(unloaded_model) > 0: - soft_empty_cache() - else: - if vram_state != VRAMState.HIGH_VRAM: - mem_free_total, mem_free_torch = get_free_memory(device, torch_free_too=True) - if mem_free_torch > mem_free_total * 0.25: - soft_empty_cache() - return unloaded_models - -def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None, force_full_load=False): - cleanup_models_gc() - global vram_state - - inference_memory = minimum_inference_memory() - extra_mem = max(inference_memory, memory_required + extra_reserved_memory()) - if minimum_memory_required is None: - minimum_memory_required = extra_mem - else: - minimum_memory_required = max(inference_memory, minimum_memory_required + extra_reserved_memory()) - - models_temp = set() - for m in models: - models_temp.add(m) - for mm in m.model_patches_models(): - models_temp.add(mm) - - models = models_temp - - models_to_load = [] - - for x in models: - loaded_model = LoadedModel(x) - try: - loaded_model_index = current_loaded_models.index(loaded_model) - except: - loaded_model_index = None - - if loaded_model_index is not None: - loaded = current_loaded_models[loaded_model_index] - loaded.currently_used = True - models_to_load.append(loaded) - else: - if hasattr(x, "model"): - logging.info(f"Requested to load {x.model.__class__.__name__}") - models_to_load.append(loaded_model) - - for loaded_model in models_to_load: - to_unload = [] - for i in range(len(current_loaded_models)): - if loaded_model.model.is_clone(current_loaded_models[i].model): - to_unload = [i] + to_unload - for i in to_unload: - model_to_unload = current_loaded_models.pop(i) - model_to_unload.model.detach(unpatch_all=False) - model_to_unload.model_finalizer.detach() - - total_memory_required = {} - for loaded_model in models_to_load: - total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device) - - for device in total_memory_required: - if device != torch.device("cpu"): - free_memory(total_memory_required[device] * 1.1 + extra_mem, device) - - for device in total_memory_required: - if device != torch.device("cpu"): - free_mem = get_free_memory(device) - if free_mem < minimum_memory_required: - models_l = free_memory(minimum_memory_required, device) - logging.info("{} models unloaded.".format(len(models_l))) - - for loaded_model in models_to_load: - model = loaded_model.model - torch_dev = model.load_device - if is_device_cpu(torch_dev): - vram_set_state = VRAMState.DISABLED - else: - vram_set_state = vram_state - lowvram_model_memory = 0 - if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM) and not force_full_load: - loaded_memory = loaded_model.model_loaded_memory() - current_free_mem = get_free_memory(torch_dev) + loaded_memory - - lowvram_model_memory = max(128 * 1024 * 1024, (current_free_mem - minimum_memory_required), min(current_free_mem * MIN_WEIGHT_MEMORY_RATIO, current_free_mem - minimum_inference_memory())) - lowvram_model_memory = max(0.1, lowvram_model_memory - loaded_memory) - - if vram_set_state == VRAMState.NO_VRAM: - lowvram_model_memory = 0.1 - - loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights) - current_loaded_models.insert(0, loaded_model) - return - -def load_model_gpu(model): - return load_models_gpu([model]) - -def loaded_models(only_currently_used=False): - output = [] - for m in current_loaded_models: - if only_currently_used: - if not m.currently_used: - continue - - output.append(m.model) - return output - - -def cleanup_models_gc(): - do_gc = False - for i in range(len(current_loaded_models)): - cur = current_loaded_models[i] - if cur.is_dead(): - logging.info("Potential memory leak detected with model {}, doing a full garbage collect, for maximum performance avoid circular references in the model code.".format(cur.real_model().__class__.__name__)) - do_gc = True - break - - if do_gc: - gc.collect() - soft_empty_cache() - - for i in range(len(current_loaded_models)): - cur = current_loaded_models[i] - if cur.is_dead(): - logging.warning("WARNING, memory leak with model {}. Please make sure it is not being referenced from somewhere.".format(cur.real_model().__class__.__name__)) - - - -def cleanup_models(): - to_delete = [] - for i in range(len(current_loaded_models)): - if current_loaded_models[i].real_model() is None: - to_delete = [i] + to_delete - - for i in to_delete: - x = current_loaded_models.pop(i) - del x - -def dtype_size(dtype): - dtype_size = 4 - if dtype == torch.float16 or dtype == torch.bfloat16: - dtype_size = 2 - elif dtype == torch.float32: - dtype_size = 4 - else: - try: - dtype_size = dtype.itemsize - except: #Old pytorch doesn't have .itemsize - pass - return dtype_size - -def unet_offload_device(): - if vram_state == VRAMState.HIGH_VRAM: - return get_torch_device() - else: - return torch.device("cpu") - -def unet_inital_load_device(parameters, dtype): - torch_dev = get_torch_device() - if vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.SHARED: - return torch_dev - - cpu_dev = torch.device("cpu") - if DISABLE_SMART_MEMORY or vram_state == VRAMState.NO_VRAM: - return cpu_dev - - model_size = dtype_size(dtype) * parameters - - mem_dev = get_free_memory(torch_dev) - mem_cpu = get_free_memory(cpu_dev) - if mem_dev > mem_cpu and model_size < mem_dev: - return torch_dev - else: - return cpu_dev - -def maximum_vram_for_weights(device=None): - return (get_total_memory(device) * 0.88 - minimum_inference_memory()) - -def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32], weight_dtype=None): - if model_params < 0: - model_params = 1000000000000000000000 - if args.fp32_unet: - return torch.float32 - if args.fp64_unet: - return torch.float64 - if args.bf16_unet: - return torch.bfloat16 - if args.fp16_unet: - return torch.float16 - if args.fp8_e4m3fn_unet: - return torch.float8_e4m3fn - if args.fp8_e5m2_unet: - return torch.float8_e5m2 - if args.fp8_e8m0fnu_unet: - return torch.float8_e8m0fnu - - fp8_dtype = None - if weight_dtype in FLOAT8_TYPES: - fp8_dtype = weight_dtype - - if fp8_dtype is not None: - if supports_fp8_compute(device): #if fp8 compute is supported the casting is most likely not expensive - return fp8_dtype - - free_model_memory = maximum_vram_for_weights(device) - if model_params * 2 > free_model_memory: - return fp8_dtype - - if PRIORITIZE_FP16 or weight_dtype == torch.float16: - if torch.float16 in supported_dtypes and should_use_fp16(device=device, model_params=model_params): - return torch.float16 - - for dt in supported_dtypes: - if dt == torch.float16 and should_use_fp16(device=device, model_params=model_params): - if torch.float16 in supported_dtypes: - return torch.float16 - if dt == torch.bfloat16 and should_use_bf16(device, model_params=model_params): - if torch.bfloat16 in supported_dtypes: - return torch.bfloat16 - - for dt in supported_dtypes: - if dt == torch.float16 and should_use_fp16(device=device, model_params=model_params, manual_cast=True): - if torch.float16 in supported_dtypes: - return torch.float16 - if dt == torch.bfloat16 and should_use_bf16(device, model_params=model_params, manual_cast=True): - if torch.bfloat16 in supported_dtypes: - return torch.bfloat16 - - return torch.float32 - -# None means no manual cast -def unet_manual_cast(weight_dtype, inference_device, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]): - if weight_dtype == torch.float32 or weight_dtype == torch.float64: - return None - - fp16_supported = should_use_fp16(inference_device, prioritize_performance=False) - if fp16_supported and weight_dtype == torch.float16: - return None - - bf16_supported = should_use_bf16(inference_device) - if bf16_supported and weight_dtype == torch.bfloat16: - return None - - fp16_supported = should_use_fp16(inference_device, prioritize_performance=True) - if PRIORITIZE_FP16 and fp16_supported and torch.float16 in supported_dtypes: - return torch.float16 - - for dt in supported_dtypes: - if dt == torch.float16 and fp16_supported: - return torch.float16 - if dt == torch.bfloat16 and bf16_supported: - return torch.bfloat16 - - return torch.float32 - -def text_encoder_offload_device(): - if args.gpu_only: - return get_torch_device() - else: - return torch.device("cpu") - -def text_encoder_device(): - if args.gpu_only: - return get_torch_device() - elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM: - if should_use_fp16(prioritize_performance=False): - return get_torch_device() - else: - return torch.device("cpu") - else: - return torch.device("cpu") - -def text_encoder_initial_device(load_device, offload_device, model_size=0): - if load_device == offload_device or model_size <= 1024 * 1024 * 1024: - return offload_device - - if is_device_mps(load_device): - return load_device - - mem_l = get_free_memory(load_device) - mem_o = get_free_memory(offload_device) - if mem_l > (mem_o * 0.5) and model_size * 1.2 < mem_l: - return load_device - else: - return offload_device - -def text_encoder_dtype(device=None): - if args.fp8_e4m3fn_text_enc: - return torch.float8_e4m3fn - elif args.fp8_e5m2_text_enc: - return torch.float8_e5m2 - elif args.fp16_text_enc: - return torch.float16 - elif args.bf16_text_enc: - return torch.bfloat16 - elif args.fp32_text_enc: - return torch.float32 - - if is_device_cpu(device): - return torch.float16 - - return torch.float16 - - -def intermediate_device(): - if args.gpu_only: - return get_torch_device() - else: - return torch.device("cpu") - -def vae_device(): - if args.cpu_vae: - return torch.device("cpu") - return get_torch_device() - -def vae_offload_device(): - if args.gpu_only: - return get_torch_device() - else: - return torch.device("cpu") - -def vae_dtype(device=None, allowed_dtypes=[]): - if args.fp16_vae: - return torch.float16 - elif args.bf16_vae: - return torch.bfloat16 - elif args.fp32_vae: - return torch.float32 - - for d in allowed_dtypes: - if d == torch.float16 and should_use_fp16(device): - return d - - # NOTE: bfloat16 seems to work on AMD for the VAE but is extremely slow in some cases compared to fp32 - # slowness still a problem on pytorch nightly 2.9.0.dev20250720+rocm6.4 tested on RDNA3 - # also a problem on RDNA4 except fp32 is also slow there. - # This is due to large bf16 convolutions being extremely slow. - if d == torch.bfloat16 and ((not is_amd()) or amd_min_version(device, min_rdna_version=4)) and should_use_bf16(device): - return d - - return torch.float32 - -def get_autocast_device(dev): - if hasattr(dev, 'type'): - return dev.type - return "cuda" - -def supports_dtype(device, dtype): #TODO - if dtype == torch.float32: - return True - if is_device_cpu(device): - return False - if dtype == torch.float16: - return True - if dtype == torch.bfloat16: - return True - return False - -def supports_cast(device, dtype): #TODO - if dtype == torch.float32: - return True - if dtype == torch.float16: - return True - if directml_enabled: #TODO: test this - return False - if dtype == torch.bfloat16: - return True - if is_device_mps(device): - return False - if dtype == torch.float8_e4m3fn: - return True - if dtype == torch.float8_e5m2: - return True - return False - -def pick_weight_dtype(dtype, fallback_dtype, device=None): - if dtype is None: - dtype = fallback_dtype - elif dtype_size(dtype) > dtype_size(fallback_dtype): - dtype = fallback_dtype - - if not supports_cast(device, dtype): - dtype = fallback_dtype - - return dtype - -def device_supports_non_blocking(device): - if args.force_non_blocking: - return True - if is_device_mps(device): - return False #pytorch bug? mps doesn't support non blocking - if is_intel_xpu(): #xpu does support non blocking but it is slower on iGPUs for some reason so disable by default until situation changes - return False - if args.deterministic: #TODO: figure out why deterministic breaks non blocking from gpu to cpu (previews) - return False - if directml_enabled: - return False - return True - -def device_should_use_non_blocking(device): - if not device_supports_non_blocking(device): - return False - return False - # return True #TODO: figure out why this causes memory issues on Nvidia and possibly others - -def force_channels_last(): - if args.force_channels_last: - return True - - #TODO - return False - - -STREAMS = {} -NUM_STREAMS = 1 -if args.async_offload: - NUM_STREAMS = 2 - logging.info("Using async weight offloading with {} streams".format(NUM_STREAMS)) - -stream_counters = {} -def get_offload_stream(device): - stream_counter = stream_counters.get(device, 0) - if NUM_STREAMS <= 1: - return None - - if device in STREAMS: - ss = STREAMS[device] - s = ss[stream_counter] - stream_counter = (stream_counter + 1) % len(ss) - if is_device_cuda(device): - ss[stream_counter].wait_stream(torch.cuda.current_stream()) - elif is_device_xpu(device): - ss[stream_counter].wait_stream(torch.xpu.current_stream()) - stream_counters[device] = stream_counter - return s - elif is_device_cuda(device): - ss = [] - for k in range(NUM_STREAMS): - ss.append(torch.cuda.Stream(device=device, priority=0)) - STREAMS[device] = ss - s = ss[stream_counter] - stream_counter = (stream_counter + 1) % len(ss) - stream_counters[device] = stream_counter - return s - elif is_device_xpu(device): - ss = [] - for k in range(NUM_STREAMS): - ss.append(torch.xpu.Stream(device=device, priority=0)) - STREAMS[device] = ss - s = ss[stream_counter] - stream_counter = (stream_counter + 1) % len(ss) - stream_counters[device] = stream_counter - return s - return None - -def sync_stream(device, stream): - if stream is None: - return - if is_device_cuda(device): - torch.cuda.current_stream().wait_stream(stream) - elif is_device_xpu(device): - torch.xpu.current_stream().wait_stream(stream) - -def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False, stream=None): - if device is None or weight.device == device: - if not copy: - if dtype is None or weight.dtype == dtype: - return weight - if stream is not None: - with stream: - return weight.to(dtype=dtype, copy=copy) - return weight.to(dtype=dtype, copy=copy) - - if stream is not None: - with stream: - r = torch.empty_like(weight, dtype=dtype, device=device) - r.copy_(weight, non_blocking=non_blocking) - else: - r = torch.empty_like(weight, dtype=dtype, device=device) - r.copy_(weight, non_blocking=non_blocking) - return r - -def cast_to_device(tensor, device, dtype, copy=False): - non_blocking = device_supports_non_blocking(device) - return cast_to(tensor, dtype=dtype, device=device, non_blocking=non_blocking, copy=copy) - -def sage_attention_enabled(): - return args.use_sage_attention - -def flash_attention_enabled(): - return args.use_flash_attention - -def xformers_enabled(): - global directml_enabled - global cpu_state - if cpu_state != CPUState.GPU: - return False - if is_intel_xpu(): - return False - if is_ascend_npu(): - return False - if is_mlu(): - return False - if is_ixuca(): - return False - if directml_enabled: - return False - return XFORMERS_IS_AVAILABLE - - -def xformers_enabled_vae(): - enabled = xformers_enabled() - if not enabled: - return False - - return XFORMERS_ENABLED_VAE - -def pytorch_attention_enabled(): - global ENABLE_PYTORCH_ATTENTION - return ENABLE_PYTORCH_ATTENTION - -def pytorch_attention_enabled_vae(): - if is_amd(): - return False # enabling pytorch attention on AMD currently causes crash when doing high res - return pytorch_attention_enabled() - -def pytorch_attention_flash_attention(): - global ENABLE_PYTORCH_ATTENTION - if ENABLE_PYTORCH_ATTENTION: - #TODO: more reliable way of checking for flash attention? - if is_nvidia(): - return True - if is_intel_xpu(): - return True - if is_ascend_npu(): - return True - if is_mlu(): - return True - if is_amd(): - return True #if you have pytorch attention enabled on AMD it probably supports at least mem efficient attention - if is_ixuca(): - return True - return False - -def force_upcast_attention_dtype(): - upcast = args.force_upcast_attention - - macos_version = mac_version() - if macos_version is not None and ((14, 5) <= macos_version): # black image bug on recent versions of macOS, I don't think it's ever getting fixed - upcast = True - - if upcast: - return {torch.float16: torch.float32} - else: - return None - -def get_free_memory(dev=None, torch_free_too=False): - global directml_enabled - if dev is None: - dev = get_torch_device() - - if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'): - mem_free_total = psutil.virtual_memory().available - mem_free_torch = mem_free_total - else: - if directml_enabled: - mem_free_total = 1024 * 1024 * 1024 #TODO - mem_free_torch = mem_free_total - elif is_intel_xpu(): - stats = torch.xpu.memory_stats(dev) - mem_active = stats['active_bytes.all.current'] - mem_reserved = stats['reserved_bytes.all.current'] - mem_free_xpu = torch.xpu.get_device_properties(dev).total_memory - mem_reserved - mem_free_torch = mem_reserved - mem_active - mem_free_total = mem_free_xpu + mem_free_torch - elif is_ascend_npu(): - stats = torch.npu.memory_stats(dev) - mem_active = stats['active_bytes.all.current'] - mem_reserved = stats['reserved_bytes.all.current'] - mem_free_npu, _ = torch.npu.mem_get_info(dev) - mem_free_torch = mem_reserved - mem_active - mem_free_total = mem_free_npu + mem_free_torch - elif is_mlu(): - stats = torch.mlu.memory_stats(dev) - mem_active = stats['active_bytes.all.current'] - mem_reserved = stats['reserved_bytes.all.current'] - mem_free_mlu, _ = torch.mlu.mem_get_info(dev) - mem_free_torch = mem_reserved - mem_active - mem_free_total = mem_free_mlu + mem_free_torch - else: - stats = torch.cuda.memory_stats(dev) - mem_active = stats['active_bytes.all.current'] - mem_reserved = stats['reserved_bytes.all.current'] - mem_free_cuda, _ = torch.cuda.mem_get_info(dev) - mem_free_torch = mem_reserved - mem_active - mem_free_total = mem_free_cuda + mem_free_torch - - if torch_free_too: - return (mem_free_total, mem_free_torch) - else: - return mem_free_total - -def cpu_mode(): - global cpu_state - return cpu_state == CPUState.CPU - -def mps_mode(): - global cpu_state - return cpu_state == CPUState.MPS - -def is_device_type(device, type): - if hasattr(device, 'type'): - if (device.type == type): - return True - return False - -def is_device_cpu(device): - return is_device_type(device, 'cpu') - -def is_device_mps(device): - return is_device_type(device, 'mps') - -def is_device_xpu(device): - return is_device_type(device, 'xpu') - -def is_device_cuda(device): - return is_device_type(device, 'cuda') - -def is_directml_enabled(): - global directml_enabled - if directml_enabled: - return True - - return False - -def should_use_fp16(device=None, model_params=0, prioritize_performance=True, manual_cast=False): - if device is not None: - if is_device_cpu(device): - return False - - if args.force_fp16: - return True - - if FORCE_FP32: - return False - - if is_directml_enabled(): - return True - - if (device is not None and is_device_mps(device)) or mps_mode(): - return True - - if cpu_mode(): - return False - - if is_intel_xpu(): - if torch_version_numeric < (2, 3): - return True - else: - return torch.xpu.get_device_properties(device).has_fp16 - - if is_ascend_npu(): - return True - - if is_mlu(): - return True - - if is_ixuca(): - return True - - if torch.version.hip: - return True - - props = torch.cuda.get_device_properties(device) - if props.major >= 8: - return True - - if props.major < 6: - return False - - #FP16 is confirmed working on a 1080 (GP104) and on latest pytorch actually seems faster than fp32 - nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050", "p40", "p100", "p6", "p4"] - for x in nvidia_10_series: - if x in props.name.lower(): - if WINDOWS or manual_cast: - return True - else: - return False #weird linux behavior where fp32 is faster - - if manual_cast: - free_model_memory = maximum_vram_for_weights(device) - if (not prioritize_performance) or model_params * 4 > free_model_memory: - return True - - if props.major < 7: - return False - - #FP16 is just broken on these cards - nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600", "MX550", "MX450", "CMP 30HX", "T2000", "T1000", "T1200"] - for x in nvidia_16_series: - if x in props.name: - return False - - return True - -def should_use_bf16(device=None, model_params=0, prioritize_performance=True, manual_cast=False): - if device is not None: - if is_device_cpu(device): #TODO ? bf16 works on CPU but is extremely slow - return False - - if FORCE_FP32: - return False - - if directml_enabled: - return False - - if (device is not None and is_device_mps(device)) or mps_mode(): - if mac_version() < (14,): - return False - return True - - if cpu_mode(): - return False - - if is_intel_xpu(): - if torch_version_numeric < (2, 3): - return True - else: - return torch.xpu.is_bf16_supported() - - if is_ascend_npu(): - return True - - if is_ixuca(): - return True - - if is_amd(): - arch = torch.cuda.get_device_properties(device).gcnArchName - if any((a in arch) for a in ["gfx1030", "gfx1031", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]): # RDNA2 and older don't support bf16 - if manual_cast: - return True - return False - - props = torch.cuda.get_device_properties(device) - - if is_mlu(): - if props.major > 3: - return True - - if props.major >= 8: - return True - - bf16_works = torch.cuda.is_bf16_supported() - - if bf16_works and manual_cast: - free_model_memory = maximum_vram_for_weights(device) - if (not prioritize_performance) or model_params * 4 > free_model_memory: - return True - - return False - -def supports_fp8_compute(device=None): - if SUPPORT_FP8_OPS: - return True - - if not is_nvidia(): - return False - - props = torch.cuda.get_device_properties(device) - if props.major >= 9: - return True - if props.major < 8: - return False - if props.minor < 9: - return False - - if torch_version_numeric < (2, 3): - return False - - if WINDOWS: - if torch_version_numeric < (2, 4): - return False - - return True - -def extended_fp16_support(): - # TODO: check why some models work with fp16 on newer torch versions but not on older - if torch_version_numeric < (2, 7): - return False - - return True - -def soft_empty_cache(force=False): - global cpu_state - if cpu_state == CPUState.MPS: - torch.mps.empty_cache() - elif is_intel_xpu(): - torch.xpu.empty_cache() - elif is_ascend_npu(): - torch.npu.empty_cache() - elif is_mlu(): - torch.mlu.empty_cache() - elif torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.ipc_collect() - -def unload_all_models(): - free_memory(1e30, get_torch_device()) - - -#TODO: might be cleaner to put this somewhere else -import threading - -class InterruptProcessingException(Exception): - pass - -interrupt_processing_mutex = threading.RLock() - -interrupt_processing = False -def interrupt_current_processing(value=True): - global interrupt_processing - global interrupt_processing_mutex - with interrupt_processing_mutex: - interrupt_processing = value - -def processing_interrupted(): - global interrupt_processing - global interrupt_processing_mutex - with interrupt_processing_mutex: - return interrupt_processing - -def throw_exception_if_processing_interrupted(): - global interrupt_processing - global interrupt_processing_mutex - with interrupt_processing_mutex: - if interrupt_processing: - interrupt_processing = False - raise InterruptProcessingException() - - -################### ComfyUI Core server.py #################### -import os -import sys -import asyncio -import traceback - -import nodes -import folder_paths -import execution -import uuid -import urllib -import json -import glob -import struct -import ssl -import socket -import ipaddress -from PIL import Image, ImageOps -from PIL.PngImagePlugin import PngInfo -from io import BytesIO - -import aiohttp -from aiohttp import web -import logging - -import mimetypes -from comfy.cli_args import args -import comfy.utils -import comfy.model_management -from comfy_api import feature_flags -import node_helpers -from comfyui_version import __version__ -from app.frontend_management import FrontendManager -from comfy_api.internal import _ComfyNodeInternal - -from app.user_manager import UserManager -from app.model_manager import ModelFileManager -from app.custom_node_manager import CustomNodeManager -from typing import Optional, Union -from api_server.routes.internal.internal_routes import InternalRoutes -from protocol import BinaryEventTypes - -# Import cache control middleware -from middleware.cache_middleware import cache_control - -async def send_socket_catch_exception(function, message): - try: - await function(message) - except (aiohttp.ClientError, aiohttp.ClientPayloadError, ConnectionResetError, BrokenPipeError, ConnectionError) as err: - logging.warning("send error: {}".format(err)) - -@web.middleware -async def compress_body(request: web.Request, handler): - accept_encoding = request.headers.get("Accept-Encoding", "") - response: web.Response = await handler(request) - if not isinstance(response, web.Response): - return response - if response.content_type not in ["application/json", "text/plain"]: - return response - if response.body and "gzip" in accept_encoding: - response.enable_compression() - return response - - -def create_cors_middleware(allowed_origin: str): - @web.middleware - async def cors_middleware(request: web.Request, handler): - if request.method == "OPTIONS": - # Pre-flight request. Reply successfully: - response = web.Response() - else: - response = await handler(request) - - response.headers['Access-Control-Allow-Origin'] = allowed_origin - response.headers['Access-Control-Allow-Methods'] = 'POST, GET, DELETE, PUT, OPTIONS' - response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization' - response.headers['Access-Control-Allow-Credentials'] = 'true' - return response - - return cors_middleware - -def is_loopback(host): - if host is None: - return False - try: - if ipaddress.ip_address(host).is_loopback: - return True - else: - return False - except: - pass - - loopback = False - for family in (socket.AF_INET, socket.AF_INET6): - try: - r = socket.getaddrinfo(host, None, family, socket.SOCK_STREAM) - for family, _, _, _, sockaddr in r: - if not ipaddress.ip_address(sockaddr[0]).is_loopback: - return loopback - else: - loopback = True - except socket.gaierror: - pass - - return loopback - - -def create_origin_only_middleware(): - @web.middleware - async def origin_only_middleware(request: web.Request, handler): - #this code is used to prevent the case where a random website can queue comfy workflows by making a POST to 127.0.0.1 which browsers don't prevent for some dumb reason. - #in that case the Host and Origin hostnames won't match - #I know the proper fix would be to add a cookie but this should take care of the problem in the meantime - if 'Host' in request.headers and 'Origin' in request.headers: - host = request.headers['Host'] - origin = request.headers['Origin'] - host_domain = host.lower() - parsed = urllib.parse.urlparse(origin) - origin_domain = parsed.netloc.lower() - host_domain_parsed = urllib.parse.urlsplit('//' + host_domain) - - #limit the check to when the host domain is localhost, this makes it slightly less safe but should still prevent the exploit - loopback = is_loopback(host_domain_parsed.hostname) - - if parsed.port is None: #if origin doesn't have a port strip it from the host to handle weird browsers, same for host - host_domain = host_domain_parsed.hostname - if host_domain_parsed.port is None: - origin_domain = parsed.hostname - - if loopback and host_domain is not None and origin_domain is not None and len(host_domain) > 0 and len(origin_domain) > 0: - if host_domain != origin_domain: - logging.warning("WARNING: request with non matching host and origin {} != {}, returning 403".format(host_domain, origin_domain)) - return web.Response(status=403) - - if request.method == "OPTIONS": - response = web.Response() - else: - response = await handler(request) - - return response - - return origin_only_middleware - -class PromptServer(): - def __init__(self, loop): - PromptServer.instance = self - - mimetypes.init() - mimetypes.add_type('application/javascript; charset=utf-8', '.js') - mimetypes.add_type('image/webp', '.webp') - - self.user_manager = UserManager() - self.model_file_manager = ModelFileManager() - self.custom_node_manager = CustomNodeManager() - self.internal_routes = InternalRoutes(self) - self.supports = ["custom_nodes_from_web"] - self.prompt_queue = execution.PromptQueue(self) - self.loop = loop - self.messages = asyncio.Queue() - self.client_session:Optional[aiohttp.ClientSession] = None - self.number = 0 - - middlewares = [cache_control] - if args.enable_compress_response_body: - middlewares.append(compress_body) - - if args.enable_cors_header: - middlewares.append(create_cors_middleware(args.enable_cors_header)) - else: - middlewares.append(create_origin_only_middleware()) - - max_upload_size = round(args.max_upload_size * 1024 * 1024) - self.app = web.Application(client_max_size=max_upload_size, middlewares=middlewares) - self.sockets = dict() - self.sockets_metadata = dict() - self.web_root = ( - FrontendManager.init_frontend(args.front_end_version) - if args.front_end_root is None - else args.front_end_root - ) - logging.info(f"[Prompt Server] web root: {self.web_root}") - routes = web.RouteTableDef() - self.routes = routes - self.last_node_id = None - self.client_id = None - - self.on_prompt_handlers = [] - - @routes.get('/ws') - async def websocket_handler(request): - ws = web.WebSocketResponse() - await ws.prepare(request) - sid = request.rel_url.query.get('clientId', '') - if sid: - # Reusing existing session, remove old - self.sockets.pop(sid, None) - else: - sid = uuid.uuid4().hex - - # Store WebSocket for backward compatibility - self.sockets[sid] = ws - # Store metadata separately - self.sockets_metadata[sid] = {"feature_flags": {}} - - try: - # Send initial state to the new client - await self.send("status", {"status": self.get_queue_info(), "sid": sid}, sid) - # On reconnect if we are the currently executing client send the current node - if self.client_id == sid and self.last_node_id is not None: - await self.send("executing", { "node": self.last_node_id }, sid) - - # Flag to track if we've received the first message - first_message = True - - async for msg in ws: - if msg.type == aiohttp.WSMsgType.ERROR: - logging.warning('ws connection closed with exception %s' % ws.exception()) - elif msg.type == aiohttp.WSMsgType.TEXT: - try: - data = json.loads(msg.data) - # Check if first message is feature flags - if first_message and data.get("type") == "feature_flags": - # Store client feature flags - client_flags = data.get("data", {}) - self.sockets_metadata[sid]["feature_flags"] = client_flags - - # Send server feature flags in response - await self.send( - "feature_flags", - feature_flags.get_server_features(), - sid, - ) - - logging.debug( - f"Feature flags negotiated for client {sid}: {client_flags}" - ) - first_message = False - except json.JSONDecodeError: - logging.warning( - f"Invalid JSON received from client {sid}: {msg.data}" - ) - except Exception as e: - logging.error(f"Error processing WebSocket message: {e}") - finally: - self.sockets.pop(sid, None) - self.sockets_metadata.pop(sid, None) - return ws - - @routes.get("/") - async def get_root(request): - response = web.FileResponse(os.path.join(self.web_root, "index.html")) - response.headers['Cache-Control'] = 'no-cache' - response.headers["Pragma"] = "no-cache" - response.headers["Expires"] = "0" - return response - - @routes.get("/embeddings") - def get_embeddings(request): - embeddings = folder_paths.get_filename_list("embeddings") - return web.json_response(list(map(lambda a: os.path.splitext(a)[0], embeddings))) - - @routes.get("/models") - def list_model_types(request): - model_types = list(folder_paths.folder_names_and_paths.keys()) - - return web.json_response(model_types) - - @routes.get("/models/{folder}") - async def get_models(request): - folder = request.match_info.get("folder", None) - if not folder in folder_paths.folder_names_and_paths: - return web.Response(status=404) - files = folder_paths.get_filename_list(folder) - return web.json_response(files) - - @routes.get("/extensions") - async def get_extensions(request): - files = glob.glob(os.path.join( - glob.escape(self.web_root), 'extensions/**/*.js'), recursive=True) - - extensions = list(map(lambda f: "/" + os.path.relpath(f, self.web_root).replace("\\", "/"), files)) - - for name, dir in nodes.EXTENSION_WEB_DIRS.items(): - files = glob.glob(os.path.join(glob.escape(dir), '**/*.js'), recursive=True) - extensions.extend(list(map(lambda f: "/extensions/" + urllib.parse.quote( - name) + "/" + os.path.relpath(f, dir).replace("\\", "/"), files))) - - return web.json_response(extensions) - - def get_dir_by_type(dir_type): - if dir_type is None: - dir_type = "input" - - if dir_type == "input": - type_dir = folder_paths.get_input_directory() - elif dir_type == "temp": - type_dir = folder_paths.get_temp_directory() - elif dir_type == "output": - type_dir = folder_paths.get_output_directory() - - return type_dir, dir_type - - def compare_image_hash(filepath, image): - hasher = node_helpers.hasher() - - # function to compare hashes of two images to see if it already exists, fix to #3465 - if os.path.exists(filepath): - a = hasher() - b = hasher() - with open(filepath, "rb") as f: - a.update(f.read()) - b.update(image.file.read()) - image.file.seek(0) - return a.hexdigest() == b.hexdigest() - return False - - def image_upload(post, image_save_function=None): - image = post.get("image") - overwrite = post.get("overwrite") - image_is_duplicate = False - - image_upload_type = post.get("type") - upload_dir, image_upload_type = get_dir_by_type(image_upload_type) - - if image and image.file: - filename = image.filename - if not filename: - return web.Response(status=400) - - subfolder = post.get("subfolder", "") - full_output_folder = os.path.join(upload_dir, os.path.normpath(subfolder)) - filepath = os.path.abspath(os.path.join(full_output_folder, filename)) - - if os.path.commonpath((upload_dir, filepath)) != upload_dir: - return web.Response(status=400) - - if not os.path.exists(full_output_folder): - os.makedirs(full_output_folder) - - split = os.path.splitext(filename) - - if overwrite is not None and (overwrite == "true" or overwrite == "1"): - pass - else: - i = 1 - while os.path.exists(filepath): - if compare_image_hash(filepath, image): #compare hash to prevent saving of duplicates with same name, fix for #3465 - image_is_duplicate = True - break - filename = f"{split[0]} ({i}){split[1]}" - filepath = os.path.join(full_output_folder, filename) - i += 1 - - if not image_is_duplicate: - if image_save_function is not None: - image_save_function(image, post, filepath) - else: - with open(filepath, "wb") as f: - f.write(image.file.read()) - - return web.json_response({"name" : filename, "subfolder": subfolder, "type": image_upload_type}) - else: - return web.Response(status=400) - - @routes.post("/upload/image") - async def upload_image(request): - post = await request.post() - return image_upload(post) - - - @routes.post("/upload/mask") - async def upload_mask(request): - post = await request.post() - - def image_save_function(image, post, filepath): - original_ref = json.loads(post.get("original_ref")) - filename, output_dir = folder_paths.annotated_filepath(original_ref['filename']) - - if not filename: - return web.Response(status=400) - - # validation for security: prevent accessing arbitrary path - if filename[0] == '/' or '..' in filename: - return web.Response(status=400) - - if output_dir is None: - type = original_ref.get("type", "output") - output_dir = folder_paths.get_directory_by_type(type) - - if output_dir is None: - return web.Response(status=400) - - if original_ref.get("subfolder", "") != "": - full_output_dir = os.path.join(output_dir, original_ref["subfolder"]) - if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir: - return web.Response(status=403) - output_dir = full_output_dir - - file = os.path.join(output_dir, filename) - - if os.path.isfile(file): - with Image.open(file) as original_pil: - metadata = PngInfo() - if hasattr(original_pil,'text'): - for key in original_pil.text: - metadata.add_text(key, original_pil.text[key]) - original_pil = original_pil.convert('RGBA') - mask_pil = Image.open(image.file).convert('RGBA') - - # alpha copy - new_alpha = mask_pil.getchannel('A') - original_pil.putalpha(new_alpha) - original_pil.save(filepath, compress_level=4, pnginfo=metadata) - - return image_upload(post, image_save_function) - - @routes.get("/view") - async def view_image(request): - if "filename" in request.rel_url.query: - filename = request.rel_url.query["filename"] - filename, output_dir = folder_paths.annotated_filepath(filename) - - if not filename: - return web.Response(status=400) - - # validation for security: prevent accessing arbitrary path - if filename[0] == '/' or '..' in filename: - return web.Response(status=400) - - if output_dir is None: - type = request.rel_url.query.get("type", "output") - output_dir = folder_paths.get_directory_by_type(type) - - if output_dir is None: - return web.Response(status=400) - - if "subfolder" in request.rel_url.query: - full_output_dir = os.path.join(output_dir, request.rel_url.query["subfolder"]) - if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir: - return web.Response(status=403) - output_dir = full_output_dir - - filename = os.path.basename(filename) - file = os.path.join(output_dir, filename) - - if os.path.isfile(file): - if 'preview' in request.rel_url.query: - with Image.open(file) as img: - preview_info = request.rel_url.query['preview'].split(';') - image_format = preview_info[0] - if image_format not in ['webp', 'jpeg'] or 'a' in request.rel_url.query.get('channel', ''): - image_format = 'webp' - - quality = 90 - if preview_info[-1].isdigit(): - quality = int(preview_info[-1]) - - buffer = BytesIO() - if image_format in ['jpeg'] or request.rel_url.query.get('channel', '') == 'rgb': - img = img.convert("RGB") - img.save(buffer, format=image_format, quality=quality) - buffer.seek(0) - - return web.Response(body=buffer.read(), content_type=f'image/{image_format}', - headers={"Content-Disposition": f"filename=\"{filename}\""}) - - if 'channel' not in request.rel_url.query: - channel = 'rgba' - else: - channel = request.rel_url.query["channel"] - - if channel == 'rgb': - with Image.open(file) as img: - if img.mode == "RGBA": - r, g, b, a = img.split() - new_img = Image.merge('RGB', (r, g, b)) - else: - new_img = img.convert("RGB") - - buffer = BytesIO() - new_img.save(buffer, format='PNG') - buffer.seek(0) - - return web.Response(body=buffer.read(), content_type='image/png', - headers={"Content-Disposition": f"filename=\"{filename}\""}) - - elif channel == 'a': - with Image.open(file) as img: - if img.mode == "RGBA": - _, _, _, a = img.split() - else: - a = Image.new('L', img.size, 255) - - # alpha img - alpha_img = Image.new('RGBA', img.size) - alpha_img.putalpha(a) - alpha_buffer = BytesIO() - alpha_img.save(alpha_buffer, format='PNG') - alpha_buffer.seek(0) - - return web.Response(body=alpha_buffer.read(), content_type='image/png', - headers={"Content-Disposition": f"filename=\"{filename}\""}) - else: - # Get content type from mimetype, defaulting to 'application/octet-stream' - content_type = mimetypes.guess_type(filename)[0] or 'application/octet-stream' - - # For security, force certain mimetypes to download instead of display - if content_type in {'text/html', 'text/html-sandboxed', 'application/xhtml+xml', 'text/javascript', 'text/css'}: - content_type = 'application/octet-stream' # Forces download - - return web.FileResponse( - file, - headers={ - "Content-Disposition": f"filename=\"{filename}\"", - "Content-Type": content_type - } - ) - - return web.Response(status=404) - - @routes.get("/view_metadata/{folder_name}") - async def view_metadata(request): - folder_name = request.match_info.get("folder_name", None) - if folder_name is None: - return web.Response(status=404) - if not "filename" in request.rel_url.query: - return web.Response(status=404) - - filename = request.rel_url.query["filename"] - if not filename.endswith(".safetensors"): - return web.Response(status=404) - - safetensors_path = folder_paths.get_full_path(folder_name, filename) - if safetensors_path is None: - return web.Response(status=404) - out = comfy.utils.safetensors_header(safetensors_path, max_size=1024*1024) - if out is None: - return web.Response(status=404) - dt = json.loads(out) - if not "__metadata__" in dt: - return web.Response(status=404) - return web.json_response(dt["__metadata__"]) - - @routes.get("/system_stats") - async def system_stats(request): - device = comfy.model_management.get_torch_device() - device_name = comfy.model_management.get_torch_device_name(device) - cpu_device = comfy.model_management.torch.device("cpu") - ram_total = comfy.model_management.get_total_memory(cpu_device) - ram_free = comfy.model_management.get_free_memory(cpu_device) - vram_total, torch_vram_total = comfy.model_management.get_total_memory(device, torch_total_too=True) - vram_free, torch_vram_free = comfy.model_management.get_free_memory(device, torch_free_too=True) - required_frontend_version = FrontendManager.get_required_frontend_version() - - system_stats = { - "system": { - "os": os.name, - "ram_total": ram_total, - "ram_free": ram_free, - "comfyui_version": __version__, - "required_frontend_version": required_frontend_version, - "python_version": sys.version, - "pytorch_version": comfy.model_management.torch_version, - "embedded_python": os.path.split(os.path.split(sys.executable)[0])[1] == "python_embeded", - "argv": sys.argv - }, - "devices": [ - { - "name": device_name, - "type": device.type, - "index": device.index, - "vram_total": vram_total, - "vram_free": vram_free, - "torch_vram_total": torch_vram_total, - "torch_vram_free": torch_vram_free, - } - ] - } - return web.json_response(system_stats) - - @routes.get("/features") - async def get_features(request): - return web.json_response(feature_flags.get_server_features()) - - @routes.get("/prompt") - async def get_prompt(request): - return web.json_response(self.get_queue_info()) - - def node_info(node_class): - obj_class = nodes.NODE_CLASS_MAPPINGS[node_class] - if issubclass(obj_class, _ComfyNodeInternal): - return obj_class.GET_NODE_INFO_V1() - info = {} - info['input'] = obj_class.INPUT_TYPES() - info['input_order'] = {key: list(value.keys()) for (key, value) in obj_class.INPUT_TYPES().items()} - info['output'] = obj_class.RETURN_TYPES - info['output_is_list'] = obj_class.OUTPUT_IS_LIST if hasattr(obj_class, 'OUTPUT_IS_LIST') else [False] * len(obj_class.RETURN_TYPES) - info['output_name'] = obj_class.RETURN_NAMES if hasattr(obj_class, 'RETURN_NAMES') else info['output'] - info['name'] = node_class - info['display_name'] = nodes.NODE_DISPLAY_NAME_MAPPINGS[node_class] if node_class in nodes.NODE_DISPLAY_NAME_MAPPINGS.keys() else node_class - info['description'] = obj_class.DESCRIPTION if hasattr(obj_class,'DESCRIPTION') else '' - info['python_module'] = getattr(obj_class, "RELATIVE_PYTHON_MODULE", "nodes") - info['category'] = 'sd' - if hasattr(obj_class, 'OUTPUT_NODE') and obj_class.OUTPUT_NODE == True: - info['output_node'] = True - else: - info['output_node'] = False - - if hasattr(obj_class, 'CATEGORY'): - info['category'] = obj_class.CATEGORY - - if hasattr(obj_class, 'OUTPUT_TOOLTIPS'): - info['output_tooltips'] = obj_class.OUTPUT_TOOLTIPS - - if getattr(obj_class, "DEPRECATED", False): - info['deprecated'] = True - if getattr(obj_class, "EXPERIMENTAL", False): - info['experimental'] = True - - if hasattr(obj_class, 'API_NODE'): - info['api_node'] = obj_class.API_NODE - return info - - @routes.get("/object_info") - async def get_object_info(request): - with folder_paths.cache_helper: - out = {} - for x in nodes.NODE_CLASS_MAPPINGS: - try: - out[x] = node_info(x) - except Exception: - logging.error(f"[ERROR] An error occurred while retrieving information for the '{x}' node.") - logging.error(traceback.format_exc()) - return web.json_response(out) - - @routes.get("/object_info/{node_class}") - async def get_object_info_node(request): - node_class = request.match_info.get("node_class", None) - out = {} - if (node_class is not None) and (node_class in nodes.NODE_CLASS_MAPPINGS): - out[node_class] = node_info(node_class) - return web.json_response(out) - - @routes.get("/history") - async def get_history(request): - max_items = request.rel_url.query.get("max_items", None) - if max_items is not None: - max_items = int(max_items) - - offset = request.rel_url.query.get("offset", None) - if offset is not None: - offset = int(offset) - else: - offset = -1 - - return web.json_response(self.prompt_queue.get_history(max_items=max_items, offset=offset)) - - @routes.get("/history/{prompt_id}") - async def get_history_prompt_id(request): - prompt_id = request.match_info.get("prompt_id", None) - return web.json_response(self.prompt_queue.get_history(prompt_id=prompt_id)) - - @routes.get("/queue") - async def get_queue(request): - queue_info = {} - current_queue = self.prompt_queue.get_current_queue_volatile() - queue_info['queue_running'] = current_queue[0] - queue_info['queue_pending'] = current_queue[1] - return web.json_response(queue_info) - - @routes.post("/prompt") - async def post_prompt(request): - logging.info("got prompt") - json_data = await request.json() - json_data = self.trigger_on_prompt(json_data) - - if "number" in json_data: - number = float(json_data['number']) - else: - number = self.number - if "front" in json_data: - if json_data['front']: - number = -number - - self.number += 1 - - if "prompt" in json_data: - prompt = json_data["prompt"] - prompt_id = str(json_data.get("prompt_id", uuid.uuid4())) - - partial_execution_targets = None - if "partial_execution_targets" in json_data: - partial_execution_targets = json_data["partial_execution_targets"] - - valid = await execution.validate_prompt(prompt_id, prompt, partial_execution_targets) - extra_data = {} - if "extra_data" in json_data: - extra_data = json_data["extra_data"] - - if "client_id" in json_data: - extra_data["client_id"] = json_data["client_id"] - if valid[0]: - outputs_to_execute = valid[2] - self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute)) - response = {"prompt_id": prompt_id, "number": number, "node_errors": valid[3]} - return web.json_response(response) - else: - logging.warning("invalid prompt: {}".format(valid[1])) - return web.json_response({"error": valid[1], "node_errors": valid[3]}, status=400) - else: - error = { - "type": "no_prompt", - "message": "No prompt provided", - "details": "No prompt provided", - "extra_info": {} - } - return web.json_response({"error": error, "node_errors": {}}, status=400) - - @routes.post("/queue") - async def post_queue(request): - json_data = await request.json() - if "clear" in json_data: - if json_data["clear"]: - self.prompt_queue.wipe_queue() - if "delete" in json_data: - to_delete = json_data['delete'] - for id_to_delete in to_delete: - delete_func = lambda a: a[1] == id_to_delete - self.prompt_queue.delete_queue_item(delete_func) - - return web.Response(status=200) - - @routes.post("/interrupt") - async def post_interrupt(request): - try: - json_data = await request.json() - except json.JSONDecodeError: - json_data = {} - - # Check if a specific prompt_id was provided for targeted interruption - prompt_id = json_data.get('prompt_id') - if prompt_id: - currently_running, _ = self.prompt_queue.get_current_queue() - - # Check if the prompt_id matches any currently running prompt - should_interrupt = False - for item in currently_running: - # item structure: (number, prompt_id, prompt, extra_data, outputs_to_execute) - if item[1] == prompt_id: - logging.info(f"Interrupting prompt {prompt_id}") - should_interrupt = True - break - - if should_interrupt: - nodes.interrupt_processing() - else: - logging.info(f"Prompt {prompt_id} is not currently running, skipping interrupt") - else: - # No prompt_id provided, do a global interrupt - logging.info("Global interrupt (no prompt_id specified)") - nodes.interrupt_processing() - - return web.Response(status=200) - - @routes.post("/free") - async def post_free(request): - json_data = await request.json() - unload_models = json_data.get("unload_models", False) - free_memory = json_data.get("free_memory", False) - if unload_models: - self.prompt_queue.set_flag("unload_models", unload_models) - if free_memory: - self.prompt_queue.set_flag("free_memory", free_memory) - return web.Response(status=200) - - @routes.post("/history") - async def post_history(request): - json_data = await request.json() - if "clear" in json_data: - if json_data["clear"]: - self.prompt_queue.wipe_history() - if "delete" in json_data: - to_delete = json_data['delete'] - for id_to_delete in to_delete: - self.prompt_queue.delete_history_item(id_to_delete) - - return web.Response(status=200) - - async def setup(self): - timeout = aiohttp.ClientTimeout(total=None) # no timeout - self.client_session = aiohttp.ClientSession(timeout=timeout) - - def add_routes(self): - self.user_manager.add_routes(self.routes) - self.model_file_manager.add_routes(self.routes) - self.custom_node_manager.add_routes(self.routes, self.app, nodes.LOADED_MODULE_DIRS.items()) - self.app.add_subapp('/internal', self.internal_routes.get_app()) - - # Prefix every route with /api for easier matching for delegation. - # This is very useful for frontend dev server, which need to forward - # everything except serving of static files. - # Currently both the old endpoints without prefix and new endpoints with - # prefix are supported. - api_routes = web.RouteTableDef() - for route in self.routes: - # Custom nodes might add extra static routes. Only process non-static - # routes to add /api prefix. - if isinstance(route, web.RouteDef): - api_routes.route(route.method, "/api" + route.path)(route.handler, **route.kwargs) - self.app.add_routes(api_routes) - self.app.add_routes(self.routes) - - # Add routes from web extensions. - for name, dir in nodes.EXTENSION_WEB_DIRS.items(): - self.app.add_routes([web.static('/extensions/' + name, dir)]) - - workflow_templates_path = FrontendManager.templates_path() - if workflow_templates_path: - self.app.add_routes([ - web.static('/templates', workflow_templates_path) - ]) - - # Serve embedded documentation from the package - embedded_docs_path = FrontendManager.embedded_docs_path() - if embedded_docs_path: - self.app.add_routes([ - web.static('/docs', embedded_docs_path) - ]) - - self.app.add_routes([ - web.static('/', self.web_root), - ]) - - def get_queue_info(self): - prompt_info = {} - exec_info = {} - exec_info['queue_remaining'] = self.prompt_queue.get_tasks_remaining() - prompt_info['exec_info'] = exec_info - return prompt_info - - async def send(self, event, data, sid=None): - if event == BinaryEventTypes.UNENCODED_PREVIEW_IMAGE: - await self.send_image(data, sid=sid) - elif event == BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA: - # data is (preview_image, metadata) - preview_image, metadata = data - await self.send_image_with_metadata(preview_image, metadata, sid=sid) - elif isinstance(data, (bytes, bytearray)): - await self.send_bytes(event, data, sid) - else: - await self.send_json(event, data, sid) - - def encode_bytes(self, event, data): - if not isinstance(event, int): - raise RuntimeError(f"Binary event types must be integers, got {event}") - - packed = struct.pack(">I", event) - message = bytearray(packed) - message.extend(data) - return message - - async def send_image(self, image_data, sid=None): - image_type = image_data[0] - image = image_data[1] - max_size = image_data[2] - if max_size is not None: - if hasattr(Image, 'Resampling'): - resampling = Image.Resampling.BILINEAR - else: - resampling = Image.Resampling.LANCZOS - - image = ImageOps.contain(image, (max_size, max_size), resampling) - type_num = 1 - if image_type == "JPEG": - type_num = 1 - elif image_type == "PNG": - type_num = 2 - - bytesIO = BytesIO() - header = struct.pack(">I", type_num) - bytesIO.write(header) - image.save(bytesIO, format=image_type, quality=95, compress_level=1) - preview_bytes = bytesIO.getvalue() - await self.send_bytes(BinaryEventTypes.PREVIEW_IMAGE, preview_bytes, sid=sid) - - async def send_image_with_metadata(self, image_data, metadata=None, sid=None): - image_type = image_data[0] - image = image_data[1] - max_size = image_data[2] - if max_size is not None: - if hasattr(Image, 'Resampling'): - resampling = Image.Resampling.BILINEAR - else: - resampling = Image.Resampling.LANCZOS - - image = ImageOps.contain(image, (max_size, max_size), resampling) - - mimetype = "image/png" if image_type == "PNG" else "image/jpeg" - - # Prepare metadata - if metadata is None: - metadata = {} - metadata["image_type"] = mimetype - - # Serialize metadata as JSON - import json - metadata_json = json.dumps(metadata).encode('utf-8') - metadata_length = len(metadata_json) - - # Prepare image data - bytesIO = BytesIO() - image.save(bytesIO, format=image_type, quality=95, compress_level=1) - image_bytes = bytesIO.getvalue() - - # Combine metadata and image - combined_data = bytearray() - combined_data.extend(struct.pack(">I", metadata_length)) - combined_data.extend(metadata_json) - combined_data.extend(image_bytes) - - await self.send_bytes(BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA, combined_data, sid=sid) - - async def send_bytes(self, event, data, sid=None): - message = self.encode_bytes(event, data) - - if sid is None: - sockets = list(self.sockets.values()) - for ws in sockets: - await send_socket_catch_exception(ws.send_bytes, message) - elif sid in self.sockets: - await send_socket_catch_exception(self.sockets[sid].send_bytes, message) - - async def send_json(self, event, data, sid=None): - message = {"type": event, "data": data} - - if sid is None: - sockets = list(self.sockets.values()) - for ws in sockets: - await send_socket_catch_exception(ws.send_json, message) - elif sid in self.sockets: - await send_socket_catch_exception(self.sockets[sid].send_json, message) - - def send_sync(self, event, data, sid=None): - self.loop.call_soon_threadsafe( - self.messages.put_nowait, (event, data, sid)) - - def queue_updated(self): - self.send_sync("status", { "status": self.get_queue_info() }) - - async def publish_loop(self): - while True: - msg = await self.messages.get() - await self.send(*msg) - - async def start(self, address, port, verbose=True, call_on_start=None): - await self.start_multi_address([(address, port)], call_on_start=call_on_start) - - async def start_multi_address(self, addresses, call_on_start=None, verbose=True): - runner = web.AppRunner(self.app, access_log=None) - await runner.setup() - ssl_ctx = None - scheme = "http" - if args.tls_keyfile and args.tls_certfile: - ssl_ctx = ssl.SSLContext(protocol=ssl.PROTOCOL_TLS_SERVER, verify_mode=ssl.CERT_NONE) - ssl_ctx.load_cert_chain(certfile=args.tls_certfile, - keyfile=args.tls_keyfile) - scheme = "https" - - if verbose: - logging.info("Starting server\n") - for addr in addresses: - address = addr[0] - port = addr[1] - site = web.TCPSite(runner, address, port, ssl_context=ssl_ctx) - await site.start() - - if not hasattr(self, 'address'): - self.address = address #TODO: remove this - self.port = port - - if ':' in address: - address_print = "[{}]".format(address) - else: - address_print = address - - if verbose: - logging.info("To see the GUI go to: {}://{}:{}".format(scheme, address_print, port)) - - if call_on_start is not None: - call_on_start(scheme, self.address, self.port) - - def add_on_prompt_handler(self, handler): - self.on_prompt_handlers.append(handler) - - def trigger_on_prompt(self, json_data): - for handler in self.on_prompt_handlers: - try: - json_data = handler(json_data) - except Exception: - logging.warning("[ERROR] An error occurred during the on_prompt_handler processing") - logging.warning(traceback.format_exc()) - - return json_data - - def send_progress_text( - self, text: Union[bytes, bytearray, str], node_id: str, sid=None - ): - if isinstance(text, str): - text = text.encode("utf-8") - node_id_bytes = str(node_id).encode("utf-8") - - # Pack the node_id length as a 4-byte unsigned integer, followed by the node_id bytes - message = struct.pack(">I", len(node_id_bytes)) + node_id_bytes + text - - self.send_sync(BinaryEventTypes.TEXT, message, sid) -########################## Comfy Core main.py ########################## -import comfy.options -comfy.options.enable_args_parsing() - -import os -import importlib.util -import folder_paths -import time -from comfy.cli_args import args -from app.logger import setup_logger -import itertools -import utils.extra_config -import logging -import sys -from comfy_execution.progress import get_progress_state -from comfy_execution.utils import get_executing_context -from comfy_api import feature_flags - -if __name__ == "__main__": - #NOTE: These do not do anything on core ComfyUI, they are for custom nodes. - os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1' - os.environ['DO_NOT_TRACK'] = '1' - -setup_logger(log_level=args.verbose, use_stdout=args.log_stdout) - -def apply_custom_paths(): - # extra model paths - extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml") - if os.path.isfile(extra_model_paths_config_path): - utils.extra_config.load_extra_path_config(extra_model_paths_config_path) - - if args.extra_model_paths_config: - for config_path in itertools.chain(*args.extra_model_paths_config): - utils.extra_config.load_extra_path_config(config_path) - - # --output-directory, --input-directory, --user-directory - if args.output_directory: - output_dir = os.path.abspath(args.output_directory) - logging.info(f"Setting output directory to: {output_dir}") - folder_paths.set_output_directory(output_dir) - - # These are the default folders that checkpoints, clip and vae models will be saved to when using CheckpointSave, etc.. nodes - folder_paths.add_model_folder_path("checkpoints", os.path.join(folder_paths.get_output_directory(), "checkpoints")) - folder_paths.add_model_folder_path("clip", os.path.join(folder_paths.get_output_directory(), "clip")) - folder_paths.add_model_folder_path("vae", os.path.join(folder_paths.get_output_directory(), "vae")) - folder_paths.add_model_folder_path("diffusion_models", - os.path.join(folder_paths.get_output_directory(), "diffusion_models")) - folder_paths.add_model_folder_path("loras", os.path.join(folder_paths.get_output_directory(), "loras")) - - if args.input_directory: - input_dir = os.path.abspath(args.input_directory) - logging.info(f"Setting input directory to: {input_dir}") - folder_paths.set_input_directory(input_dir) - - if args.user_directory: - user_dir = os.path.abspath(args.user_directory) - logging.info(f"Setting user directory to: {user_dir}") - folder_paths.set_user_directory(user_dir) - - -def execute_prestartup_script(): - if args.disable_all_custom_nodes and len(args.whitelist_custom_nodes) == 0: - return - - def execute_script(script_path): - module_name = os.path.splitext(script_path)[0] - try: - spec = importlib.util.spec_from_file_location(module_name, script_path) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return True - except Exception as e: - logging.error(f"Failed to execute startup-script: {script_path} / {e}") - return False - - node_paths = folder_paths.get_folder_paths("custom_nodes") - for custom_node_path in node_paths: - possible_modules = os.listdir(custom_node_path) - node_prestartup_times = [] - - for possible_module in possible_modules: - module_path = os.path.join(custom_node_path, possible_module) - if os.path.isfile(module_path) or module_path.endswith(".disabled") or module_path == "__pycache__": - continue - - script_path = os.path.join(module_path, "prestartup_script.py") - if os.path.exists(script_path): - if args.disable_all_custom_nodes and possible_module not in args.whitelist_custom_nodes: - logging.info(f"Prestartup Skipping {possible_module} due to disable_all_custom_nodes and whitelist_custom_nodes") - continue - time_before = time.perf_counter() - success = execute_script(script_path) - node_prestartup_times.append((time.perf_counter() - time_before, module_path, success)) - if len(node_prestartup_times) > 0: - logging.info("\nPrestartup times for custom nodes:") - for n in sorted(node_prestartup_times): - if n[2]: - import_message = "" - else: - import_message = " (PRESTARTUP FAILED)" - logging.info("{:6.1f} seconds{}: {}".format(n[0], import_message, n[1])) - logging.info("") - -apply_custom_paths() -execute_prestartup_script() - - -# Main code -import asyncio -import shutil -import threading -import gc - - -if os.name == "nt": - os.environ['MIMALLOC_PURGE_DELAY'] = '0' - -if __name__ == "__main__": - if args.default_device is not None: - default_dev = args.default_device - devices = list(range(32)) - devices.remove(default_dev) - devices.insert(0, default_dev) - devices = ','.join(map(str, devices)) - os.environ['CUDA_VISIBLE_DEVICES'] = str(devices) - os.environ['HIP_VISIBLE_DEVICES'] = str(devices) - - if args.cuda_device is not None: - os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device) - os.environ['HIP_VISIBLE_DEVICES'] = str(args.cuda_device) - logging.info("Set cuda device to: {}".format(args.cuda_device)) - - if args.oneapi_device_selector is not None: - os.environ['ONEAPI_DEVICE_SELECTOR'] = args.oneapi_device_selector - logging.info("Set oneapi device selector to: {}".format(args.oneapi_device_selector)) - - if args.deterministic: - if 'CUBLAS_WORKSPACE_CONFIG' not in os.environ: - os.environ['CUBLAS_WORKSPACE_CONFIG'] = ":4096:8" - - import cuda_malloc - -if 'torch' in sys.modules: - logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.") - -import comfy.utils - -import execution -import server -from protocol import BinaryEventTypes -import nodes -import comfy.model_management -import comfyui_version -import app.logger -import hook_breaker_ac10a0 - -def cuda_malloc_warning(): - device = comfy.model_management.get_torch_device() - device_name = comfy.model_management.get_torch_device_name(device) - cuda_malloc_warning = False - if "cudaMallocAsync" in device_name: - for b in cuda_malloc.blacklist: - if b in device_name: - cuda_malloc_warning = True - if cuda_malloc_warning: - logging.warning("\nWARNING: this card most likely does not support cuda-malloc, if you get \"CUDA error\" please run ComfyUI with: --disable-cuda-malloc\n") - - -def prompt_worker(q, server_instance): - current_time: float = 0.0 - cache_type = execution.CacheType.CLASSIC - if args.cache_lru > 0: - cache_type = execution.CacheType.LRU - elif args.cache_none: - cache_type = execution.CacheType.DEPENDENCY_AWARE - - e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_size=args.cache_lru) - last_gc_collect = 0 - need_gc = False - gc_collect_interval = 10.0 - - while True: - timeout = 1000.0 - if need_gc: - timeout = max(gc_collect_interval - (current_time - last_gc_collect), 0.0) - - queue_item = q.get(timeout=timeout) - if queue_item is not None: - item, item_id = queue_item - execution_start_time = time.perf_counter() - prompt_id = item[1] - server_instance.last_prompt_id = prompt_id - - e.execute(item[2], prompt_id, item[3], item[4]) - need_gc = True - q.task_done(item_id, - e.history_result, - status=execution.PromptQueue.ExecutionStatus( - status_str='success' if e.success else 'error', - completed=e.success, - messages=e.status_messages)) - if server_instance.client_id is not None: - server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id}, server_instance.client_id) - - current_time = time.perf_counter() - execution_time = current_time - execution_start_time - - # Log Time in a more readable way after 10 minutes - if execution_time > 600: - execution_time = time.strftime("%H:%M:%S", time.gmtime(execution_time)) - logging.info(f"Prompt executed in {execution_time}") - else: - logging.info("Prompt executed in {:.2f} seconds".format(execution_time)) - - flags = q.get_flags() - free_memory = flags.get("free_memory", False) - - if flags.get("unload_models", free_memory): - comfy.model_management.unload_all_models() - need_gc = True - last_gc_collect = 0 - - if free_memory: - e.reset() - need_gc = True - last_gc_collect = 0 - - if need_gc: - current_time = time.perf_counter() - if (current_time - last_gc_collect) > gc_collect_interval: - gc.collect() - comfy.model_management.soft_empty_cache() - last_gc_collect = current_time - need_gc = False - hook_breaker_ac10a0.restore_functions() - - -async def run(server_instance, address='', port=8188, verbose=True, call_on_start=None): - addresses = [] - for addr in address.split(","): - addresses.append((addr, port)) - await asyncio.gather( - server_instance.start_multi_address(addresses, call_on_start, verbose), server_instance.publish_loop() - ) - -def hijack_progress(server_instance): - def hook(value, total, preview_image, prompt_id=None, node_id=None): - executing_context = get_executing_context() - if prompt_id is None and executing_context is not None: - prompt_id = executing_context.prompt_id - if node_id is None and executing_context is not None: - node_id = executing_context.node_id - comfy.model_management.throw_exception_if_processing_interrupted() - if prompt_id is None: - prompt_id = server_instance.last_prompt_id - if node_id is None: - node_id = server_instance.last_node_id - progress = {"value": value, "max": total, "prompt_id": prompt_id, "node": node_id} - get_progress_state().update_progress(node_id, value, total, preview_image) - - server_instance.send_sync("progress", progress, server_instance.client_id) - if preview_image is not None: - # Only send old method if client doesn't support preview metadata - if not feature_flags.supports_feature( - server_instance.sockets_metadata, - server_instance.client_id, - "supports_preview_metadata", - ): - server_instance.send_sync( - BinaryEventTypes.UNENCODED_PREVIEW_IMAGE, - preview_image, - server_instance.client_id, - ) - - comfy.utils.set_progress_bar_global_hook(hook) - - -def cleanup_temp(): - temp_dir = folder_paths.get_temp_directory() - if os.path.exists(temp_dir): - shutil.rmtree(temp_dir, ignore_errors=True) - - -def setup_database(): - try: - from app.database.db import init_db, dependencies_available - if dependencies_available(): - init_db() - except Exception as e: - logging.error(f"Failed to initialize database. Please ensure you have installed the latest requirements. If the error persists, please report this as in future the database will be required: {e}") - - -def start_comfyui(asyncio_loop=None): - """ - Starts the ComfyUI server using the provided asyncio event loop or creates a new one. - Returns the event loop, server instance, and a function to start the server asynchronously. - """ - if args.temp_directory: - temp_dir = os.path.join(os.path.abspath(args.temp_directory), "temp") - logging.info(f"Setting temp directory to: {temp_dir}") - folder_paths.set_temp_directory(temp_dir) - cleanup_temp() - - if args.windows_standalone_build: - try: - import new_updater - new_updater.update_windows_updater() - except: - pass - - if not asyncio_loop: - asyncio_loop = asyncio.new_event_loop() - asyncio.set_event_loop(asyncio_loop) - prompt_server = server.PromptServer(asyncio_loop) - - hook_breaker_ac10a0.save_functions() - asyncio_loop.run_until_complete(nodes.init_extra_nodes( - init_custom_nodes=(not args.disable_all_custom_nodes) or len(args.whitelist_custom_nodes) > 0, - init_api_nodes=not args.disable_api_nodes - )) - hook_breaker_ac10a0.restore_functions() - - cuda_malloc_warning() - setup_database() - - prompt_server.add_routes() - hijack_progress(prompt_server) - - threading.Thread(target=prompt_worker, daemon=True, args=(prompt_server.prompt_queue, prompt_server,)).start() - - if args.quick_test_for_ci: - exit(0) - - os.makedirs(folder_paths.get_temp_directory(), exist_ok=True) - call_on_start = None - if args.auto_launch: - def startup_server(scheme, address, port): - import webbrowser - if os.name == 'nt' and address == '0.0.0.0': - address = '127.0.0.1' - if ':' in address: - address = "[{}]".format(address) - webbrowser.open(f"{scheme}://{address}:{port}") - call_on_start = startup_server - - async def start_all(): - await prompt_server.setup() - await run(prompt_server, address=args.listen, port=args.port, verbose=not args.dont_print_server, call_on_start=call_on_start) - - # Returning these so that other code can integrate with the ComfyUI loop and server - return asyncio_loop, prompt_server, start_all - - -if __name__ == "__main__": - # Running directly, just start ComfyUI. - logging.info("Python version: {}".format(sys.version)) - logging.info("ComfyUI version: {}".format(comfyui_version.__version__)) - - if sys.version_info.major == 3 and sys.version_info.minor < 10: - logging.warning("WARNING: You are using a python version older than 3.10, please upgrade to a newer one. 3.12 and above is recommended.") - - event_loop, _, start_all_func = start_comfyui() - try: - x = start_all_func() - app.logger.print_startup_warnings() - event_loop.run_until_complete(x) - except KeyboardInterrupt: - logging.info("\nStopped server") - - cleanup_temp() -###################### Comfy Core execution.py ###################### -import copy -import heapq -import inspect -import logging -import sys -import threading -import time -import traceback -from enum import Enum -from typing import List, Literal, NamedTuple, Optional, Union -import asyncio - -import torch - -import comfy.model_management -import nodes -from comfy_execution.caching import ( - BasicCache, - CacheKeySetID, - CacheKeySetInputSignature, - DependencyAwareCache, - HierarchicalCache, - LRUCache, -) -from comfy_execution.graph import ( - DynamicPrompt, - ExecutionBlocker, - ExecutionList, - get_input_info, -) -from comfy_execution.graph_utils import GraphBuilder, is_link -from comfy_execution.validation import validate_node_input -from comfy_execution.progress import get_progress_state, reset_progress_state, add_progress_handler, WebUIProgressHandler -from comfy_execution.utils import CurrentNodeContext -from comfy_api.internal import _ComfyNodeInternal, _NodeOutputInternal, first_real_override, is_class, make_locked_method_func -from comfy_api.latest import io - - -class ExecutionResult(Enum): - SUCCESS = 0 - FAILURE = 1 - PENDING = 2 - -class DuplicateNodeError(Exception): - pass - -class IsChangedCache: - def __init__(self, prompt_id: str, dynprompt: DynamicPrompt, outputs_cache: BasicCache): - self.prompt_id = prompt_id - self.dynprompt = dynprompt - self.outputs_cache = outputs_cache - self.is_changed = {} - - async def get(self, node_id): - if node_id in self.is_changed: - return self.is_changed[node_id] - - node = self.dynprompt.get_node(node_id) - class_type = node["class_type"] - class_def = nodes.NODE_CLASS_MAPPINGS[class_type] - has_is_changed = False - is_changed_name = None - if issubclass(class_def, _ComfyNodeInternal) and first_real_override(class_def, "fingerprint_inputs") is not None: - has_is_changed = True - is_changed_name = "fingerprint_inputs" - elif hasattr(class_def, "IS_CHANGED"): - has_is_changed = True - is_changed_name = "IS_CHANGED" - if not has_is_changed: - self.is_changed[node_id] = False - return self.is_changed[node_id] - - if "is_changed" in node: - self.is_changed[node_id] = node["is_changed"] - return self.is_changed[node_id] - - # Intentionally do not use cached outputs here. We only want constants in IS_CHANGED - input_data_all, _, hidden_inputs = get_input_data(node["inputs"], class_def, node_id, None) - try: - is_changed = await _async_map_node_over_list(self.prompt_id, node_id, class_def, input_data_all, is_changed_name) - is_changed = await resolve_map_node_over_list_results(is_changed) - node["is_changed"] = [None if isinstance(x, ExecutionBlocker) else x for x in is_changed] - except Exception as e: - logging.warning("WARNING: {}".format(e)) - node["is_changed"] = float("NaN") - finally: - self.is_changed[node_id] = node["is_changed"] - return self.is_changed[node_id] - - -class CacheType(Enum): - CLASSIC = 0 - LRU = 1 - DEPENDENCY_AWARE = 2 - - -class CacheSet: - def __init__(self, cache_type=None, cache_size=None): - if cache_type == CacheType.DEPENDENCY_AWARE: - self.init_dependency_aware_cache() - logging.info("Disabling intermediate node cache.") - elif cache_type == CacheType.LRU: - if cache_size is None: - cache_size = 0 - self.init_lru_cache(cache_size) - logging.info("Using LRU cache") - else: - self.init_classic_cache() - - self.all = [self.outputs, self.ui, self.objects] - - # Performs like the old cache -- dump data ASAP - def init_classic_cache(self): - self.outputs = HierarchicalCache(CacheKeySetInputSignature) - self.ui = HierarchicalCache(CacheKeySetInputSignature) - self.objects = HierarchicalCache(CacheKeySetID) - - def init_lru_cache(self, cache_size): - self.outputs = LRUCache(CacheKeySetInputSignature, max_size=cache_size) - self.ui = LRUCache(CacheKeySetInputSignature, max_size=cache_size) - self.objects = HierarchicalCache(CacheKeySetID) - - # only hold cached items while the decendents have not executed - def init_dependency_aware_cache(self): - self.outputs = DependencyAwareCache(CacheKeySetInputSignature) - self.ui = DependencyAwareCache(CacheKeySetInputSignature) - self.objects = DependencyAwareCache(CacheKeySetID) - - def recursive_debug_dump(self): - result = { - "outputs": self.outputs.recursive_debug_dump(), - "ui": self.ui.recursive_debug_dump(), - } - return result - -SENSITIVE_EXTRA_DATA_KEYS = ("auth_token_comfy_org", "api_key_comfy_org") - -def get_input_data(inputs, class_def, unique_id, outputs=None, dynprompt=None, extra_data={}): - is_v3 = issubclass(class_def, _ComfyNodeInternal) - if is_v3: - valid_inputs, schema = class_def.INPUT_TYPES(include_hidden=False, return_schema=True) - else: - valid_inputs = class_def.INPUT_TYPES() - input_data_all = {} - missing_keys = {} - hidden_inputs_v3 = {} - for x in inputs: - input_data = inputs[x] - _, input_category, input_info = get_input_info(class_def, x, valid_inputs) - def mark_missing(): - missing_keys[x] = True - input_data_all[x] = (None,) - if is_link(input_data) and (not input_info or not input_info.get("rawLink", False)): - input_unique_id = input_data[0] - output_index = input_data[1] - if outputs is None: - mark_missing() - continue # This might be a lazily-evaluated input - cached_output = outputs.get(input_unique_id) - if cached_output is None: - mark_missing() - continue - if output_index >= len(cached_output): - mark_missing() - continue - obj = cached_output[output_index] - input_data_all[x] = obj - elif input_category is not None: - input_data_all[x] = [input_data] - - if is_v3: - if schema.hidden: - if io.Hidden.prompt in schema.hidden: - hidden_inputs_v3[io.Hidden.prompt] = dynprompt.get_original_prompt() if dynprompt is not None else {} - if io.Hidden.dynprompt in schema.hidden: - hidden_inputs_v3[io.Hidden.dynprompt] = dynprompt - if io.Hidden.extra_pnginfo in schema.hidden: - hidden_inputs_v3[io.Hidden.extra_pnginfo] = extra_data.get('extra_pnginfo', None) - if io.Hidden.unique_id in schema.hidden: - hidden_inputs_v3[io.Hidden.unique_id] = unique_id - if io.Hidden.auth_token_comfy_org in schema.hidden: - hidden_inputs_v3[io.Hidden.auth_token_comfy_org] = extra_data.get("auth_token_comfy_org", None) - if io.Hidden.api_key_comfy_org in schema.hidden: - hidden_inputs_v3[io.Hidden.api_key_comfy_org] = extra_data.get("api_key_comfy_org", None) - else: - if "hidden" in valid_inputs: - h = valid_inputs["hidden"] - for x in h: - if h[x] == "PROMPT": - input_data_all[x] = [dynprompt.get_original_prompt() if dynprompt is not None else {}] - if h[x] == "DYNPROMPT": - input_data_all[x] = [dynprompt] - if h[x] == "EXTRA_PNGINFO": - input_data_all[x] = [extra_data.get('extra_pnginfo', None)] - if h[x] == "UNIQUE_ID": - input_data_all[x] = [unique_id] - if h[x] == "AUTH_TOKEN_COMFY_ORG": - input_data_all[x] = [extra_data.get("auth_token_comfy_org", None)] - if h[x] == "API_KEY_COMFY_ORG": - input_data_all[x] = [extra_data.get("api_key_comfy_org", None)] - return input_data_all, missing_keys, hidden_inputs_v3 - -map_node_over_list = None #Don't hook this please - -async def resolve_map_node_over_list_results(results): - remaining = [x for x in results if isinstance(x, asyncio.Task) and not x.done()] - if len(remaining) == 0: - return [x.result() if isinstance(x, asyncio.Task) else x for x in results] - else: - done, pending = await asyncio.wait(remaining) - for task in done: - exc = task.exception() - if exc is not None: - raise exc - return [x.result() if isinstance(x, asyncio.Task) else x for x in results] - -async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, func, allow_interrupt=False, execution_block_cb=None, pre_execute_cb=None, hidden_inputs=None): - # check if node wants the lists - input_is_list = getattr(obj, "INPUT_IS_LIST", False) - - if len(input_data_all) == 0: - max_len_input = 0 - else: - max_len_input = max(len(x) for x in input_data_all.values()) - - # get a slice of inputs, repeat last input when list isn't long enough - def slice_dict(d, i): - return {k: v[i if len(v) > i else -1] for k, v in d.items()} - - results = [] - async def process_inputs(inputs, index=None, input_is_list=False): - if allow_interrupt: - nodes.before_node_execution() - execution_block = None - for k, v in inputs.items(): - if input_is_list: - for e in v: - if isinstance(e, ExecutionBlocker): - v = e - break - if isinstance(v, ExecutionBlocker): - execution_block = execution_block_cb(v) if execution_block_cb else v - break - if execution_block is None: - if pre_execute_cb is not None and index is not None: - pre_execute_cb(index) - # V3 - if isinstance(obj, _ComfyNodeInternal) or (is_class(obj) and issubclass(obj, _ComfyNodeInternal)): - # if is just a class, then assign no resources or state, just create clone - if is_class(obj): - type_obj = obj - obj.VALIDATE_CLASS() - class_clone = obj.PREPARE_CLASS_CLONE(hidden_inputs) - # otherwise, use class instance to populate/reuse some fields - else: - type_obj = type(obj) - type_obj.VALIDATE_CLASS() - class_clone = type_obj.PREPARE_CLASS_CLONE(hidden_inputs) - f = make_locked_method_func(type_obj, func, class_clone) - # V1 - else: - f = getattr(obj, func) - if inspect.iscoroutinefunction(f): - async def async_wrapper(f, prompt_id, unique_id, list_index, args): - with CurrentNodeContext(prompt_id, unique_id, list_index): - return await f(**args) - task = asyncio.create_task(async_wrapper(f, prompt_id, unique_id, index, args=inputs)) - # Give the task a chance to execute without yielding - await asyncio.sleep(0) - if task.done(): - result = task.result() - results.append(result) - else: - results.append(task) - else: - with CurrentNodeContext(prompt_id, unique_id, index): - result = f(**inputs) - results.append(result) - else: - results.append(execution_block) - - if input_is_list: - await process_inputs(input_data_all, 0, input_is_list=input_is_list) - elif max_len_input == 0: - await process_inputs({}) - else: - for i in range(max_len_input): - input_dict = slice_dict(input_data_all, i) - await process_inputs(input_dict, i) - return results - - -def merge_result_data(results, obj): - # check which outputs need concatenating - output = [] - output_is_list = [False] * len(results[0]) - if hasattr(obj, "OUTPUT_IS_LIST"): - output_is_list = obj.OUTPUT_IS_LIST - - # merge node execution results - for i, is_list in zip(range(len(results[0])), output_is_list): - if is_list: - value = [] - for o in results: - if isinstance(o[i], ExecutionBlocker): - value.append(o[i]) - else: - value.extend(o[i]) - output.append(value) - else: - output.append([o[i] for o in results]) - return output - -async def get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=None, pre_execute_cb=None, hidden_inputs=None): - return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, hidden_inputs=hidden_inputs) - has_pending_task = any(isinstance(r, asyncio.Task) and not r.done() for r in return_values) - if has_pending_task: - return return_values, {}, False, has_pending_task - output, ui, has_subgraph = get_output_from_returns(return_values, obj) - return output, ui, has_subgraph, False - -def get_output_from_returns(return_values, obj): - results = [] - uis = [] - subgraph_results = [] - has_subgraph = False - for i in range(len(return_values)): - r = return_values[i] - if isinstance(r, dict): - if 'ui' in r: - uis.append(r['ui']) - if 'expand' in r: - # Perform an expansion, but do not append results - has_subgraph = True - new_graph = r['expand'] - result = r.get("result", None) - if isinstance(result, ExecutionBlocker): - result = tuple([result] * len(obj.RETURN_TYPES)) - subgraph_results.append((new_graph, result)) - elif 'result' in r: - result = r.get("result", None) - if isinstance(result, ExecutionBlocker): - result = tuple([result] * len(obj.RETURN_TYPES)) - results.append(result) - subgraph_results.append((None, result)) - elif isinstance(r, _NodeOutputInternal): - # V3 - if r.ui is not None: - if isinstance(r.ui, dict): - uis.append(r.ui) - else: - uis.append(r.ui.as_dict()) - if r.expand is not None: - has_subgraph = True - new_graph = r.expand - result = r.result - if r.block_execution is not None: - result = tuple([ExecutionBlocker(r.block_execution)] * len(obj.RETURN_TYPES)) - subgraph_results.append((new_graph, result)) - elif r.result is not None: - result = r.result - if r.block_execution is not None: - result = tuple([ExecutionBlocker(r.block_execution)] * len(obj.RETURN_TYPES)) - results.append(result) - subgraph_results.append((None, result)) - else: - if isinstance(r, ExecutionBlocker): - r = tuple([r] * len(obj.RETURN_TYPES)) - results.append(r) - subgraph_results.append((None, r)) - - if has_subgraph: - output = subgraph_results - elif len(results) > 0: - output = merge_result_data(results, obj) - else: - output = [] - ui = dict() - # TODO: Think there's an existing bug here - # If we're performing a subgraph expansion, we probably shouldn't be returning UI values yet. - # They'll get cached without the completed subgraphs. It's an edge case and I'm not aware of - # any nodes that use both subgraph expansion and custom UI outputs, but might be a problem in the future. - if len(uis) > 0: - ui = {k: [y for x in uis for y in x[k]] for k in uis[0].keys()} - return output, ui, has_subgraph - -def format_value(x): - if x is None: - return None - elif isinstance(x, (int, float, bool, str)): - return x - else: - return str(x) - -async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes): - unique_id = current_item - real_node_id = dynprompt.get_real_node_id(unique_id) - display_node_id = dynprompt.get_display_node_id(unique_id) - parent_node_id = dynprompt.get_parent_node_id(unique_id) - inputs = dynprompt.get_node(unique_id)['inputs'] - class_type = dynprompt.get_node(unique_id)['class_type'] - class_def = nodes.NODE_CLASS_MAPPINGS[class_type] - if caches.outputs.get(unique_id) is not None: - if server.client_id is not None: - cached_output = caches.ui.get(unique_id) or {} - server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": cached_output.get("output",None), "prompt_id": prompt_id }, server.client_id) - get_progress_state().finish_progress(unique_id) - return (ExecutionResult.SUCCESS, None, None) - - input_data_all = None - try: - if unique_id in pending_async_nodes: - results = [] - for r in pending_async_nodes[unique_id]: - if isinstance(r, asyncio.Task): - try: - results.append(r.result()) - except Exception as ex: - # An async task failed - propagate the exception up - del pending_async_nodes[unique_id] - raise ex - else: - results.append(r) - del pending_async_nodes[unique_id] - output_data, output_ui, has_subgraph = get_output_from_returns(results, class_def) - elif unique_id in pending_subgraph_results: - cached_results = pending_subgraph_results[unique_id] - resolved_outputs = [] - for is_subgraph, result in cached_results: - if not is_subgraph: - resolved_outputs.append(result) - else: - resolved_output = [] - for r in result: - if is_link(r): - source_node, source_output = r[0], r[1] - node_output = caches.outputs.get(source_node)[source_output] - for o in node_output: - resolved_output.append(o) - - else: - resolved_output.append(r) - resolved_outputs.append(tuple(resolved_output)) - output_data = merge_result_data(resolved_outputs, class_def) - output_ui = [] - has_subgraph = False - else: - get_progress_state().start_progress(unique_id) - input_data_all, missing_keys, hidden_inputs = get_input_data(inputs, class_def, unique_id, caches.outputs, dynprompt, extra_data) - if server.client_id is not None: - server.last_node_id = display_node_id - server.send_sync("executing", { "node": unique_id, "display_node": display_node_id, "prompt_id": prompt_id }, server.client_id) - - obj = caches.objects.get(unique_id) - if obj is None: - obj = class_def() - caches.objects.set(unique_id, obj) - - if issubclass(class_def, _ComfyNodeInternal): - lazy_status_present = first_real_override(class_def, "check_lazy_status") is not None - else: - lazy_status_present = getattr(obj, "check_lazy_status", None) is not None - if lazy_status_present: - required_inputs = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, "check_lazy_status", allow_interrupt=True, hidden_inputs=hidden_inputs) - required_inputs = await resolve_map_node_over_list_results(required_inputs) - required_inputs = set(sum([r for r in required_inputs if isinstance(r,list)], [])) - required_inputs = [x for x in required_inputs if isinstance(x,str) and ( - x not in input_data_all or x in missing_keys - )] - if len(required_inputs) > 0: - for i in required_inputs: - execution_list.make_input_strong_link(unique_id, i) - return (ExecutionResult.PENDING, None, None) - - def execution_block_cb(block): - if block.message is not None: - mes = { - "prompt_id": prompt_id, - "node_id": unique_id, - "node_type": class_type, - "executed": list(executed), - - "exception_message": f"Execution Blocked: {block.message}", - "exception_type": "ExecutionBlocked", - "traceback": [], - "current_inputs": [], - "current_outputs": [], - } - server.send_sync("execution_error", mes, server.client_id) - return ExecutionBlocker(None) - else: - return block - def pre_execute_cb(call_index): - # TODO - How to handle this with async functions without contextvars (which requires Python 3.12)? - GraphBuilder.set_default_prefix(unique_id, call_index, 0) - output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, hidden_inputs=hidden_inputs) - if has_pending_tasks: - pending_async_nodes[unique_id] = output_data - unblock = execution_list.add_external_block(unique_id) - async def await_completion(): - tasks = [x for x in output_data if isinstance(x, asyncio.Task)] - await asyncio.gather(*tasks, return_exceptions=True) - unblock() - asyncio.create_task(await_completion()) - return (ExecutionResult.PENDING, None, None) - if len(output_ui) > 0: - caches.ui.set(unique_id, { - "meta": { - "node_id": unique_id, - "display_node": display_node_id, - "parent_node": parent_node_id, - "real_node_id": real_node_id, - }, - "output": output_ui - }) - if server.client_id is not None: - server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": output_ui, "prompt_id": prompt_id }, server.client_id) - if has_subgraph: - cached_outputs = [] - new_node_ids = [] - new_output_ids = [] - new_output_links = [] - for i in range(len(output_data)): - new_graph, node_outputs = output_data[i] - if new_graph is None: - cached_outputs.append((False, node_outputs)) - else: - # Check for conflicts - for node_id in new_graph.keys(): - if dynprompt.has_node(node_id): - raise DuplicateNodeError(f"Attempt to add duplicate node {node_id}. Ensure node ids are unique and deterministic or use graph_utils.GraphBuilder.") - for node_id, node_info in new_graph.items(): - new_node_ids.append(node_id) - display_id = node_info.get("override_display_id", unique_id) - dynprompt.add_ephemeral_node(node_id, node_info, unique_id, display_id) - # Figure out if the newly created node is an output node - class_type = node_info["class_type"] - class_def = nodes.NODE_CLASS_MAPPINGS[class_type] - if hasattr(class_def, 'OUTPUT_NODE') and class_def.OUTPUT_NODE == True: - new_output_ids.append(node_id) - for i in range(len(node_outputs)): - if is_link(node_outputs[i]): - from_node_id, from_socket = node_outputs[i][0], node_outputs[i][1] - new_output_links.append((from_node_id, from_socket)) - cached_outputs.append((True, node_outputs)) - new_node_ids = set(new_node_ids) - for cache in caches.all: - subcache = await cache.ensure_subcache_for(unique_id, new_node_ids) - subcache.clean_unused() - for node_id in new_output_ids: - execution_list.add_node(node_id) - for link in new_output_links: - execution_list.add_strong_link(link[0], link[1], unique_id) - pending_subgraph_results[unique_id] = cached_outputs - return (ExecutionResult.PENDING, None, None) - caches.outputs.set(unique_id, output_data) - except comfy.model_management.InterruptProcessingException as iex: - logging.info("Processing interrupted") - - # skip formatting inputs/outputs - error_details = { - "node_id": real_node_id, - } - - return (ExecutionResult.FAILURE, error_details, iex) - except Exception as ex: - typ, _, tb = sys.exc_info() - exception_type = full_type_name(typ) - input_data_formatted = {} - if input_data_all is not None: - input_data_formatted = {} - for name, inputs in input_data_all.items(): - input_data_formatted[name] = [format_value(x) for x in inputs] - - logging.error(f"!!! Exception during processing !!! {ex}") - logging.error(traceback.format_exc()) - tips = "" - - if isinstance(ex, comfy.model_management.OOM_EXCEPTION): - tips = "This error means you ran out of memory on your GPU.\n\nTIPS: If the workflow worked before you might have accidentally set the batch_size to a large number." - logging.error("Got an OOM, unloading all loaded models.") - comfy.model_management.unload_all_models() - - error_details = { - "node_id": real_node_id, - "exception_message": "{}\n{}".format(ex, tips), - "exception_type": exception_type, - "traceback": traceback.format_tb(tb), - "current_inputs": input_data_formatted - } - - return (ExecutionResult.FAILURE, error_details, ex) - - get_progress_state().finish_progress(unique_id) - executed.add(unique_id) - - return (ExecutionResult.SUCCESS, None, None) - -class PromptExecutor: - def __init__(self, server, cache_type=False, cache_size=None): - self.cache_size = cache_size - self.cache_type = cache_type - self.server = server - self.reset() - - def reset(self): - self.caches = CacheSet(cache_type=self.cache_type, cache_size=self.cache_size) - self.status_messages = [] - self.success = True - - def add_message(self, event, data: dict, broadcast: bool): - data = { - **data, - "timestamp": int(time.time() * 1000), - } - self.status_messages.append((event, data)) - if self.server.client_id is not None or broadcast: - self.server.send_sync(event, data, self.server.client_id) - - def handle_execution_error(self, prompt_id, prompt, current_outputs, executed, error, ex): - node_id = error["node_id"] - class_type = prompt[node_id]["class_type"] - - # First, send back the status to the frontend depending - # on the exception type - if isinstance(ex, comfy.model_management.InterruptProcessingException): - mes = { - "prompt_id": prompt_id, - "node_id": node_id, - "node_type": class_type, - "executed": list(executed), - } - self.add_message("execution_interrupted", mes, broadcast=True) - else: - mes = { - "prompt_id": prompt_id, - "node_id": node_id, - "node_type": class_type, - "executed": list(executed), - "exception_message": error["exception_message"], - "exception_type": error["exception_type"], - "traceback": error["traceback"], - "current_inputs": error["current_inputs"], - "current_outputs": list(current_outputs), - } - self.add_message("execution_error", mes, broadcast=False) - - def execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]): - asyncio.run(self.execute_async(prompt, prompt_id, extra_data, execute_outputs)) - - async def execute_async(self, prompt, prompt_id, extra_data={}, execute_outputs=[]): - nodes.interrupt_processing(False) - - if "client_id" in extra_data: - self.server.client_id = extra_data["client_id"] - else: - self.server.client_id = None - - self.status_messages = [] - self.add_message("execution_start", { "prompt_id": prompt_id}, broadcast=False) - - with torch.inference_mode(): - dynamic_prompt = DynamicPrompt(prompt) - reset_progress_state(prompt_id, dynamic_prompt) - add_progress_handler(WebUIProgressHandler(self.server)) - is_changed_cache = IsChangedCache(prompt_id, dynamic_prompt, self.caches.outputs) - for cache in self.caches.all: - await cache.set_prompt(dynamic_prompt, prompt.keys(), is_changed_cache) - cache.clean_unused() - - cached_nodes = [] - for node_id in prompt: - if self.caches.outputs.get(node_id) is not None: - cached_nodes.append(node_id) - - comfy.model_management.cleanup_models_gc() - self.add_message("execution_cached", - { "nodes": cached_nodes, "prompt_id": prompt_id}, - broadcast=False) - pending_subgraph_results = {} - pending_async_nodes = {} # TODO - Unify this with pending_subgraph_results - executed = set() - execution_list = ExecutionList(dynamic_prompt, self.caches.outputs) - current_outputs = self.caches.outputs.all_node_ids() - for node_id in list(execute_outputs): - execution_list.add_node(node_id) - - while not execution_list.is_empty(): - node_id, error, ex = await execution_list.stage_node_execution() - if error is not None: - self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex) - break - - assert node_id is not None, "Node ID should not be None at this point" - result, error, ex = await execute(self.server, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes) - self.success = result != ExecutionResult.FAILURE - if result == ExecutionResult.FAILURE: - self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex) - break - elif result == ExecutionResult.PENDING: - execution_list.unstage_node_execution() - else: # result == ExecutionResult.SUCCESS: - execution_list.complete_node_execution() - else: - # Only execute when the while-loop ends without break - self.add_message("execution_success", { "prompt_id": prompt_id }, broadcast=False) - - ui_outputs = {} - meta_outputs = {} - all_node_ids = self.caches.ui.all_node_ids() - for node_id in all_node_ids: - ui_info = self.caches.ui.get(node_id) - if ui_info is not None: - ui_outputs[node_id] = ui_info["output"] - meta_outputs[node_id] = ui_info["meta"] - self.history_result = { - "outputs": ui_outputs, - "meta": meta_outputs, - } - self.server.last_node_id = None - if comfy.model_management.DISABLE_SMART_MEMORY: - comfy.model_management.unload_all_models() - - -async def validate_inputs(prompt_id, prompt, item, validated): - unique_id = item - if unique_id in validated: - return validated[unique_id] - - inputs = prompt[unique_id]['inputs'] - class_type = prompt[unique_id]['class_type'] - obj_class = nodes.NODE_CLASS_MAPPINGS[class_type] - - class_inputs = obj_class.INPUT_TYPES() - valid_inputs = set(class_inputs.get('required',{})).union(set(class_inputs.get('optional',{}))) - - errors = [] - valid = True - - validate_function_inputs = [] - validate_has_kwargs = False - if issubclass(obj_class, _ComfyNodeInternal): - validate_function_name = "validate_inputs" - validate_function = first_real_override(obj_class, validate_function_name) - else: - validate_function_name = "VALIDATE_INPUTS" - validate_function = getattr(obj_class, validate_function_name, None) - if validate_function is not None: - argspec = inspect.getfullargspec(validate_function) - validate_function_inputs = argspec.args - validate_has_kwargs = argspec.varkw is not None - received_types = {} - - for x in valid_inputs: - input_type, input_category, extra_info = get_input_info(obj_class, x, class_inputs) - assert extra_info is not None - if x not in inputs: - if input_category == "required": - error = { - "type": "required_input_missing", - "message": "Required input is missing", - "details": f"{x}", - "extra_info": { - "input_name": x - } - } - errors.append(error) - continue - - val = inputs[x] - info = (input_type, extra_info) - if isinstance(val, list): - if len(val) != 2: - error = { - "type": "bad_linked_input", - "message": "Bad linked input, must be a length-2 list of [node_id, slot_index]", - "details": f"{x}", - "extra_info": { - "input_name": x, - "input_config": info, - "received_value": val - } - } - errors.append(error) - continue - - o_id = val[0] - o_class_type = prompt[o_id]['class_type'] - r = nodes.NODE_CLASS_MAPPINGS[o_class_type].RETURN_TYPES - received_type = r[val[1]] - received_types[x] = received_type - if 'input_types' not in validate_function_inputs and not validate_node_input(received_type, input_type): - details = f"{x}, received_type({received_type}) mismatch input_type({input_type})" - error = { - "type": "return_type_mismatch", - "message": "Return type mismatch between linked nodes", - "details": details, - "extra_info": { - "input_name": x, - "input_config": info, - "received_type": received_type, - "linked_node": val - } - } - errors.append(error) - continue - try: - r = await validate_inputs(prompt_id, prompt, o_id, validated) - if r[0] is False: - # `r` will be set in `validated[o_id]` already - valid = False - continue - except Exception as ex: - typ, _, tb = sys.exc_info() - valid = False - exception_type = full_type_name(typ) - reasons = [{ - "type": "exception_during_inner_validation", - "message": "Exception when validating inner node", - "details": str(ex), - "extra_info": { - "input_name": x, - "input_config": info, - "exception_message": str(ex), - "exception_type": exception_type, - "traceback": traceback.format_tb(tb), - "linked_node": val - } - }] - validated[o_id] = (False, reasons, o_id) - continue - else: - try: - # Unwraps values wrapped in __value__ key. This is used to pass - # list widget value to execution, as by default list value is - # reserved to represent the connection between nodes. - if isinstance(val, dict) and "__value__" in val: - val = val["__value__"] - inputs[x] = val - - if input_type == "INT": - val = int(val) - inputs[x] = val - if input_type == "FLOAT": - val = float(val) - inputs[x] = val - if input_type == "STRING": - val = str(val) - inputs[x] = val - if input_type == "BOOLEAN": - val = bool(val) - inputs[x] = val - except Exception as ex: - error = { - "type": "invalid_input_type", - "message": f"Failed to convert an input value to a {input_type} value", - "details": f"{x}, {val}, {ex}", - "extra_info": { - "input_name": x, - "input_config": info, - "received_value": val, - "exception_message": str(ex) - } - } - errors.append(error) - continue - - if x not in validate_function_inputs and not validate_has_kwargs: - if "min" in extra_info and val < extra_info["min"]: - error = { - "type": "value_smaller_than_min", - "message": "Value {} smaller than min of {}".format(val, extra_info["min"]), - "details": f"{x}", - "extra_info": { - "input_name": x, - "input_config": info, - "received_value": val, - } - } - errors.append(error) - continue - if "max" in extra_info and val > extra_info["max"]: - error = { - "type": "value_bigger_than_max", - "message": "Value {} bigger than max of {}".format(val, extra_info["max"]), - "details": f"{x}", - "extra_info": { - "input_name": x, - "input_config": info, - "received_value": val, - } - } - errors.append(error) - continue - - if isinstance(input_type, list): - combo_options = input_type - if val not in combo_options: - input_config = info - list_info = "" - - # Don't send back gigantic lists like if they're lots of - # scanned model filepaths - if len(combo_options) > 20: - list_info = f"(list of length {len(combo_options)})" - input_config = None - else: - list_info = str(combo_options) - - error = { - "type": "value_not_in_list", - "message": "Value not in list", - "details": f"{x}: '{val}' not in {list_info}", - "extra_info": { - "input_name": x, - "input_config": input_config, - "received_value": val, - } - } - errors.append(error) - continue - - if len(validate_function_inputs) > 0 or validate_has_kwargs: - input_data_all, _, hidden_inputs = get_input_data(inputs, obj_class, unique_id) - input_filtered = {} - for x in input_data_all: - if x in validate_function_inputs or validate_has_kwargs: - input_filtered[x] = input_data_all[x] - if 'input_types' in validate_function_inputs: - input_filtered['input_types'] = [received_types] - - ret = await _async_map_node_over_list(prompt_id, unique_id, obj_class, input_filtered, validate_function_name, hidden_inputs=hidden_inputs) - ret = await resolve_map_node_over_list_results(ret) - for x in input_filtered: - for i, r in enumerate(ret): - if r is not True and not isinstance(r, ExecutionBlocker): - details = f"{x}" - if r is not False: - details += f" - {str(r)}" - - error = { - "type": "custom_validation_failed", - "message": "Custom validation failed for node", - "details": details, - "extra_info": { - "input_name": x, - } - } - errors.append(error) - continue - - if len(errors) > 0 or valid is not True: - ret = (False, errors, unique_id) - else: - ret = (True, [], unique_id) - - validated[unique_id] = ret - return ret - -def full_type_name(klass): - module = klass.__module__ - if module == 'builtins': - return klass.__qualname__ - return module + '.' + klass.__qualname__ - -async def validate_prompt(prompt_id, prompt, partial_execution_list: Union[list[str], None]): - outputs = set() - for x in prompt: - if 'class_type' not in prompt[x]: - error = { - "type": "invalid_prompt", - "message": "Cannot execute because a node is missing the class_type property.", - "details": f"Node ID '#{x}'", - "extra_info": {} - } - return (False, error, [], {}) - - class_type = prompt[x]['class_type'] - class_ = nodes.NODE_CLASS_MAPPINGS.get(class_type, None) - if class_ is None: - error = { - "type": "invalid_prompt", - "message": f"Cannot execute because node {class_type} does not exist.", - "details": f"Node ID '#{x}'", - "extra_info": {} - } - return (False, error, [], {}) - - if hasattr(class_, 'OUTPUT_NODE') and class_.OUTPUT_NODE is True: - if partial_execution_list is None or x in partial_execution_list: - outputs.add(x) - - if len(outputs) == 0: - error = { - "type": "prompt_no_outputs", - "message": "Prompt has no outputs", - "details": "", - "extra_info": {} - } - return (False, error, [], {}) - - good_outputs = set() - errors = [] - node_errors = {} - validated = {} - for o in outputs: - valid = False - reasons = [] - try: - m = await validate_inputs(prompt_id, prompt, o, validated) - valid = m[0] - reasons = m[1] - except Exception as ex: - typ, _, tb = sys.exc_info() - valid = False - exception_type = full_type_name(typ) - reasons = [{ - "type": "exception_during_validation", - "message": "Exception when validating node", - "details": str(ex), - "extra_info": { - "exception_type": exception_type, - "traceback": traceback.format_tb(tb) - } - }] - validated[o] = (False, reasons, o) - - if valid is True: - good_outputs.add(o) - else: - logging.error(f"Failed to validate prompt for output {o}:") - if len(reasons) > 0: - logging.error("* (prompt):") - for reason in reasons: - logging.error(f" - {reason['message']}: {reason['details']}") - errors += [(o, reasons)] - for node_id, result in validated.items(): - valid = result[0] - reasons = result[1] - # If a node upstream has errors, the nodes downstream will also - # be reported as invalid, but there will be no errors attached. - # So don't return those nodes as having errors in the response. - if valid is not True and len(reasons) > 0: - if node_id not in node_errors: - class_type = prompt[node_id]['class_type'] - node_errors[node_id] = { - "errors": reasons, - "dependent_outputs": [], - "class_type": class_type - } - logging.error(f"* {class_type} {node_id}:") - for reason in reasons: - logging.error(f" - {reason['message']}: {reason['details']}") - node_errors[node_id]["dependent_outputs"].append(o) - logging.error("Output will be ignored") - - if len(good_outputs) == 0: - errors_list = [] - for o, errors in errors: - for error in errors: - errors_list.append(f"{error['message']}: {error['details']}") - errors_list = "\n".join(errors_list) - - error = { - "type": "prompt_outputs_failed_validation", - "message": "Prompt outputs failed validation", - "details": errors_list, - "extra_info": {} - } - - return (False, error, list(good_outputs), node_errors) - - return (True, None, list(good_outputs), node_errors) - -MAXIMUM_HISTORY_SIZE = 10000 - -class PromptQueue: - def __init__(self, server): - self.server = server - self.mutex = threading.RLock() - self.not_empty = threading.Condition(self.mutex) - self.task_counter = 0 - self.queue = [] - self.currently_running = {} - self.history = {} - self.flags = {} - - def put(self, item): - with self.mutex: - heapq.heappush(self.queue, item) - self.server.queue_updated() - self.not_empty.notify() - - def get(self, timeout=None): - with self.not_empty: - while len(self.queue) == 0: - self.not_empty.wait(timeout=timeout) - if timeout is not None and len(self.queue) == 0: - return None - item = heapq.heappop(self.queue) - i = self.task_counter - self.currently_running[i] = copy.deepcopy(item) - self.task_counter += 1 - self.server.queue_updated() - return (item, i) - - class ExecutionStatus(NamedTuple): - status_str: Literal['success', 'error'] - completed: bool - messages: List[str] - - def task_done(self, item_id, history_result, - status: Optional['PromptQueue.ExecutionStatus']): - with self.mutex: - prompt = self.currently_running.pop(item_id) - if len(self.history) > MAXIMUM_HISTORY_SIZE: - self.history.pop(next(iter(self.history))) - - status_dict: Optional[dict] = None - if status is not None: - status_dict = copy.deepcopy(status._asdict()) - - # Remove sensitive data from extra_data before storing in history - for sensitive_val in SENSITIVE_EXTRA_DATA_KEYS: - if sensitive_val in prompt[3]: - prompt[3].pop(sensitive_val) - - self.history[prompt[1]] = { - "prompt": prompt, - "outputs": {}, - 'status': status_dict, - } - self.history[prompt[1]].update(history_result) - self.server.queue_updated() - - # Note: slow - def get_current_queue(self): - with self.mutex: - out = [] - for x in self.currently_running.values(): - out += [x] - return (out, copy.deepcopy(self.queue)) - - # read-safe as long as queue items are immutable - def get_current_queue_volatile(self): - with self.mutex: - running = [x for x in self.currently_running.values()] - queued = copy.copy(self.queue) - return (running, queued) - - def get_tasks_remaining(self): - with self.mutex: - return len(self.queue) + len(self.currently_running) - - def wipe_queue(self): - with self.mutex: - self.queue = [] - self.server.queue_updated() - - def delete_queue_item(self, function): - with self.mutex: - for x in range(len(self.queue)): - if function(self.queue[x]): - if len(self.queue) == 1: - self.wipe_queue() - else: - self.queue.pop(x) - heapq.heapify(self.queue) - self.server.queue_updated() - return True - return False - - def get_history(self, prompt_id=None, max_items=None, offset=-1, map_function=None): - with self.mutex: - if prompt_id is None: - out = {} - i = 0 - if offset < 0 and max_items is not None: - offset = len(self.history) - max_items - for k in self.history: - if i >= offset: - p = self.history[k] - if map_function is not None: - p = map_function(p) - out[k] = p - if max_items is not None and len(out) >= max_items: - break - i += 1 - return out - elif prompt_id in self.history: - p = self.history[prompt_id] - if map_function is None: - p = copy.deepcopy(p) - else: - p = map_function(p) - return {prompt_id: p} - else: - return {} - - def wipe_history(self): - with self.mutex: - self.history = {} - - def delete_history_item(self, id_to_delete): - with self.mutex: - self.history.pop(id_to_delete, None) - - def set_flag(self, name, data): - with self.mutex: - self.flags[name] = data - self.not_empty.notify() - - def get_flags(self, reset=True): - with self.mutex: - if reset: - ret = self.flags - self.flags = {} - return ret - else: - return self.flags.copy() diff --git a/memory-bank/comfyui-lineage.md b/memory-bank/comfyui-lineage.md deleted file mode 100644 index a6fddd1..0000000 --- a/memory-bank/comfyui-lineage.md +++ /dev/null @@ -1,108 +0,0 @@ -# ComfyUI Core Lineage & Integration Analysis (Updated 2025-09-29) - -## Overview - -ComfyUI‑MultiGPU extends (does not replace) ComfyUI core. Principles: -- Extend, not replace: patch specific core functions and inherit existing nodes -- Fail loudly: small, explicit patch points so core API changes surface quickly -- User agency: device placement is explicit and honored -- Multi‑device native: treat all devices as first‑class - -Current code reality: -- Phase 3 “Selective Ejection” is implemented via a per‑model flag (no global sentinel). -- Outstanding caveat: when no models are flagged, the current unload path delegates to the original core unload (unloads everything). Target is strict no‑op in this branch. - -## ComfyUI Core Foundation (Reference) - -Key concepts implemented by ComfyUI core (see memory-bank/comfy_core.py snapshot): -- Global list: `current_loaded_models` -- Model wrapper: `LoadedModel` with methods like `model_load`, `model_unload`, `model_memory_required` -- Memory utilities: `soft_empty_cache()`, `get_free_memory()`, etc. -- Prompt execution: - - `/free` endpoint sets queue flags: `unload_models`, `free_memory` (server.py) - - `main.py` prompt worker consumes flags: - - If `unload_models` (or `free_memory`): `comfy.model_management.unload_all_models()` - - If `free_memory`: `PromptExecutor.reset()` - - Then GC + `comfy.model_management.soft_empty_cache()` - -This is the canonical “Manager button” path for model + execution cache cleanup. - -## How MultiGPU Extends ComfyUI Core - -MultiGPU adds small patches and inherits nodes to enable multi‑device behavior while preserving ComfyUI’s flow. - -### 1) Device selection alignment -- File: `__init__.py` -- Patches: - - `mm.get_torch_device = get_torch_device_patched` - - `mm.text_encoder_device = text_encoder_device_patched` -- Purpose: Respect user‑selected devices supplied by MultiGPU wrappers while staying coherent with ComfyUI’s device model. - -### 2) Multi‑device VRAM cache + CPU reset -- File: `__init__.py` -- Patch: - - `mm.soft_empty_cache = soft_empty_cache_distorch2_patched` -- Behavior: - - Detects if any DisTorch2 model is active and clears allocator caches on ALL devices via `soft_empty_cache_multigpu()` (from `device_utils.py`) - - Integrates adaptive CPU memory reset; can force `PromptExecutor.reset()` on `force=True` for Manager parity - -### 3) Selective ejection (patched unload) -- File: `model_management_mgpu.py` -- Patch: - - `mm.unload_all_models = _mgpu_patched_unload_all_models` -- Behavior: - - Iterate `mm.current_loaded_models` and split into: - - `models_to_unload`: models with per‑model flag `_mgpu_unload_distorch_model == True` - - `kept_models`: all others - - If any flagged: unload only the flagged models and set `mm.current_loaded_models = kept_models` - - Current caveat: If none are flagged (all kept), code delegates to original core unload, which unloads everything (target: strict no‑op for this branch) - -### 4) Per‑model flag is set at load time (no global sentinel) -- File: `distorch_2.py` -- Where: - - In DisTorch2 wrappers (UNET/CLIP/VAE) inside `override(...)`, after calling the original loader: - - `out[0].model._mgpu_unload_distorch_model = (keep_loaded == False)` -- Rationale: - - Surgical precision at model granularity and no persistent global state - -### 5) Manager parity helper for tests/flows -- File: `model_management_mgpu.py` -- Function: - - `force_full_system_cleanup(reason="manual", force=True)` -- Behavior: - - Sets both `unload_models=True` and `free_memory=True` on the PromptQueue, just like the Manager “Free model and node cache” button - -## End‑to‑End Free Flow (Now) - -“Manager button” or parity helper triggers the same core actions: - -1) POST /free with `{"unload_models": true, "free_memory": true}` -2) `main.py` prompt worker consumes flags: - - Calls `comfy.model_management.unload_all_models()` - - MultiGPU patched unload runs: - - If any models flagged via `_mgpu_unload_distorch_model=True`: unload only those and retain others - - If none are flagged: current code delegates to original unload (unloads everything) — under review - - Calls `PromptExecutor.reset()` - - GC + `comfy.model_management.soft_empty_cache()` - - MultiGPU patched soft empty runs: - - Multi‑device allocator cache clear (CUDA/MPS/XPU/NPU/MLU/DirectML/CoreX as available) - - Optional CPU reset behavior when forced - -Intended invariant (target): -- Only flagged DisTorch2 models are ejected; unflagged (keep_loaded=True) models remain live after the full flow. - -## Behavior Notes & Next Step - -- Implemented: - - Per‑model selective ejection (Phase 3) without global sentinel - - Multi‑device allocator clearing and Manager parity semantics -- Caveat: - - If no models are flagged, current patched unload delegates to original unload (unloads everything) - - This can defeat selectiveness when all models are intended to be retained -- Next step (hardening): - - Reinstate “strict no‑op” in the all‑kept branch of `_mgpu_patched_unload_all_models` (never delegate to original unload if nothing is flagged) - - Add instrumentation around pre/post unload, post reset, post soft‑empty to ensure retained models remain alive - -## Sequence Summary - -A) Vanilla ComfyUI Manager “Free diff --git a/memory-bank/performance-benchmarks.md b/memory-bank/performance-benchmarks.md deleted file mode 100644 index 46517ee..0000000 --- a/memory-bank/performance-benchmarks.md +++ /dev/null @@ -1,281 +0,0 @@ -# Performance Benchmarks & Hardware Analysis - -## Executive Summary - -Comprehensive benchmarking across 5 model architectures and 6 hardware configurations reveals **bandwidth is king** for DisTorch2 performance. NVLink provides near-native performance while PCIe 4.0 CPU offloading offers excellent price/performance for most users. - -## Benchmark Configuration - -### Test Systems -- **PCIe 3.0 System**: i7-11700F @ 2.50GHz, DDR4-2667, older motherboard -- **PCIe 4.0 System**: Ryzen 5 7600X @ 4.70GHz, DDR5-4800, modern motherboard - -### Hardware Configurations Tested -1. **RTX 3090 (no donor)**: Baseline - 799.3 GB/s internal VRAM -2. **x8 PCIe 3.0 CPU**: 6.8 GB/s measured bandwidth -3. **x16 PCIe 4.0 CPU**: 27.2 GB/s measured bandwidth -4. **RTX 3090 (NVLINK)**: 50.8 GB/s high-speed interconnect -5. **RTX 3090 (x8)**: 4.4 GB/s P2P over limited bus -6. **GTX 1660 Ti (x4)**: 2.1 GB/s P2P over slow bus - -## Model Performance Analysis - -### QWEN Image (FP8 - 19GB Model) - -| GB Offloaded | RTX 3090 (no donor) | x8 PCIe 3.0 CPU | x16 PCIe 4.0 CPU | RTX 3090 (NVLINK) | RTX 3090 (x8) | GTX 1660 Ti (x4) | -|--------------|---------------------|-----------------|------------------|-------------------|---------------|------------------| -| 0 | 4.28s | 4.28s | 4.45s | 4.28s | 4.28s | 4.28s | -| 1.2 | 4.28s | 4.71s | 4.59s | 4.37s | 5.77s | 6.64s | -| 2.4 | 4.28s | 5.16s | 4.71s | 4.45s | 7.27s | 9.01s | -| 4.8 | 4.28s | 6.07s | 4.89s | 4.63s | 10.28s | 13.79s | -| 9.5 | 4.28s | 7.84s | 5.39s | 4.95s | 16.21s | #N/A | -| 19 | 4.28s | 11.43s | 6.30s | 5.64s | 28.33s | #N/A | - -**Key Insights**: -- **NVLink Excellence**: Only 32% slowdown at maximum offloading (5.64s vs 4.28s) -- **PCIe 4.0 Sweet Spot**: 47% slowdown at maximum offloading (6.30s vs 4.28s) -- **x8 GPU Penalty**: 562% slowdown shows P2P limitations (28.33s vs 4.28s) - -### FLUX GGUF (Q8_0 - 12GB Model) - -| GB Offloaded | RTX 3090 (no donor) | x8 PCIe 3.0 CPU | x16 PCIe 4.0 CPU | RTX 3090 (NVLINK) | RTX 3090 (x8) | GTX 1660 Ti (x4) | -|--------------|---------------------|-----------------|------------------|-------------------|---------------|------------------| -| 0 | 1.29s | 1.29s | 1.32s | 1.29s | 1.29s | 1.29s | -| 1.5 | 1.29s | 1.6s | 1.4s | 1.32s | 1.76s | 2s | -| 3 | 1.29s | 1.9s | 1.49s | 1.35s | 2.24s | 2.74s | -| 5.9 | 1.29s | 2.5s | 1.65s | 1.41s | 3.15s | #N/A | -| 11.8 | 1.29s | 3.76s | 1.99s | 1.52s | 5.04s | #N/A | - -**Key Insights**: -- **GGUF Efficiency**: Pre-quantized format reduces transfer overhead -- **Linear Scaling**: Performance scales predictably with offload amount -- **Bandwidth Correlation**: Results align with measured connection speeds - -### WAN 2.2 (FP8 Video - 14GB Model) - -| GB Offloaded | RTX 3090 (no donor) | x8 PCIe 3.0 CPU | x16 PCIe 4.0 CPU | RTX 3090 (NVLINK) | RTX 3090 (x8) | GTX 1660 Ti (x4) | -|--------------|---------------------|-----------------|------------------|-------------------|---------------|------------------| -| 0 | 111.3s | 111.3s | 111.3s | 111.3s | 111.3s | 111.3s | -| 1.7 | 111.3s | 111.3s | 111.5s | 111.1s | 112.2s | 114.0s | -| 3.4 | 111.3s | 111.9s | 111.7s | 111.0s | 114.4s | 117.2s | -| 6.7 | 111.3s | 112.9s | 111.9s | 111.5s | 118.2s | #N/A | -| 13.3 | 111.3s | 115.5s | 112.3s | 111.9s | 126.1s | #N/A | - -**Key Insights**: -- **Video Generation Resilience**: Minimal performance impact across all configurations -- **Compute-Heavy Workload**: Long inference times mask transfer latency -- **Hardware Tolerance**: Even slow connections deliver acceptable performance -- **Maximum Impact**: Only 4% slowdown with CPU offloading (115.5s vs 111.3s) - -### FLUX-KONTEXT-FP16 (22GB Model) - -| GB Offloaded | RTX 3090 (no donor) | x8 PCIe 3.0 CPU | x16 PCIe 4.0 CPU | RTX 3090 (NVLINK) | RTX 3090 (x8) | GTX 1660 Ti (x4) | -|--------------|---------------------|-----------------|------------------|-------------------|---------------|------------------| -| 0 | 2.74s | 2.74s | 2.66s | 2.74s | 2.74s | 2.74s | -| 1.4 | 2.74s | 2.78s | 2.65s | 2.52s | 2.94s | 3.17s | -| 2.8 | 2.74s | 3.06s | 2.71s | 2.53s | 3.38s | 3.84s | -| 5.6 | 2.74s | 3.63s | 2.88s | 2.61s | 4.27s | #N/A | -| 11.1 | 2.74s | 4.76s | 3.17s | 2.71s | 6.00s | #N/A | -| 22.17 | 2.74s | 7.03s | 3.81s | 2.92s | 9.54s | #N/A | - -**Key Insights**: -- **Large Model Challenge**: 22GB model tests all configurations -- **NVLink Dominance**: Only 7% slowdown at full offload (2.92s vs 2.74s) -- **CPU Viability**: 39% slowdown acceptable for capability gain (3.81s vs 2.74s) - -### QWEN Image FP16 (38GB Model - Extreme Test) - -| GB Offloaded | x8 PCIe 3.0 CPU | RTX 3090 (NVLINK) | RTX 3090 (x8) | RTX 3090 (no donor - fp8) | -|--------------|-----------------|-------------------|---------------|---------------------------| -| 0 | #N/A | #N/A | #N/A | 4.28s | -| 16 | 10.02s | 4.61s | 14.15s | 4.28s | -| 19 | 11.12s | 4.73s | 16.07s | 4.28s | -| 22 | 12.25s | 4.88s | 17.99s | 4.28s | -| 27 | 14.13s | #N/A | #N/A | 4.28s | -| 32 | 16s | #N/A | #N/A | 4.28s | -| 38 | 18.29s | #N/A | #N/A | 4.28s | - -**Key Insights**: -- **Impossible Made Possible**: 38GB model runs on any hardware -- **NVLink Superiority**: Maintains reasonable performance even at extreme scales -- **Quality vs Convenience**: FP8 offers convenience, FP16 offers ultimate quality - -## Hardware Configuration Analysis - -### Performance Hierarchy (Best to Worst) - -1. **NVLink 2x3090** (50.8 GB/s) - - **Use Case**: Professional/enthusiast dual-GPU setups - - **Performance**: Near-native across all workloads - - **Investment**: High (requires compatible cards + motherboard) - -2. **PCIe 4.0 x16 CPU** (27.2 GB/s) - - **Use Case**: Modern single-GPU systems with fast RAM - - **Performance**: Excellent for most workloads - - **Investment**: Moderate (modern motherboard + DDR5) - -3. **PCIe 3.0 x16 CPU** (15.8 GB/s theoretical) - - **Use Case**: Older systems with capability upgrade - - **Performance**: Acceptable for most workloads, some penalty - - **Investment**: Low (leverage existing hardware) - -4. **PCIe 3.0 x8 CPU** (6.8 GB/s measured) - - **Use Case**: Budget systems, older motherboards - - **Performance**: Noticeable slowdown but functional - - **Investment**: Minimal (system RAM upgrade recommended) - -5. **PCIe 3.0 x8 P2P GPU** (4.4 GB/s measured) - - **Use Case**: Dual-GPU consumer motherboards (x8/x8 split) - - **Performance**: Significant slowdown for image work - - **Investment**: Poor ROI unless already owned - -6. **PCIe 3.0 x4 P2P GPU** (2.1 GB/s measured) - - **Use Case**: Older secondary GPUs in slow slots - - **Performance**: Severe slowdown, capacity-only benefit - - **Investment**: Only for extreme VRAM needs - -## Strategic Recommendations - -### For Image Generation (FLUX, QWEN) -**Priority: Bandwidth Optimization** - -1. **Gold Standard**: NVLink 2x3090 setup - - Effectively creates 48GB VRAM pool with minimal penalty - - Suitable for professional/enthusiast workflows - - Consider refurbished 3090s for cost optimization - -2. **Modern Path**: RTX 5090/5080 + PCIe 4.0 + DDR5 - - Single GPU with fast CPU offloading - - Future-proofs with PCIe 5.0 capabilities - - Best price/performance for new builds - -3. **Budget Path**: Existing GPU + system RAM upgrade - - Maximize system RAM (64GB+) for large model storage - - Accept performance penalty for capability gain - - Most accessible entry point - -**Avoid**: x8/x8 PCIe splits for P2P unless NVLink available - -### For Video Generation (WAN, HunyuanVideo) -**Priority: Capacity Maximization** - -1. **Any Available Hardware**: Video generation is bandwidth-tolerant - - Old GPUs in x4 slots provide meaningful capacity - - CPU offloading performs nearly as well as GPU storage - - Focus on total available memory over speed - -2. **Mixed Architecture Builds**: Combine new + old hardware - - Primary: RTX 4090/5090 for compute - - Secondary: Any available GPU for model storage - - System RAM: As much as financially feasible - -3. **Evolution Strategy**: Incremental hardware additions - - Start with single GPU + CPU offloading - - Add secondary GPUs as budget allows - - Each additional device provides capacity benefit - -### Universal Low-VRAM Strategy - -**Multi-Tool Approach**: Use entire ComfyUI-MultiGPU ecosystem - -1. **Ancillary Models**: CLIP/VAE to secondary devices - ``` - CLIPLoaderMultiGPU → cuda:1 or cpu - VAELoaderMultiGPU → cuda:1 or cpu - ``` - -2. **Main Model**: DisTorch2 for UNet distribution - ``` - UNETLoaderDisTorch2MultiGPU → expert allocation - ``` - -3. **Memory Management**: Progressive offloading strategy - - Start conservative (minimal offloading) - - Increase offloading until workflow stable - - Monitor performance vs capability tradeoff - -## Performance Scaling Laws - -### Bandwidth vs Performance Relationship - -**Linear Correlation Observed**: -- **Transfer Time = (GB Offloaded × Steps) ÷ Bandwidth** -- **Total Slowdown = Baseline Time + Transfer Time** - -**Example Calculation** (QWEN 19GB, 10 steps, 19GB offloaded): -- **NVLink** (50.8 GB/s): 19×10÷50.8 = 3.7s transfer time -- **PCIe 4.0** (27.2 GB/s): 19×10÷27.2 = 7.0s transfer time -- **PCIe 3.0 x8** (6.8 GB/s): 19×10÷6.8 = 27.9s transfer time - -**Measured vs Calculated** shows strong correlation, validating model. - -### Model Architecture Impact - -**Transfer Overhead by Model Type**: - -| Model Type | Overhead Factor | Reason | -|------------|----------------|---------| -| GGUF Models | 0.8x | Pre-quantized, optimized transfers | -| FP16 SafeTensors | 1.0x | Standard transfer overhead | -| Video Models | 0.3x | Long compute masks transfer time | -| Image Models | 1.2x | Short compute exposes transfer time | - -### Hardware Utilization Patterns - -**GPU Utilization During DisTorch Operation**: -- **Compute GPU**: 95-100% during inference steps -- **Donor GPU**: 0-15% (transfer operations only) -- **System RAM**: Varies with offload amount -- **PCIe Bus**: Burst usage during layer swaps - -**Memory Pressure Thresholds**: -- **90% VRAM**: Automatic offloading triggered -- **95% System RAM**: Performance degradation likely -- **100% Available Memory**: OOM failure imminent - -## Benchmarking Methodology - -### Test Validation -- **Consistent Environment**: Same ComfyUI version, same models -- **Multiple Runs**: 3 runs averaged, outliers discarded -- **Hardware Monitoring**: GPU-Z, HWiNFO64 for validation -- **Transfer Measurement**: Custom timing instrumentation - -### Limitations -- **Single-User Testing**: Results may vary with different hardware combinations -- **Model-Specific**: Some architectures may exhibit different patterns -- **Dynamic Factors**: System load, thermal throttling not controlled -- **Sample Size**: Limited to available hardware configurations - -### Reproducibility -```python -# Benchmark configuration used -BENCHMARK_CONFIG = { - "comfyui_version": "0.3.50", - "torch_version": "2.8.0+cu128", - "model_precision": "fp16", - "steps": 10, - "guidance_scale": 7.5, - "resolution": "1024x1024" -} -``` - -## Future Benchmarking Plans - -### Next-Generation Hardware Testing -- **RTX 5090**: PCIe 5.0 validation when available -- **PCIe 5.0 Motherboards**: Maximum bandwidth testing -- **DDR5-6000+**: RAM speed impact on CPU offloading -- **AMD RDNA4**: HIP/ROCm performance characterization - -### Extended Model Coverage -- **Mixture of Experts**: Sparse model behavior analysis -- **Multimodal Models**: Text+Vision combined workloads -- **Real-Time Models**: Streaming inference requirements -- **Custom Architectures**: Community model support - -### Advanced Metrics -- **Power Efficiency**: Performance per watt analysis -- **Thermal Behavior**: Sustained performance under load -- **Quality Metrics**: Objective image/video quality measurement -- **User Experience**: Subjective workflow satisfaction surveys diff --git a/memory-bank/productContext.md b/memory-bank/productContext.md deleted file mode 100644 index 5f61dd7..0000000 --- a/memory-bank/productContext.md +++ /dev/null @@ -1,116 +0,0 @@ -# Product Context: Why ComfyUI-MultiGPU Exists - -## The Problem Space - -### The VRAM Crisis -Modern AI models are experiencing explosive growth in size: -- **FLUX.1-dev**: 23.8GB (exceeds most consumer cards) -- **WAN 2.2**: 14GB+ (video generation demands) -- **Hunyuan Video**: 25GB+ (next-gen video models) -- **QWEN Image**: Up to 38GB in FP16 (professional image editing) - -Meanwhile, consumer hardware remains constrained: -- **RTX 4090**: 24GB VRAM (can't fit largest models) -- **RTX 3090**: 24GB VRAM (aging but still powerful) -- **RTX 4080/4070**: 16GB/12GB (mainstream but limited) -- **Budget Cards**: 8GB or less (significant portion of user base) - -### The Workflow Limitation -ComfyUI's default behavior loads entire models onto the primary GPU: -- **Latent space competition**: Model storage vs computation space -- **Resolution limits**: Large models prevent high-resolution generation -- **Batch size restrictions**: Memory consumed by static weights -- **OOM failures**: Workflows simply fail to run - -### The Speed vs. Memory Dilemma -Existing solutions force uncomfortable tradeoffs: -- **--lowvram mode**: Dynamic but unpredictable, quality issues with LoRAs -- **Quantization**: Quality loss, limited model support -- **Model switching**: Slow, workflow interruption -- **Single-GPU limitation**: Unused hardware sitting idle - -## The Vision - -### Unified Compute Pool -Transform multi-GPU setups from "main + unused" to "unified compute": -- **Primary GPU**: 100% dedicated to computation/latent processing -- **Secondary GPUs**: High-speed model storage (NVLINK, PCIe) -- **System RAM**: Extended model storage with optimized transfers -- **Mixed Architectures**: Old cards find new life as storage - -### Deterministic Memory Management -Replace dynamic allocation with user-controlled distribution: -- **Static Mapping**: Model layers assigned to specific devices -- **Predictable Performance**: Known transfer costs and timing -- **Quality Preservation**: Full-precision LoRA patching on compute device -- **Workflow Reliability**: Consistent behavior across runs - -### Hardware Democracy -Enable AI generation across hardware tiers: -- **Budget Systems**: 8GB card + system RAM for large models -- **Enthusiast Builds**: 2x3090 effectively becomes 48GB unified pool -- **Mixed Setups**: 4090 + old 1080 Ti = expanded capability -- **Enterprise**: Workstation-grade hardware optimization - -## User Experience Goals - -### For Low-VRAM Users -- **Model Access**: Run any model regardless of VRAM size -- **Resolution Freedom**: Generate at previously impossible dimensions -- **Batch Processing**: Multiple images/frames without OOM -- **Quality Maintenance**: No forced quantization or quality loss - -### For Multi-GPU Users -- **Hardware Utilization**: Every GPU contributes meaningfully -- **Performance Optimization**: NVLink, PCIe bandwidth maximization -- **Flexible Distribution**: Fine-grained control over model placement -- **Scaling Benefits**: More hardware = more capability - -### For Workflow Creators -- **Predictability**: Consistent memory usage patterns -- **Configurability**: Expert modes for precise control -- **Compatibility**: Works with existing ComfyUI workflows -- **Documentation**: Clear performance expectations - -## The Market Reality - -### Community Demand -Issues and feedback reveal consistent patterns: -- **"Only cuda:0 visible"**: Multi-GPU setup confusion -- **"Out of memory"**: VRAM exhaustion with large models -- **"Slow generation"**: Inefficient memory management -- **"Can't run X model"**: Hardware limitations blocking workflows - -### Hardware Evolution -Consumer GPU landscape trends: -- **VRAM Stagnation**: 24GB ceiling for years -- **Model Growth**: Exponential size increases -- **Price Pressure**: High-end cards increasingly expensive -- **Mixed Installations**: Users combining new + old hardware - -### Ecosystem Position -ComfyUI's role in AI generation: -- **Node-based workflows**: Flexible but memory-hungry -- **Model diversity**: Supports every major architecture -- **Community-driven**: Custom nodes enable specialization -- **Production use**: Professional workflows demand reliability - -## Success Metrics - -### Technical Success -- **Model Loading**: Any model loads on any hardware combination -- **Performance Predictability**: Benchmarked speed vs. memory tradeoffs -- **Stability**: No crashes or memory leaks in extended use -- **Compatibility**: Works across operating systems and configurations - -### User Success -- **Workflow Enablement**: Previously impossible workflows now work -- **Hardware Investment**: Old GPUs gain new utility -- **Resolution/Batch Scaling**: Tangible output quality improvements -- **Community Growth**: Increasing adoption and positive feedback - -### Ecosystem Success -- **ComfyUI Integration**: Seamless operation with core functionality -- **Developer Adoption**: Other custom nodes build on our patterns -- **Hardware Vendor Recognition**: Acknowledged in optimization discussions -- **Production Deployment**: Used in commercial/professional settings diff --git a/memory-bank/progress.md b/memory-bank/progress.md deleted file mode 100644 index 725e9da..0000000 --- a/memory-bank/progress.md +++ /dev/null @@ -1,215 +0,0 @@ -# Project Progress & Status (Updated 2025-09-30) - -## Production Status: v2.5.0 Release Candidate - -**Overall Assessment**: PRODUCTION READY -**Code Quality**: 8.5/10 - Clean, refactored, comprehensive -**Stability**: 9/10 - Verified working in production -**Performance**: 8/10 - Validated across hardware tiers -**Community**: 7.5/10 - Active adoption, growing ecosystem - -## What Works (Verified in Production) ✅ - -### Core MultiGPU Infrastructure -- **Dynamic Class Override System** (City96 pattern): Inheritance-based node wrapping, auto-adapts to ComfyCore -- **Universal Device Detection**: CPU, CUDA, MPS, XPU, NPU, MLU, DirectML, CoreX -- **Multi-Device VRAM Management**: `soft_empty_cache_multigpu()` clears allocator caches across all devices -- **Automatic Node Registration**: Detects available custom nodes and creates compatible MultiGPU variants - -### DisTorch2 Distributed Loading (Refactored) -- **Universal SafeTensor Support**: Works with any safetensor-based model -- **Load-Patch-Distribute Pipeline**: Quality-preserving LoRA patching on compute device before distribution -- **Three Allocation Modes**: Bytes (cuda:0,4gb;cpu,2gb), Ratios (cuda:0,50%;cpu,50%), Fractions (automatic) -- **CLIP Head Preservation**: Unified allocation function with CLIP-specific head handling -- **~10% Performance Improvement** over DisTorch V1 - -### Selective Unloading (Verified Working) ✅ -**Verified in Production Logs** (2025-09-30): -``` -[CATEGORIZE_SUMMARY] kept_models: 2, models_to_unload: 1, total: 3 -[SELECTIVE_UNLOAD] Proceeding with selective unload: retaining 2, unloading 1 -[REMAINING_MODEL] 0: AutoencodingEngine -[REMAINING_MODEL] 1: FluxClipModel_ -``` - -**Components**: -1. **Per-Model Flag System**: `_mgpu_unload_distorch_model` set during load based on `keep_loaded` parameter -2. **Patched unload_all_models**: Categorizes models, selectively unloads flagged ones, rebuilds `mm.current_loaded_models` -3. **GC Anchor System**: Prevents premature garbage collection of retained models -4. **Manager Parity**: `force_full_system_cleanup()` mirrors ComfyUI-Manager "Free model and node cache" - -### Hardware Configuration Support -- **NVLink**: 5-7% slowdown (near-native) -- **PCIe 4.0 x16**: 40-50% slowdown (excellent) -- **PCIe 3.0 x16**: 70-80% slowdown (good) -- **PCIe 4.0 x8**: 80-100% slowdown (acceptable) -- **PCIe 3.0 x8**: 150-200% slowdown (workable) -- **PCIe 3.0 x4**: 300-400% slowdown (last resort) - -### External Integrations -- ✅ **ComfyUI-GGUF**: DisTorch-enabled quantized model nodes -- ✅ **WanVideoWrapper**: MultiGPU video generation -- ✅ **Florence2**: Vision model support -- ✅ **HunyuanVideoWrapper**: Native VAE + device selection -- ✅ **LTXVideo**: Video generation -- ✅ **MMAudio**: Audio synthesis -- ✅ **PuLID**: Identity preservation - -### Documentation -- Comprehensive README with architecture overview -- 20+ example JSON workflows -- Performance benchmarks and hardware recommendations -- Troubleshooting guides - -## Recent Achievements (v2.5.0) - -### Code Refactoring (-219 lines total) -1. **DisTorch2 Allocation Consolidation** (-179 lines): - - Unified `analyze_safetensor_loading()` and `analyze_safetensor_loading_clip()` into single function - - CLIP head preservation via helper function `_extract_clip_head_blocks()` - - Eliminated 85% code duplication - - Single source of truth for allocation logic - -2. **Production Cleanup** (-40 lines): - - Removed diagnostic instrumentation from `model_management_mgpu.py` - - Deleted `_mgpu_instrumented_soft_empty_cache()` wrapper (debug artifact) - - Clear separation: device_utils.py = functional, model_management = lifecycle - -### Architecture Improvements -- **Comprehensive Logging**: Production-grade telemetry at every major operation -- **Clean Module Boundaries**: Single responsibility, clear dependency direction -- **No Debug Cruft**: All diagnostic code removed, only production logging remains -- **Verified Working**: Selective unload tested and confirmed in production - -## Development Roadmap - -### Immediate (This Week) -- [x] Refactor DisTorch2 allocation functions -- [x] Remove diagnostic code -- [x] Verify selective unload working -- [x] Update memory bank documentation -- [ ] Final v2.5.0 testing pass -- [ ] GitHub release notes and changelog - -### Short-term (2-4 Weeks) -- **Integration Expansion**: - - Mochi video model support - - Community-requested custom node integrations - - Issue triage and resolution - -- **Documentation**: - - Tutorial series refresh - - Hardware selection guide - - Configuration validation tools - -### Medium-term (2-3 Months) -- **User Experience**: - - Allocation string generator with validation - - Hardware profiler (bandwidth/VRAM/latency) - - Performance prediction tools - -- **Professional Features**: - - Batch processing optimization - - Quality metrics and parity validation - - Performance dashboard - -### Long-term (6-12 Months) -- **Research & Advanced Features**: - - Model parallelism experiments - - Pipeline parallelism - - Streaming inference for video - - Multi-node/cloud orchestration - -## Known Limitations & Workarounds - -### Hardware Constraints -- **DirectML Performance**: Functional but slower than native CUDA -- **CPU Offload Overhead**: PCIe bandwidth becomes bottleneck in extreme offload scenarios -- **Memory Pressure**: Adaptive thresholds may trigger premature unloads under extreme pressure - -### API Dependencies -- **ComfyCore Changes**: Fail-loudly approach surfaces API changes immediately -- **Custom Node Evolution**: Ongoing monitoring of integration points required - -### Documentation Gaps -- Advanced configuration recipes for edge cases -- Hardware-specific optimization guides (in progress) -- Video tutorial series (planned) - -## Quality Assurance - -### Technical Validation ✅ -- **Bit-exact Quality Parity**: Maintains identical output to single-GPU -- **Performance Predictability**: Consistent with hardware bandwidth tiers -- **Zero Regressions**: Selective unload working correctly -- **Comprehensive Logging**: Production debugging capabilities - -### Model Validation ✅ -- FLUX (1.dev, schnell, GGUF variants) -- WAN Video (1.3B, 2.0, 2.2) -- QWEN VL (image understanding) -- HunyuanVideo (text-to-video) -- Florence2 (vision tasks) -- SDXL, SD1.5 (classic models) - -### Community Feedback -- Active GitHub issues and discussions -- Integration requests from other node developers -- Positive feedback on performance and stability -- Actionable feature requests - -## Success Metrics - -### Technical -- ✅ Selective unload verified working in production -- ✅ Clean refactored codebase (-219 lines) -- ✅ Universal device support maintained -- ✅ Performance validated across 6 hardware tiers - -### User Impact -- ✅ Previously impossible workflows now run reliably -- ✅ Clear guidance for low-VRAM and multi-GPU users -- ✅ Reduced support load through better documentation -- ✅ Growing community adoption - -### Ecosystem -- ✅ 10+ custom node integrations -- ✅ Recognition in optimization discussions -- ✅ Community validation across hardware configs - -## Evolution of Design Decisions - -### Architectural Choices -1. **Dynamic Class Override** → Minimal code, automatic compatibility -2. **Load-Patch-Distribute** → Quality preservation, no precision loss -3. **Per-Model Flags** → Granular control without global state -4. **Fail-Loudly** → Immediate API change detection - -### Memory Management -1. **Conservative Defaults** → User control, explicit behavior -2. **Transparent Logging** → Production debugging capability -3. **Multi-Device Native** → All devices treated equally -4. **Adaptive Thresholds** → Automatic OOM prevention - -### Integration Strategy -1. **Inheritance-Based** → City96 pattern, minimal patch surface -2. **Three Core Patches** → Device selection, cache clearing, selective unload -3. **Single Source of Truth** → device_utils.py for device management - -## Next Actions - -1. **Final v2.5.0 Testing**: Edge case validation, regression tests -2. **Release Preparation**: Changelog, GitHub release notes, announcement -3. **Community Engagement**: Issue triage, feature requests, integrations -4. **Documentation**: Tutorial refresh, hardware guides, troubleshooting - -## Summary - -ComfyUI-MultiGPU v2.5.0 represents production maturity: -- Clean, refactored codebase with comprehensive logging -- Verified working selective unload system -- Universal device support across 7 accelerator types -- Quality-preserving distributed inference -- Active community with growing ecosystem - -The architecture is stable, performant, and ready for production deployment. diff --git a/memory-bank/projectbrief.md b/memory-bank/projectbrief.md deleted file mode 100644 index 3a13692..0000000 --- a/memory-bank/projectbrief.md +++ /dev/null @@ -1,54 +0,0 @@ -# ComfyUI-MultiGPU Project Brief - -## Project Identity -**Name**: ComfyUI-MultiGPU -**Maintainer**: John Pollock (@pollockjj) -**Current Version**: 2.4.7 (Production Grade) -**Repository**: https://github.com/pollockjj/ComfyUI-MultiGPU - -## Core Mission -Transform ComfyUI from single-GPU to multi-device AI inference platform. Stop using expensive compute cards for model storage - unleash them on maximum latent space instead. - -## What We Build -A ComfyUI custom_node that provides: -- **Universal Multi-Device Support**: CUDA, CPU, XPU, NPU, MLU, MPS, DirectML -- **Advanced Memory Management**: DisTorch2 distributed model loading -- **Device-Aware Node Wrapping**: MultiGPU versions of all major ComfyUI loaders -- **Production-Grade Stability**: 300+ commits, 90 resolved issues - -## Evolution Timeline -- **Aug 2024**: Basic multi-GPU device selection (Alexander Dzhoganov) -- **Dec 2024**: City96 architectural revolution (400+ lines → 50 lines via inheritance) -- **Jan 2025**: DisTorch V1 (GGUF virtual VRAM) -- **Aug 2025**: DisTorch V2.0 (Universal .safetensor support) -- **Sep 2025**: Production maturity (Version 2.4.7) - -## Core Problems Solved -1. **VRAM Limitations**: Run 38GB models on 24GB cards -2. **Hardware Utilization**: Turn mixed GPU setups into unified compute pool -3. **Memory Management**: Deterministic model distribution vs dynamic --lowvram -4. **Workflow Scaling**: Enable previously impossible resolutions/batch sizes - -## Primary User Segments -- **Low-VRAM Users**: 8GB-16GB cards accessing large models -- **Multi-GPU Enthusiasts**: 2x3090, mixed architecture setups -- **Production Users**: Consistent performance requirements -- **Video Generation**: WAN, HunyuanVideo, LTX workflows - -## Technical Foundation -- **Dynamic Class Override System**: Elegant inheritance-based node wrapping -- **Load-Patch-Distribute (LPD)**: Load on compute → patch LoRAs → distribute at FP16 -- **Virtual VRAM**: CPU/GPU memory appears as extended VRAM pool -- **Expert Allocation Modes**: Bytes, ratios, and fraction-based distribution - -## Success Metrics -- **Community Adoption**: 300+ commits, active issue resolution -- **Performance Validation**: Benchmarked across hardware configurations -- **Ecosystem Integration**: Supports 15+ model loader types -- **Stability**: Production deployments running complex workflows - -## Development Philosophy -- **Work WITH ComfyUI**: Leverage existing patterns, don't fight core -- **Fail Loudly**: No defensive coding - we want to know when ComfyCore changes -- **Self-Documenting Code**: Structure and names tell the story -- **Inheritance Over Composition**: Dynamic class overrides, not manual definitions diff --git a/memory-bank/reference-logs/d2_testing_result.txt b/memory-bank/reference-logs/d2_testing_result.txt deleted file mode 100644 index ae60be4..0000000 --- a/memory-bank/reference-logs/d2_testing_result.txt +++ /dev/null @@ -1,430 +0,0 @@ - -[MultiGPU_Memory_Management] malloc_trim(0) begin -[MultiGPU Model Management] 2025-09-28T16:55:25.058Z mem_mgmt_pre-malloc-trim cpu|45.88 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:25.083Z mem_mgmt_post-malloc-trim cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU_Memory_Management] malloc_trim(0) released memory -[MultiGPU Model Management] soft_empty_cache_multigpu: devices to clear = ['cpu', 'cuda:0', 'cuda:1'] -[MultiGPU Model Management] Clearing CUDA cache on cuda:0 (idx=0) -[MultiGPU Model Management] 2025-09-28T16:55:25.084Z general_pre-empty:cuda:0 cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] Cleared CUDA cache (and IPC if available) on cuda:0 -[MultiGPU Model Management] 2025-09-28T16:55:25.086Z general_post-empty:cuda:0 cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] Clearing CUDA cache on cuda:1 (idx=1) -[MultiGPU Model Management] 2025-09-28T16:55:25.086Z general_pre-empty:cuda:1 cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] Cleared CUDA cache (and IPC if available) on cuda:1 -[MultiGPU Model Management] 2025-09-28T16:55:25.087Z general_post-empty:cuda:1 cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:25.087Z general_post-soft-empty cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] Force flag active: triggering executor cache reset (CPU) -[MultiGPU Model Management] 2025-09-28T16:55:25.088Z executor_reset_pre-trigger (forced_soft_empty) cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU_Memory_Management] Triggering PromptExecutor cache reset. Reason: forced_soft_empty -[MultiGPU_Leak_Analyzer] High pressure detected: patchers=27, cpu_mem=49.3%. Analyzing referrers. -[MultiGPU_Leak_Analyzer] Patcher #0 id=137954885156944 referrers=4 - Ref 0: list(len=27) mod=builtins - Ref 1: list(len=5) mod=builtins - Ref 2: tuple mod=builtins - Ref 3: VAE mod=comfy.sd -[MultiGPU_Leak_Analyzer] Patcher #1 id=137952743399760 referrers=5 - Ref 0: list(len=27) mod=builtins - Ref 1: list(len=5) mod=builtins - Ref 2: CLIP mod=comfy.sd - Ref 3: list(len=1) mod=builtins - Ref 4: set mod=builtins -[MultiGPU_Leak_Analyzer] Patcher #2 id=137949390436688 referrers=4 - Ref 0: list(len=27) mod=builtins - Ref 1: list(len=5) mod=builtins - Ref 2: tuple mod=builtins - Ref 3: CLIP mod=comfy.sd -[MultiGPU_Leak_Analyzer] Patcher #3 id=137954821603840 referrers=3 - Ref 0: list(len=27) mod=builtins - Ref 1: list(len=5) mod=builtins - Ref 2: VAE mod=comfy.sd -[MultiGPU_Leak_Analyzer] Patcher #4 id=137952743397888 referrers=4 - Ref 0: list(len=27) mod=builtins - Ref 1: list(len=5) mod=builtins - Ref 2: tuple mod=builtins - Ref 3: VAE mod=comfy.sd -[MultiGPU Model Management] 2025-09-28T16:55:25.299Z distorch_prune_start cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] [PRUNE_DEBUG] Starting prune - current_loaded_models count: 9 -[MultiGPU Model Management] [PRUNE_DEBUG] Model 0: SDXLClipModel, keep_loaded=False, hash=cd38a1f8, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 1: AutoencodingEngine, keep_loaded=False, hash=be20bc8e, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 2: Flux, keep_loaded=False, hash=d78780cf, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 3: FluxClipModel_, keep_loaded=False, hash=12b44a2d, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 4: BaseModel, keep_loaded=False, hash=8b7f0c0a, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 5: Flux, keep_loaded=False, hash=72754c29, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 6: FluxClipModel_, keep_loaded=False, hash=7190f578, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 7: QwenImageTEModel_, keep_loaded=False, hash=9b313bd7, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 8: QwenImageTEModel_, keep_loaded=False, hash=34f9ab09, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Active hashes V2: 9, Store has: 0 -[MultiGPU Model Management] [PRUNE_DEBUG] No stale allocation entries to prune -[MultiGPU Model Management] [PRUNE_DEBUG] No stale settings entries to prune -[MultiGPU Model Management] [PRUNE_DEBUG] After pruning - V2 allocation store has: 0 entries -[MultiGPU Model Management] 2025-09-28T16:55:25.346Z distorch_prune_end cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:25.347Z mem_mgmt_pre-history-clear cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:25.347Z mem_mgmt_post-history-clear cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU_Memory_Management] malloc_trim(0) begin -[MultiGPU Model Management] 2025-09-28T16:55:25.348Z mem_mgmt_pre-malloc-trim cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:25.349Z mem_mgmt_post-malloc-trim cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU_Memory_Management] malloc_trim(0) released memory -[MultiGPU Model Management] 2025-09-28T16:55:25.350Z executor_reset_post-trigger (forced_soft_empty) cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:25.350Z patched_soft_empty_end cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:25.351Z patched_load_models_gpu_pre-original-call cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] [MultiGPU_LoadedModel_Patch] Set base model Patcher 137952743399760. -[MultiGPU Model Management] 2025-09-28T16:55:25.353Z safetensor:cd38a1f8_pre-load cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU_DisTorch2_CLIP] CLIP Compute Device: cuda:1 -[MultiGPU_DisTorch2_CLIP] Expert String Examples: - Direct(byte) Mode - cuda:0,500mb;cuda:1,3.0g;cpu,5gb* -> '*' cpu = over/underflow device, put 0.50gb on cuda0, 3.00gb on cuda1, and 5.00gb (or the rest) on cpu - Ratio(%) Mode - cuda:0,8%;cuda:1,8%;cpu,4% -> 8:8:4 ratio, put 40% on cuda0, 40% on cuda1, and 20% on cpu -=============================================== - DisTorch2 Model Virtual VRAM Analysis -=============================================== -Object Role Original(GB) Total(GB) Virt(GB) ------------------------------------------------ -cuda:1 recip 23.56GB 25.56GB +2.00GB -cpu donor 93.98GB 91.98GB -2.00GB ------------------------------------------------ -model model 1.52GB 0.00GB -2.00GB -[MultiGPU_DisTorch2_CLIP] Final CLIP Allocation String: -cuda:1,0.0000;cpu,0.0213;cuda:0,0.0 -================================================== - DisTorch2 CLIP Model Device Allocations -================================================== -Device VRAM GB Dev % Model GB Dist % --------------------------------------------------- -cuda:0 23.56 0.0% 0.00 0.0% -cuda:1 23.56 0.0% 0.00 0.0% -cpu 93.98 2.1% 2.00 100.0% --------------------------------------------------- - DisTorch2 CLIP Model Layer Distribution --------------------------------------------------- -Layer Type Layers Memory (MB) % Total --------------------------------------------------- -Embedding 4 193.30 12.4% -LayerNorm 90 0.39 0.0% -Linear 266 1367.11 87.6% --------------------------------------------------- -[MultiGPU_DisTorch2_CLIP] Preserving 4 head layer(s) (193.30 MB) on compute device: cuda:1 -DisTorch2 CLIP Model Final Device/Layer Assignments --------------------------------------------------- -Device Layers Memory (MB) % Total --------------------------------------------------- -cuda:1 94 193.69 12.4% -cpu 266 1367.11 87.6% --------------------------------------------------- -[MultiGPU DisTorch V2] DisTorch loading completed. -[MultiGPU DisTorch V2] Total memory: 1560.80MB -[MultiGPU Model Management] 2025-09-28T16:55:25.367Z safetensor:cd38a1f8_post-load cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:25.367Z patched_load_models_gpu_post-original-call cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:25.670Z patched_load_models_gpu_start cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] Incoming models summary: SDXL:4.78GB req on cuda:0 -[MultiGPU Model Management] Non-Zero incoming DisTorch2 model detected. Initiating proactive unload. -[MultiGPU Model Management] Need calc on cuda:0: effective_needed=1.10GB, free_now=3.72GB, need_bytes=0.00GB -[MultiGPU Model Management] No unloads; 25% torch-cache rule triggered on: cpu. Calling soft_empty_cache() -[MultiGPU Model Management] 2025-09-28T16:55:25.678Z patched_soft_empty_start:force=True cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:25.678Z distorch_prune_start cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] [PRUNE_DEBUG] Starting prune - current_loaded_models count: 9 -[MultiGPU Model Management] [PRUNE_DEBUG] Model 0: SDXLClipModel, keep_loaded=False, hash=cd38a1f8, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 1: AutoencodingEngine, keep_loaded=False, hash=be20bc8e, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 2: Flux, keep_loaded=False, hash=d78780cf, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 3: FluxClipModel_, keep_loaded=False, hash=12b44a2d, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 4: BaseModel, keep_loaded=False, hash=8b7f0c0a, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 5: Flux, keep_loaded=False, hash=72754c29, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 6: FluxClipModel_, keep_loaded=False, hash=7190f578, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 7: QwenImageTEModel_, keep_loaded=False, hash=9b313bd7, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 8: QwenImageTEModel_, keep_loaded=False, hash=34f9ab09, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Active hashes V2: 9, Store has: 0 -[MultiGPU Model Management] [PRUNE_DEBUG] No stale allocation entries to prune -[MultiGPU Model Management] [PRUNE_DEBUG] No stale settings entries to prune -[MultiGPU Model Management] [PRUNE_DEBUG] After pruning - V2 allocation store has: 0 entries -[MultiGPU Model Management] 2025-09-28T16:55:25.725Z distorch_prune_end cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] [DETECT_DEBUG] Checking DisTorch2 active status - loaded models: 9, store entries: 4 -[MultiGPU Model Management] [DETECT_DEBUG] Model 0: SDXLClipModel, hash=cd38a1f8, in_store=True, alloc_value='#cuda:1;2.0;cpu', keep_loaded=False -[MultiGPU Model Management] [DETECT_DEBUG] DisTorch2 ACTIVE detected on model: SDXLClipModel -[MultiGPU Model Management] [DETECT_DEBUG] Final DisTorch2 active status: True -[MultiGPU Model Management] DisTorch2 active: clearing allocator caches on all devices (VRAM) -[MultiGPU Model Management] soft_empty_cache_multigpu: starting GC and multi-device cache clear -[MultiGPU Model Management] 2025-09-28T16:55:25.727Z general_pre-soft-empty cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:25.728Z general_pre-gc cpu|45.38 cuda:0|19.83 cuda:1|19.08 -[MultiGPU_Lifecycle] [pre-gc] Tracked ModelPatchers=27, approx CPU RAM=64192.58 MB -[MultiGPU_Lifecycle] [post-gc] Tracked ModelPatchers=27, approx CPU RAM=64192.58 MB -[MultiGPU Model Management] 2025-09-28T16:55:25.997Z general_post-gc cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] soft_empty_cache_multigpu: garbage collection complete -[MultiGPU_Memory_Management] malloc_trim(0) begin -[MultiGPU Model Management] 2025-09-28T16:55:25.998Z mem_mgmt_pre-malloc-trim cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:25.999Z mem_mgmt_post-malloc-trim cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU_Memory_Management] malloc_trim(0) released memory -[MultiGPU Model Management] soft_empty_cache_multigpu: devices to clear = ['cpu', 'cuda:0', 'cuda:1'] -[MultiGPU Model Management] Clearing CUDA cache on cuda:0 (idx=0) -[MultiGPU Model Management] 2025-09-28T16:55:26.000Z general_pre-empty:cuda:0 cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] Cleared CUDA cache (and IPC if available) on cuda:0 -[MultiGPU Model Management] 2025-09-28T16:55:26.002Z general_post-empty:cuda:0 cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] Clearing CUDA cache on cuda:1 (idx=1) -[MultiGPU Model Management] 2025-09-28T16:55:26.002Z general_pre-empty:cuda:1 cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] Cleared CUDA cache (and IPC if available) on cuda:1 -[MultiGPU Model Management] 2025-09-28T16:55:26.003Z general_post-empty:cuda:1 cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:26.003Z general_post-soft-empty cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] Force flag active: triggering executor cache reset (CPU) -[MultiGPU Model Management] 2025-09-28T16:55:26.004Z executor_reset_pre-trigger (forced_soft_empty) cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU_Memory_Management] Triggering PromptExecutor cache reset. Reason: forced_soft_empty -[MultiGPU_Leak_Analyzer] High pressure detected: patchers=27, cpu_mem=49.3%. Analyzing referrers. -[MultiGPU_Leak_Analyzer] Patcher #0 id=137954885156944 referrers=4 - Ref 0: list(len=27) mod=builtins - Ref 1: list(len=5) mod=builtins - Ref 2: tuple mod=builtins - Ref 3: VAE mod=comfy.sd -[MultiGPU_Leak_Analyzer] Patcher #1 id=137952743399760 referrers=3 - Ref 0: list(len=27) mod=builtins - Ref 1: list(len=5) mod=builtins - Ref 2: CLIP mod=comfy.sd -[MultiGPU_Leak_Analyzer] Patcher #2 id=137949390436688 referrers=4 - Ref 0: list(len=27) mod=builtins - Ref 1: list(len=5) mod=builtins - Ref 2: tuple mod=builtins - Ref 3: CLIP mod=comfy.sd -[MultiGPU_Leak_Analyzer] Patcher #3 id=137954821603840 referrers=3 - Ref 0: list(len=27) mod=builtins - Ref 1: list(len=5) mod=builtins - Ref 2: VAE mod=comfy.sd -[MultiGPU_Leak_Analyzer] Patcher #4 id=137952743397888 referrers=4 - Ref 0: list(len=27) mod=builtins - Ref 1: list(len=5) mod=builtins - Ref 2: tuple mod=builtins - Ref 3: VAE mod=comfy.sd -[MultiGPU Model Management] 2025-09-28T16:55:26.180Z distorch_prune_start cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] [PRUNE_DEBUG] Starting prune - current_loaded_models count: 9 -[MultiGPU Model Management] [PRUNE_DEBUG] Model 0: SDXLClipModel, keep_loaded=False, hash=cd38a1f8, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 1: AutoencodingEngine, keep_loaded=False, hash=be20bc8e, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 2: Flux, keep_loaded=False, hash=d78780cf, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 3: FluxClipModel_, keep_loaded=False, hash=12b44a2d, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 4: BaseModel, keep_loaded=False, hash=8b7f0c0a, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 5: Flux, keep_loaded=False, hash=72754c29, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 6: FluxClipModel_, keep_loaded=False, hash=7190f578, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 7: QwenImageTEModel_, keep_loaded=False, hash=9b313bd7, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 8: QwenImageTEModel_, keep_loaded=False, hash=34f9ab09, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Active hashes V2: 9, Store has: 0 -[MultiGPU Model Management] [PRUNE_DEBUG] No stale allocation entries to prune -[MultiGPU Model Management] [PRUNE_DEBUG] No stale settings entries to prune -[MultiGPU Model Management] [PRUNE_DEBUG] After pruning - V2 allocation store has: 0 entries -[MultiGPU Model Management] 2025-09-28T16:55:26.227Z distorch_prune_end cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:26.227Z mem_mgmt_pre-history-clear cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:26.228Z mem_mgmt_post-history-clear cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU_Memory_Management] malloc_trim(0) begin -[MultiGPU Model Management] 2025-09-28T16:55:26.228Z mem_mgmt_pre-malloc-trim cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:26.230Z mem_mgmt_post-malloc-trim cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU_Memory_Management] malloc_trim(0) released memory -[MultiGPU Model Management] 2025-09-28T16:55:26.230Z executor_reset_post-trigger (forced_soft_empty) cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:26.231Z patched_soft_empty_end cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:26.231Z patched_load_models_gpu_pre-original-call cpu|45.37 cuda:0|19.83 cuda:1|19.08 -[MultiGPU Model Management] [MultiGPU_LoadedModel_Patch] Set base model Patcher 137952610352304. -Requested to load SDXL -[MultiGPU Model Management] 2025-09-28T16:55:26.456Z patched_soft_empty_start:force=False cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:26.456Z distorch_prune_start cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU Model Management] [PRUNE_DEBUG] Starting prune - current_loaded_models count: 8 -[MultiGPU Model Management] [PRUNE_DEBUG] Model 0: SDXLClipModel, keep_loaded=False, hash=cd38a1f8, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 1: AutoencodingEngine, keep_loaded=False, hash=be20bc8e, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 2: Flux, keep_loaded=False, hash=d78780cf, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 3: FluxClipModel_, keep_loaded=False, hash=12b44a2d, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 4: Flux, keep_loaded=False, hash=72754c29, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 5: FluxClipModel_, keep_loaded=False, hash=7190f578, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 6: QwenImageTEModel_, keep_loaded=False, hash=9b313bd7, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 7: QwenImageTEModel_, keep_loaded=False, hash=34f9ab09, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Active hashes V2: 8, Store has: 0 -[MultiGPU Model Management] [PRUNE_DEBUG] No stale allocation entries to prune -[MultiGPU Model Management] [PRUNE_DEBUG] No stale settings entries to prune -[MultiGPU Model Management] [PRUNE_DEBUG] After pruning - V2 allocation store has: 0 entries -[MultiGPU Model Management] 2025-09-28T16:55:26.500Z distorch_prune_end cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU Model Management] [DETECT_DEBUG] Checking DisTorch2 active status - loaded models: 8, store entries: 4 -[MultiGPU Model Management] [DETECT_DEBUG] Model 0: SDXLClipModel, hash=cd38a1f8, in_store=True, alloc_value='#cuda:1;2.0;cpu', keep_loaded=False -[MultiGPU Model Management] [DETECT_DEBUG] DisTorch2 ACTIVE detected on model: SDXLClipModel -[MultiGPU Model Management] [DETECT_DEBUG] Final DisTorch2 active status: True -[MultiGPU Model Management] DisTorch2 active: clearing allocator caches on all devices (VRAM) -[MultiGPU Model Management] soft_empty_cache_multigpu: starting GC and multi-device cache clear -[MultiGPU Model Management] 2025-09-28T16:55:26.502Z general_pre-soft-empty cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:26.503Z general_pre-gc cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU_Lifecycle] [pre-gc] Tracked ModelPatchers=27, approx CPU RAM=64811.94 MB -[MultiGPU_Lifecycle] [post-gc] Tracked ModelPatchers=27, approx CPU RAM=64811.94 MB -[MultiGPU Model Management] 2025-09-28T16:55:26.769Z general_post-gc cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU Model Management] soft_empty_cache_multigpu: garbage collection complete -[MultiGPU_Memory_Management] malloc_trim(0) begin -[MultiGPU Model Management] 2025-09-28T16:55:26.770Z mem_mgmt_pre-malloc-trim cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:26.772Z mem_mgmt_post-malloc-trim cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU_Memory_Management] malloc_trim(0) released memory -[MultiGPU Model Management] soft_empty_cache_multigpu: devices to clear = ['cpu', 'cuda:0', 'cuda:1'] -[MultiGPU Model Management] Clearing CUDA cache on cuda:0 (idx=0) -[MultiGPU Model Management] 2025-09-28T16:55:26.772Z general_pre-empty:cuda:0 cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU Model Management] Cleared CUDA cache (and IPC if available) on cuda:0 -[MultiGPU Model Management] 2025-09-28T16:55:26.788Z general_post-empty:cuda:0 cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU Model Management] Clearing CUDA cache on cuda:1 (idx=1) -[MultiGPU Model Management] 2025-09-28T16:55:26.788Z general_pre-empty:cuda:1 cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU Model Management] Cleared CUDA cache (and IPC if available) on cuda:1 -[MultiGPU Model Management] 2025-09-28T16:55:26.789Z general_post-empty:cuda:1 cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:26.789Z general_post-soft-empty cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:26.790Z patched_soft_empty_end cpu|45.37 cuda:0|19.23 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:26.795Z safetensor:5d907277_pre-load cpu|45.37 cuda:0|19.23 cuda:1|19.08 -=============================================== - DisTorch2 Model Virtual VRAM Analysis -=============================================== -Object Role Original(GB) Total(GB) Virt(GB) ------------------------------------------------ -cuda:0 recip 23.56GB 24.66GB +1.10GB -cpu donor 93.98GB 92.88GB -1.10GB ------------------------------------------------ -model model 4.78GB 3.68GB -1.10GB -================================================== -[MultiGPU DisTorch V2] Final Allocation String: -cuda:0,0.1563;cpu,0.0117;cuda:1,0.0 -================================================== - DisTorch2 Model Device Allocations -================================================== -Device VRAM GB Dev % Model GB Dist % --------------------------------------------------- -cuda:0 23.56 15.6% 3.68 77.0% -cuda:1 23.56 0.0% 0.00 0.0% -cpu 93.98 1.2% 1.10 23.0% --------------------------------------------------- - DisTorch2 Model Layer Distribution --------------------------------------------------- -Layer Type Layers Memory (MB) % Total --------------------------------------------------- -Linear 743 4260.26 87.0% -Conv2d 51 635.67 13.0% -GroupNorm 46 0.17 0.0% -LayerNorm 210 0.95 0.0% --------------------------------------------------- -DisTorch2 Model Final Device/Layer Assignments --------------------------------------------------- -Device Layers Memory (MB) % Total --------------------------------------------------- -cuda:0 (<0.01%) 261 2.34 0.0% -cuda:0 584 3769.60 77.0% -cpu 205 1125.10 23.0% --------------------------------------------------- -[MultiGPU DisTorch V2] DisTorch loading completed. -[MultiGPU DisTorch V2] Total memory: 4897.05MB -[MultiGPU Model Management] 2025-09-28T16:55:28.156Z safetensor:5d907277_post-load cpu|44.34 cuda:0|22.91 cuda:1|19.08 -[MultiGPU Model Management] 2025-09-28T16:55:28.157Z patched_load_models_gpu_post-original-call cpu|44.34 cuda:0|22.91 cuda:1|19.08 -100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [00:13<00:00, 1.49it/s] -[MultiGPU Model Management] 2025-09-28T16:55:41.636Z patched_load_models_gpu_start cpu|44.29 cuda:0|22.91 cuda:1|19.08 -[MultiGPU Model Management] Incoming models summary: AutoencoderKL:0.16GB req on cuda:1 -[MultiGPU Model Management] 2025-09-28T16:55:41.638Z patched_load_models_gpu_pre-original-call cpu|44.29 cuda:0|22.91 cuda:1|19.08 -[MultiGPU Model Management] [MultiGPU_LoadedModel_Patch] Set base model Patcher 137952743397888. -Requested to load AutoencoderKL -[MultiGPU Model Management] 2025-09-28T16:55:41.755Z patched_soft_empty_start:force=False cpu|44.49 cuda:0|22.91 cuda:1|18.74 -[MultiGPU Model Management] 2025-09-28T16:55:41.755Z distorch_prune_start cpu|44.49 cuda:0|22.91 cuda:1|18.74 -[MultiGPU Model Management] [PRUNE_DEBUG] Starting prune - current_loaded_models count: 7 -[MultiGPU Model Management] [PRUNE_DEBUG] Model 0: SDXL, keep_loaded=False, hash=5d907277, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 1: AutoencodingEngine, keep_loaded=False, hash=be20bc8e, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 2: Flux, keep_loaded=False, hash=d78780cf, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 3: FluxClipModel_, keep_loaded=False, hash=12b44a2d, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 4: Flux, keep_loaded=False, hash=72754c29, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 5: QwenImageTEModel_, keep_loaded=False, hash=9b313bd7, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 6: QwenImageTEModel_, keep_loaded=False, hash=34f9ab09, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Active hashes V2: 7, Store has: 0 -[MultiGPU Model Management] [PRUNE_DEBUG] No stale allocation entries to prune -[MultiGPU Model Management] [PRUNE_DEBUG] No stale settings entries to prune -[MultiGPU Model Management] [PRUNE_DEBUG] After pruning - V2 allocation store has: 0 entries -[MultiGPU Model Management] 2025-09-28T16:55:41.804Z distorch_prune_end cpu|44.49 cuda:0|22.91 cuda:1|18.74 -[MultiGPU Model Management] [DETECT_DEBUG] Checking DisTorch2 active status - loaded models: 7, store entries: 4 -[MultiGPU Model Management] [DETECT_DEBUG] Model 0: SDXL, hash=5d907277, in_store=True, alloc_value='#cuda:0;1.1;cpu', keep_loaded=False -[MultiGPU Model Management] [DETECT_DEBUG] DisTorch2 ACTIVE detected on model: SDXL -[MultiGPU Model Management] [DETECT_DEBUG] Final DisTorch2 active status: True -[MultiGPU Model Management] DisTorch2 active: clearing allocator caches on all devices (VRAM) -[MultiGPU Model Management] soft_empty_cache_multigpu: starting GC and multi-device cache clear -[MultiGPU Model Management] 2025-09-28T16:55:41.809Z general_pre-soft-empty cpu|44.49 cuda:0|22.91 cuda:1|18.74 -[MultiGPU Model Management] 2025-09-28T16:55:41.810Z general_pre-gc cpu|44.49 cuda:0|22.91 cuda:1|18.74 -[MultiGPU_Lifecycle] [pre-gc] Tracked ModelPatchers=27, approx CPU RAM=61393.56 MB -[MultiGPU_Lifecycle] [post-gc] Tracked ModelPatchers=27, approx CPU RAM=61393.56 MB -[MultiGPU Model Management] 2025-09-28T16:55:42.078Z general_post-gc cpu|44.48 cuda:0|22.91 cuda:1|18.74 -[MultiGPU Model Management] soft_empty_cache_multigpu: garbage collection complete -[MultiGPU_Memory_Management] malloc_trim(0) begin -[MultiGPU Model Management] 2025-09-28T16:55:42.079Z mem_mgmt_pre-malloc-trim cpu|44.48 cuda:0|22.91 cuda:1|18.74 -[MultiGPU Model Management] 2025-09-28T16:55:42.136Z mem_mgmt_post-malloc-trim cpu|42.20 cuda:0|22.91 cuda:1|18.74 -[MultiGPU_Memory_Management] malloc_trim(0) released memory -[MultiGPU Model Management] soft_empty_cache_multigpu: devices to clear = ['cpu', 'cuda:0', 'cuda:1'] -[MultiGPU Model Management] Clearing CUDA cache on cuda:0 (idx=0) -[MultiGPU Model Management] 2025-09-28T16:55:42.136Z general_pre-empty:cuda:0 cpu|42.20 cuda:0|22.91 cuda:1|18.74 -[MultiGPU Model Management] Cleared CUDA cache (and IPC if available) on cuda:0 -[MultiGPU Model Management] 2025-09-28T16:55:42.149Z general_post-empty:cuda:0 cpu|42.20 cuda:0|22.91 cuda:1|18.74 -[MultiGPU Model Management] Clearing CUDA cache on cuda:1 (idx=1) -[MultiGPU Model Management] 2025-09-28T16:55:42.150Z general_pre-empty:cuda:1 cpu|42.20 cuda:0|22.91 cuda:1|18.74 -[MultiGPU Model Management] Cleared CUDA cache (and IPC if available) on cuda:1 -[MultiGPU Model Management] 2025-09-28T16:55:42.150Z general_post-empty:cuda:1 cpu|42.20 cuda:0|22.91 cuda:1|18.74 -[MultiGPU Model Management] 2025-09-28T16:55:42.151Z general_post-soft-empty cpu|42.20 cuda:0|22.91 cuda:1|18.74 -[MultiGPU Model Management] 2025-09-28T16:55:42.151Z patched_soft_empty_end cpu|42.20 cuda:0|22.91 cuda:1|18.74 -[MultiGPU Model Management] 2025-09-28T16:55:42.153Z safetensor:626f5bc4_pre-load cpu|42.20 cuda:0|22.91 cuda:1|18.74 -loaded completely 179.03548431396484 159.55708122253418 True -[MultiGPU Model Management] 2025-09-28T16:55:42.195Z safetensor:626f5bc4_post-load cpu|42.20 cuda:0|22.91 cuda:1|18.90 -[MultiGPU Model Management] 2025-09-28T16:55:42.195Z patched_load_models_gpu_post-original-call cpu|42.20 cuda:0|22.91 cuda:1|18.90 -Prompt executed in 327.39 seconds -[MultiGPU Model Management] [UNLOAD_DEBUG] Patched unload_all_models called - initial model count: 8 -[MultiGPU Model Management] [UNLOAD_DEBUG] Model 0: AutoencoderKL, keep_loaded=False -[MultiGPU Model Management] [UNLOAD_DEBUG] Adding to kept_models: AutoencoderKL -[MultiGPU Model Management] [GC_ANCHOR] Added retention anchor for AutoencoderKL, reason: keep_loaded_test, total anchors: 1 -[MultiGPU Model Management] [UNLOAD_DEBUG] Model 1: SDXL, keep_loaded=False -[MultiGPU Model Management] [UNLOAD_DEBUG] Adding to kept_models: SDXL -[MultiGPU Model Management] [GC_ANCHOR] Added retention anchor for SDXL, reason: keep_loaded_test, total anchors: 2 -[MultiGPU Model Management] [UNLOAD_DEBUG] Model 2: AutoencodingEngine, keep_loaded=False -[MultiGPU Model Management] [UNLOAD_DEBUG] Adding to kept_models: AutoencodingEngine -[MultiGPU Model Management] [GC_ANCHOR] Added retention anchor for AutoencodingEngine, reason: keep_loaded_test, total anchors: 3 -[MultiGPU Model Management] [UNLOAD_DEBUG] Model 3: Flux, keep_loaded=False -[MultiGPU Model Management] [UNLOAD_DEBUG] Adding to kept_models: Flux -[MultiGPU Model Management] [GC_ANCHOR] Added retention anchor for Flux, reason: keep_loaded_test, total anchors: 4 -[MultiGPU Model Management] [UNLOAD_DEBUG] Model 4: FluxClipModel_, keep_loaded=False -[MultiGPU Model Management] [UNLOAD_DEBUG] Adding to kept_models: FluxClipModel_ -[MultiGPU Model Management] [GC_ANCHOR] Added retention anchor for FluxClipModel_, reason: keep_loaded_test, total anchors: 5 -[MultiGPU Model Management] [UNLOAD_DEBUG] Model 5: Flux, keep_loaded=False -[MultiGPU Model Management] [UNLOAD_DEBUG] Adding to kept_models: Flux -[MultiGPU Model Management] [GC_ANCHOR] Added retention anchor for Flux, reason: keep_loaded_test, total anchors: 6 -[MultiGPU Model Management] [UNLOAD_DEBUG] Model 6: QwenImageTEModel_, keep_loaded=False -[MultiGPU Model Management] [UNLOAD_DEBUG] Adding to kept_models: QwenImageTEModel_ -[MultiGPU Model Management] [GC_ANCHOR] Added retention anchor for QwenImageTEModel_, reason: keep_loaded_test, total anchors: 7 -[MultiGPU Model Management] [UNLOAD_DEBUG] Model 7: QwenImageTEModel_, keep_loaded=False -[MultiGPU Model Management] [UNLOAD_DEBUG] Adding to kept_models: QwenImageTEModel_ -[MultiGPU Model Management] [GC_ANCHOR] Added retention anchor for QwenImageTEModel_, reason: keep_loaded_test, total anchors: 8 -[MultiGPU Model Management] [UNLOAD_DEBUG] Final counts - kept_models: 8, models_to_unload: 0 -[MultiGPU Model Management] Found 8 model(s) to retain, unloading 0 model(s) -[MultiGPU Model Management] [UNLOAD_DEBUG] Updated mm.current_loaded_models, new count: 8 -[MultiGPU Model Management] Successfully retained 8 model(s) during unload -[MultiGPU Model Management] [MultiGPU_LoadedModel_Patch] Clone Patcher 137950125562944 GC'd. LoadedModel already gone or missing _switch_parent. -[MultiGPU Model Management] 2025-09-28T16:55:43.520Z patched_soft_empty_start:force=False cpu|23.56 cuda:0|22.91 cuda:1|18.90 -[MultiGPU Model Management] 2025-09-28T16:55:43.524Z distorch_prune_start cpu|23.56 cuda:0|22.91 cuda:1|18.90 -[MultiGPU Model Management] [PRUNE_DEBUG] Starting prune - current_loaded_models count: 8 -[MultiGPU Model Management] [PRUNE_DEBUG] Model 0: AutoencoderKL, keep_loaded=False, hash=626f5bc4, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 1: SDXL, keep_loaded=False, hash=5d907277, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 2: AutoencodingEngine, keep_loaded=False, hash=be20bc8e, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 3: Flux, keep_loaded=False, hash=d78780cf, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 4: FluxClipModel_, keep_loaded=False, hash=12b44a2d, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 5: Flux, keep_loaded=False, hash=72754c29, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 6: QwenImageTEModel_, keep_loaded=False, hash=9b313bd7, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Model 7: QwenImageTEModel_, keep_loaded=False, hash=34f9ab09, has_v2_allocation=False -[MultiGPU Model Management] [PRUNE_DEBUG] Active hashes V2: 8, Store has: 0 -[MultiGPU Model Management] [PRUNE_DEBUG] No stale allocation entries to prune -[MultiGPU Model Management] [PRUNE_DEBUG] No stale settings entries to prune -[MultiGPU Model Management] [PRUNE_DEBUG] After pruning - V2 allocation store has: 0 entries -[MultiGPU Model Management] 2025-09-28T16:55:43.598Z distorch_prune_end cpu|23.56 cuda:0|22.91 cuda:1|18.90 -[MultiGPU Model Management] [DETECT_DEBUG] Checking DisTorch2 active status - loaded models: 8, store entries: 4 -[MultiGPU Model Management] [DETECT_DEBUG] Model 0: AutoencoderKL, hash=626f5bc4, in_store=False, alloc_value='', keep_loaded=False -[MultiGPU Model Management] [DETECT_DEBUG] Model 1: SDXL, hash=5d907277, in_store=True, alloc_value='#cuda:0;1.1;cpu', keep_loaded=False -[MultiGPU Model Management] [DETECT_DEBUG] DisTorch2 ACTIVE detected on model: SDXL -[MultiGPU Model Management] [DETECT_DEBUG] Final DisTorch2 active status: True -[MultiGPU Model Management] DisTorch2 active: clearing allocator caches on all devices (VRAM) -[MultiGPU Model Management] soft_empty_cache_multigpu: starting GC and multi-device cache clear -[MultiGPU Model Management] 2025-09-28T16:55:43.607Z general_pre-soft-empty cpu|23.56 cuda:0|22.91 cuda:1|18.90 -[MultiGPU Model Management] 2025-09-28T16:55:43.608Z general_pre-gc cpu|23.56 cuda:0|22.91 cuda:1|18.90 -[MultiGPU_Lifecycle] [pre-gc] Tracked ModelPatchers=13, approx CPU RAM=9581.31 MB -[MultiGPU_Lifecycle] [post-gc] Tracked ModelPatchers=13, approx CPU RAM=9581.31 MB -[MultiGPU Model Management] 2025-09-28T16:55:43.929Z general_post-gc cpu|23.54 cuda:0|22.91 cuda:1|18.90 -[MultiGPU Model Management] soft_empty_cache_multigpu: garbage collection complete -[MultiGPU_Memory_Management] malloc_trim(0) begin -[MultiGPU Model Management] 2025-09-28T16:55:43.931Z mem_mgmt_pre-malloc-trim cpu|23.54 cuda:0|22.91 cuda:1|18.90 -[MultiGPU Model Management] 2025-09-28T16:55:44.319Z mem_mgmt_post-malloc-trim cpu|10.43 cuda:0|22.91 cuda:1|18.90 -[MultiGPU_Memory_Management] malloc_trim(0) released memory -[MultiGPU Model Management] soft_empty_cache_multigpu: devices to clear = ['cpu', 'cuda:0', 'cuda:1'] -[MultiGPU Model Management] Clearing CUDA cache on cuda:0 (idx=0) -[MultiGPU Model Management] 2025-09-28T16:55:44.320Z general_pre-empty:cuda:0 cpu|10.43 cuda:0|22.91 cuda:1|18.90 -[MultiGPU Model Management] Cleared CUDA cache (and IPC if available) on cuda:0 -[MultiGPU Model Management] 2025-09-28T16:55:44.391Z general_post-empty:cuda:0 cpu|10.41 cuda:0|22.91 cuda:1|18.90 -[MultiGPU Model Management] Clearing CUDA cache on cuda:1 (idx=1) -[MultiGPU Model Management] 2025-09-28T16:55:44.391Z general_pre-empty:cuda:1 cpu|10.41 cuda:0|22.91 cuda:1|18.90 -[MultiGPU Model Management] Cleared CUDA cache (and IPC if available) on cuda:1 -[MultiGPU Model Management] 2025-09-28T16:55:44.392Z general_post-empty:cuda:1 cpu|10.41 cuda:0|22.91 cuda:1|18.90 -[MultiGPU Model Management] 2025-09-28T16:55:44.392Z general_post-soft-empty cpu|10.41 cuda:0|22.91 cuda:1|18.90 -[MultiGPU Model Management] 2025-09-28T16:55:44.393Z patched_soft_empty_end cpu|10.41 cuda:0|22.91 cuda:1|18.90 diff --git a/memory-bank/reference-logs/d2_testing_stripped.json b/memory-bank/reference-logs/d2_testing_stripped.json deleted file mode 100644 index a7b707d..0000000 --- a/memory-bank/reference-logs/d2_testing_stripped.json +++ /dev/null @@ -1,171 +0,0 @@ -{ - "10": { - "type": "CheckpointLoaderSimpleDisTorch2MultiGPU", - "widgets_values": [ - "safetensor_testing/realDream_15SD15.safetensors", - "cuda:0", - 1, - "cpu", - "", - false - ] - }, - "17": { - "type": "CheckpointLoaderAdvancedDisTorch2MultiGPU", - "widgets_values": [ - "Juggernaut-XL_v9_RunDiffusionPhoto_v2.safetensors", - "cuda:0", - 1.1, - "cpu", - "cuda:1", - 2, - "cpu", - "cuda:1", - "", - "", - false - ] - }, - "29": { - "type": "CheckpointLoaderAdvancedMultiGPU", - "widgets_values": [ - "safetensor_testing/realisticVisionV60B1_v51VAE.safetensors", - "cuda:0", - "cuda:1", - "cuda:1" - ] - }, - "30": { - "type": "CheckpointLoaderSimpleMultiGPU", - "widgets_values": [ - "safetensor_testing/realDream_15SD15.safetensors", - "cuda:0" - ] - }, - "40": { - "type": "UNETLoaderDisTorch2MultiGPU", - "widgets_values": [ - "qwen_image_fp8_e4m3fn.safetensors", - "fp8_e4m3fn", - "cuda:0", - 16, - "cpu", - "", - false - ] - }, - "41": { - "type": "VAELoaderMultiGPU", - "widgets_values": [ - "qwen_image_vae.safetensors", - "cuda:1" - ] - }, - "42": { - "type": "CLIPLoaderMultiGPU", - "widgets_values": [ - "qwen_2.5_vl_7b_fp8_scaled.safetensors", - "qwen_image", - "cuda:1" - ] - }, - "53": { - "type": "UNETLoaderMultiGPU", - "widgets_values": [ - "flux1-dev-fp8.safetensors", - "default", - "cuda:0" - ] - }, - "54": { - "type": "VAELoaderMultiGPU", - "widgets_values": [ - "ae.safetensors", - "cuda:1" - ] - }, - "55": { - "type": "DualCLIPLoaderMultiGPU", - "widgets_values": [ - "t5xxl_fp8_e4m3fn.safetensors", - "clip_l.safetensors", - "flux", - "cuda:1" - ] - }, - "72": { - "type": "UnetLoaderGGUFMultiGPU", - "widgets_values": [ - "flux1-dev-Q8_0.gguf", - "cuda:0" - ] - }, - "73": { - "type": "DualCLIPLoaderGGUFMultiGPU", - "widgets_values": [ - "t5-v1_1-xxl-encoder-Q8_0.gguf", - "clip_l.safetensors", - "flux", - "cuda:1" - ] - }, - "88": { - "type": "CLIPLoaderGGUFMultiGPU", - "widgets_values": [ - "Qwen2.5-VL-7B-Instruct-Q4_K_S.gguf", - "qwen_image", - "cuda:1" - ] - }, - "90": { - "type": "UNETLoader", - "widgets_values": [ - "WanVideo/2_2/Wan2_2-I2V-A14B-HIGH_fp8_e4m3fn_scaled_KJ.safetensors", - "default" - ] - }, - "108": { - "type": "VAELoaderMultiGPU", - "widgets_values": [ - "ae.safetensors", - "cuda:1" - ] - }, - "120": { - "type": "CLIPLoaderMultiGPU", - "widgets_values": [ - "umt5_xxl_fp8_e4m3fn_scaled.safetensors", - "wan", - "cuda:1" - ] - }, - "126": { - "type": "UnetLoaderGGUFDisTorch2MultiGPU", - "widgets_values": [ - "Wan2.2-T2V-A14B-HighNoise-Q8_0.gguf", - "cuda:0", - 47.5, - "cpu", - "", - true - ] - }, - "130": { - "type": "VAELoaderMultiGPU", - "widgets_values": [ - "wan_2.1_vae.safetensors", - "cuda:1" - ] - }, - "135": { - "type": "UnetLoaderGGUFDisTorch2MultiGPU", - "widgets_values": [ - "Wan2.2-T2V-A14B-LowNoise-Q8_0.gguf", - "cuda:0", - 14, - "cpu", - "", - true - ] - } -} \ No newline at end of file diff --git a/memory-bank/systemPatterns.md b/memory-bank/systemPatterns.md deleted file mode 100644 index c3b1c58..0000000 --- a/memory-bank/systemPatterns.md +++ /dev/null @@ -1,466 +0,0 @@ -# System Architecture & Patterns (Updated 2025-09-29) - -## Core Architecture - -### Dynamic Class Override System -Foundation Pattern: City96's elegant inheritance-based approach (Dec 2024 revolution) - -```python -def override_class(original_class, device_param="device"): - class MultiGPUClass(original_class): - @classmethod - def INPUT_TYPES(cls): - inputs = original_class.INPUT_TYPES() - inputs["required"][device_param] = (get_device_list(),) - return inputs - - def override(self, *args, **kwargs): - device = kwargs.pop(device_param, None) - mm.text_encoder_device = device - return original_class.FUNCTION(self, *args, **kwargs) - - return MultiGPUClass -``` - -Key Benefits: -- 50 lines vs 400+: Eliminated manual class definitions -- Universal Support: Works with any ComfyUI loader node -- Maintenance: Auto-adapts to ComfyCore changes -- Consistency: Unified behavior across all MultiGPU nodes - -### Load-Patch-Distribute (LPD) Method -DisTorch2 Core Process: - -```python -# 1. LOAD - Always on compute device first -tensor = load_tensor_on_compute_device(tensor_name) - -# 2. PATCH - Apply all LoRAs at full precision -if lora_patches: - tensor = apply_lora_patches(tensor, lora_patches, precision=torch.float16) - -# 3. DISTRIBUTE - Move to target device after patching -final_tensor = tensor.to(target_device) -``` - -Design Principles: -- Quality First: No precision loss during LoRA application -- Deterministic: Same allocation every time -- ComfyUI Native: Works with existing ComfyCore patterns - -## Memory Management Architecture - -### Virtual VRAM System -Concept: Make CPU/secondary GPU memory appear as extended VRAM - -```python -class VirtualVRAM: - def __init__(self, compute_device, donor_device, virtual_gb): - self.compute_device = compute_device # e.g., "cuda:0" - self.donor_device = donor_device # e.g., "cpu" or "cuda:1" - self.virtual_gb = virtual_gb # Extended memory pool - - def allocate_layers(self, model_layers, allocation_string): - # "cuda:0,2.5gb;cpu,*" -> assign layers based on cumulative memory -``` - -### Expert Allocation Modes - -Bytes Mode (Recommended): -```python -# "cuda:0,2.5gb;cuda:1,3.0g;cpu,*" -def parse_bytes_allocation(allocation_string): - devices = [] - for device_spec in allocation_string.split(';'): - device_name, memory_spec = device_spec.split(',') - if memory_spec == '*': - memory_bytes = float('inf') # Overflow device - else: - memory_bytes = parse_memory_string(memory_spec) # 2.5gb -> bytes - devices.append((device_name, memory_bytes)) - return devices -``` - -Ratio Mode (llama.cpp style): -```python -# "cuda:0,25%;cpu,75%" -> 1:3 split -def parse_ratio_allocation(allocation_string): - total_ratio = sum(float(spec.split(',')[1].rstrip('%')) for spec in allocation_string.split(';')) - device_ratios = [] - for device_spec in allocation_string.split(';'): - device_name, ratio_spec = device_spec.split(',') - ratio = float(ratio_spec.rstrip('%')) / total_ratio - device_ratios.append((device_name, ratio)) - return device_ratios -``` - -### Selective Ejection Pipeline (v2.5.0 - VERIFIED WORKING) - -**Load-time Flagging** (per-model transient): -```python -# In DisTorch2 wrapper after real loader returns -if hasattr(out[0], 'model') and hasattr(out[0].model, '_mgpu_keep_loaded'): - keep_loaded = out[0].model._mgpu_keep_loaded - out[0].model._mgpu_unload_distorch_model = (not keep_loaded) -``` -Purpose: Mark specific DisTorch models for ejection when user unchecks "keep loaded" - -**Manager-Parity Cleanup Trigger**: -```python -def force_full_system_cleanup(reason="manual", force=True): - pq.set_flag("unload_models", True) # Exactly what Manager's - pq.set_flag("free_memory", True) # "Free model and node cache" does -``` - -**Selective Unloading** (patched `mm.unload_all_models`): -```python -def _mgpu_patched_unload_all_models(): - # Categorize models by flag - models_to_unload = [lm for lm in mm.current_loaded_models - if getattr(lm.model, '_mgpu_unload_distorch_model', False)] - kept_models = [lm for lm in mm.current_loaded_models - if not getattr(lm.model, '_mgpu_unload_distorch_model', False)] - - if kept_models: - # Selective unload: eject flagged, retain others - for lm in models_to_unload: - lm.model_unload(unpatch_weights=True) - - # Add GC anchors to prevent premature collection - for lm in kept_models: - add_retention_anchor(lm.model, "keep_loaded_protection") - - # Rebuild with kept models only - mm.current_loaded_models = kept_models - else: - # No models to keep - standard cleanup - _mgpu_original_unload_all_models() -``` - -**Multi-Device VRAM + CPU Management** (patched `mm.soft_empty_cache`): -```python -def soft_empty_cache_distorch2_patched(force=False): - # 1. Detect DisTorch2 activity - is_distorch_active = any(model_hash in safetensor_allocation_store - for model in mm.current_loaded_models) - - # 2. VRAM allocator management - if is_distorch_active: - soft_empty_cache_multigpu() # Clear all device caches - else: - original_soft_empty_cache(force) # Standard single-device - - # 3. Adaptive CPU memory management - check_cpu_memory_threshold() - - # 4. Forced executor reset (Manager parity) - if force: - trigger_executor_cache_reset(reason="forced_soft_empty", force=True) -``` - -**Verified Working** (Production Logs 2025-09-30): -``` -[CATEGORIZE_SUMMARY] kept_models: 2, models_to_unload: 1, total: 3 -[SELECTIVE_UNLOAD] Proceeding with selective unload: retaining 2, unloading 1 -[UNLOAD_EXECUTE] Unloading model: Flux -[REMAINING_MODEL] 0: AutoencodingEngine -[REMAINING_MODEL] 1: FluxClipModel_ -``` - -### Device Detection & Management - -Multi-Device Enumeration: -```python -def get_device_list(): - devices = ["cpu"] # Always available - if torch.cuda.is_available(): - devices.extend([f"cuda:{i}" for i in range(torch.cuda.device_count())]) - # XPU/NPU/MLU/MPS/DirectML/CoreX detection... - return devices -``` - -Device Bandwidth Intelligence (from benchmarking): -1. NVLINK (~50.8 GB/s) -2. PCIe 4.0 x16 (~27.2 GB/s) -3. PCIe 3.0 x8 (~6.8 GB/s) -4. PCIe 3.0 x4 (~2.1 GB/s) - -## Integration Patterns - -### ComfyCore Alignment -Philosophy: Work WITH ComfyUI, not against it - -```python -# GOOD: Use ComfyCore's device management -current_device = mm.get_torch_device() -mm.text_encoder_device = target_device - -# AVOID: Direct PyTorch device manipulation -torch.cuda.set_device(device_id) # Bypasses ComfyCore -``` - -### Node Registration System -```python -# Dynamic registration based on available dependencies -if "ComfyUI-GGUF" in installed_modules: - NODE_CLASS_MAPPINGS["UnetLoaderGGUFDisTorch2MultiGPU"] = create_gguf_distorch_node() -``` - -### Dependency Detection -```python -def check_module_availability(module_paths): - for path in module_paths: - if os.path.exists(os.path.join(custom_nodes_dir, path)): - return True - return False -``` - -## Performance Optimization Patterns - -### Layer Transfer Optimization -```python -def optimized_layer_transfer(layer, source_device, target_device): - if source_device == target_device: - return layer - non_blocking = "cuda" in source_device and "cuda" in target_device - if source_device == "cpu" and "cuda" in target_device: - layer = layer.pin_memory() - return layer.to(target_device, non_blocking=non_blocking) -``` - -### Memory Pressure Management -```python -def should_auto_offload(model_size_gb, vram_available_gb, threshold=0.9): - return model_size_gb > (vram_available_gb * threshold) - -def calculate_offload_amount(model_size_gb, target_vram_usage_gb): - return max(0, model_size_gb - target_vram_usage_gb) -``` - -## Error Handling Philosophy - -### Fail Loudly Pattern -```python -# GOOD: Let ComfyCore changes surface immediately -def load_model(self, model_name, device): - return original_loader.load_unet(model_name, device) - -# AVOID: Defensive coding that masks issues -try: - return original_loader.load_unet(model_name, device) -except AttributeError: - return fallback_method() -``` - -### Integration Validation -```python -def validate_comfycore_integration(): - required_attrs = ['FUNCTION', 'INPUT_TYPES', 'RETURN_TYPES'] - for attr in required_attrs: - if not hasattr(target_class, attr): - raise AttributeError(f"ComfyCore node missing {attr} - API changed") -``` - -## Code Style Patterns - -### Self-Documenting Code -```python -def override_class_with_device_selection(original_class, device_param_name="device"): - compute_device = kwargs.get(device_param_name, mm.get_torch_device()) -``` - -### Minimal Comments Philosophy -Prefer structure and naming to convey intent; use comments for non-obvious constraints/assumptions. - -## Architectural Decision Records - -### Why Dynamic Class Override vs Manual Definitions -Decision: Use inheritance-based class override (City96 approach) -Rationale: -- Reduces code from 400+ lines to ~50 lines -- Auto-adapts to ComfyCore changes -- Eliminates maintenance burden of manual node definitions -- Provides consistent behavior across all node types - -### Why Load-Patch-Distribute vs Direct Distribution -Decision: Always load on compute device first, then distribute -Rationale: -- Ensures LoRA patches applied at full precision -- Maintains quality parity with single-GPU workflows -- Predictable behavior regardless of target device -- Works with ComfyCore’s existing patching mechanisms - -### Why Expert Modes vs Automatic Only -Decision: Provide both automatic and expert allocation modes -Rationale: -- Automatic mode enables low-VRAM users immediately -- Expert modes allow optimization for specific hardware -- Performance depends on bandwidth topology; experts need control - -### Why Universal Device Support vs CUDA-Only -Decision: Support CPU, XPU, NPU, MLU, MPS, DirectML alongside CUDA -Rationale: -- ComfyUI’s user base spans diverse hardware -- Future-proof for emerging accelerators -- Hardware democracy principle - -### Why Per-Model Flag vs Global Sentinel (Updated) -Decision: Use per-model `_mgpu_unload_distorch_model` instead of a global “DISTORCH2_UNLOAD_MODEL” sentinel -Rationale: -- Surgical precision at model granularity -- No persistent or cross-workflow state -- Cleaner semantics under ComfyUI’s queue/flag model - -Hardened unloading rule (target to re-apply): -- If no models are flagged for ejection, `mm.unload_all_models` must be a strict no-op to preserve retained models across the full Manager-parity flow. - -## Testing & Validation Patterns - -### Hardware Configuration Testing -```python -HARDWARE_CONFIGS = [ - {"compute": "cuda:0", "donor": "cpu", "connection": "PCIe 4.0 x16"}, - {"compute": "cuda:0", "donor": "cuda:1", "connection": "NVLink"}, - {"compute": "cuda:0", "donor": "cuda:1", "connection": "PCIe 3.0 x8"}, - {"compute": "cuda:0", "donor": "cuda:1", "connection": "PCIe 3.0 x4"}, -] -``` - -### Model Compatibility Validation -```python -TEST_MODELS = [ - {"name": "FLUX.1-dev", "format": ".safetensors", "size_gb": 23.8}, - {"name": "WAN 2.2", "format": ".safetensors", "size_gb": 14.0}, - {"name": "FLUX-GGUF", "format": ".gguf", "size_gb": 11.8}, - {"name": "QWEN Image", "format": ".safetensors", "size_gb": 38.0}, -] -``` - -### Performance Regression Testing -```python -def benchmark_allocation_performance(model, hardware_config, allocation_configs): - baseline_time = benchmark_single_gpu(model) - for allocation in allocation_configs: - distributed_time = benchmark_distributed(model, hardware_config, allocation) - performance_ratio = distributed_time / baseline_time - assert performance_ratio < expected_slowdown_threshold(hardware_config) -``` - -## Recent Refactorings (v2.5.0) - -### DisTorch2 Allocation Consolidation (-179 lines) -**Problem**: 85% code duplication between `analyze_safetensor_loading()` and `analyze_safetensor_loading_clip()` - -**Solution**: Unified function with CLIP support flag -```python -def _extract_clip_head_blocks(raw_block_list, compute_device): - """Helper: Identify and pre-assign CLIP head blocks to compute device""" - head_keywords = ['embed', 'wte', 'wpe', 'token_embedding', 'position_embedding'] - head_blocks = [] - distributable_blocks = [] - block_assignments = {} - - for module_size, module_name, module_object, params in raw_block_list: - if any(kw in module_name.lower() for kw in head_keywords): - head_blocks.append((module_size, module_name, module_object, params)) - block_assignments[module_name] = compute_device - else: - distributable_blocks.append((module_size, module_name, module_object, params)) - - return head_blocks, distributable_blocks, block_assignments, head_memory - -def analyze_safetensor_loading(model_patcher, allocations_string, is_clip=False): - """Unified allocation function with CLIP head preservation support""" - # Common allocation logic... - - if is_clip: - head_blocks, distributable_raw, block_assignments, head_memory = \ - _extract_clip_head_blocks(raw_block_list, compute_device) - # Adjust compute_device quota for head blocks - donor_quotas[compute_device] -= head_memory - else: - distributable_raw = raw_block_list - block_assignments = {} - - # Continue with unified distribution logic... -``` - -**Benefits**: -- Single source of truth for allocation -- CLIP special case isolated in 20-line helper -- Easier to maintain and debug -- Same behavior, cleaner architecture - -### Production Cleanup (-40 lines) -**Removed**: Diagnostic instrumentation wrapper `_mgpu_instrumented_soft_empty_cache()` - -**Rationale**: Pure debug logging with no production function - removed to clean codebase - -**Result**: Clear separation between device_utils.py (functional) and model_management_mgpu.py (lifecycle) - -## Module Architecture (Post-Refactoring) - -### Core Module Separation -Problem Solved: Eliminated circular import `device_utils.py` ↔ `distorch_2.py` - -Solution: `model_management_mgpu.py` as central model lifecycle hub - -### Module Responsibilities - -device_utils.py (Base Layer): -- Device enumeration and detection -- VRAM cache management (`soft_empty_cache_multigpu`) -- Pure hardware abstraction – no model tracking - -model_management_mgpu.py (Core Layer): -- Model lifecycle tracking and memory logging -- Cleanup orchestration (`force_full_system_cleanup`, `trigger_executor_cache_reset`, `check_cpu_memory_threshold`) -- Patched unload path (selective ejection) - -distorch_2.py/distorch.py (Feature Layer): -- DisTorch distribution algorithms and allocation analysis -- Per-model flagging (`_mgpu_unload_distorch_model`) during DisTorch loads -- Imports FROM Core/Base only - -UI Layer: nodes.py, checkpoint_multigpu.py -- Device-aware user interfaces and node definitions - -Assembly: __init__.py -- Final integration/patch registration (`mm.soft_empty_cache` patch, node maps) - -### Import Flow Architecture -``` - ┌─────────────────┐ - │ __init__.py │ ← Assembly Layer - └─────────────────┘ - ↑ - ┌─────────────────┐ - │ UI Layer │ ← nodes.py, checkpoint_multigpu.py - └─────────────────┘ - ↑ - ┌─────────────────┐ - │ Feature Layer │ ← distorch_2.py, distorch.py - └─────────────────┘ - ↑ - ┌─────────────────┐ - │ Core Layer │ ← model_management_mgpu.py - └─────────────────┘ - ↑ - ┌─────────────────┐ - │ Base Layer │ ← device_utils.py - └─────────────────┘ -``` - -### Architectural Validation -Rule: Dependencies only flow UPWARD. Violations create circular imports. - -Prevention: Before any import, verify it respects the layer hierarchy. - -### Function Migration Record -Moved from device_utils.py to model_management_mgpu.py: -- `multigpu_memory_log` – memory state logging -- `trigger_executor_cache_reset` – CPU memory management -- `check_cpu_memory_threshold` – adaptive cleanup triggers -- `force_full_system_cleanup` – Manager-parity free flow - -Rationale: These belong to model lifecycle/cleanup, not hardware enumeration. diff --git a/memory-bank/techContext.md b/memory-bank/techContext.md deleted file mode 100644 index 349923b..0000000 --- a/memory-bank/techContext.md +++ /dev/null @@ -1,175 +0,0 @@ -# Technical Context & Dependencies (Updated 2025-09-29) - -## Core Technology Stack - -### Python Environment -Requirements: -- Python 3.10+ recommended -- PyTorch 2.x (CUDA/HIP/XPU backends as available) -- ComfyUI as host framework - -### Framework Dependencies - -Required (ComfyUI Core) -```python -import torch -import comfy.model_management as mm -import comfy.model_patcher -import comfy.utils -import folder_paths -``` - -Optional (External Custom Nodes) -```python -# ComfyUI-GGUF Integration -try: - from ComfyUI_GGUF import nodes as gguf_nodes - GGUF_AVAILABLE = True -except ImportError: - GGUF_AVAILABLE = False - -# WanVideoWrapper Integration -try: - import ComfyUI_WanVideoWrapper.nodes as wanvideo_nodes - WANVIDEO_AVAILABLE = True -except ImportError: - WANVIDEO_AVAILABLE = False -``` - -## Device Support Matrix - -Primary Support (tested) -- CUDA (NVIDIA) -- CPU -- MPS (Apple Metal) - -Extended/Community -- XPU (Intel) -- NPU (Ascend) -- MLU (Cambricon) -- DirectML (Windows) -- CoreX/IXUCA - -## Integration Architecture (Current Patch Points) - -This project extends ComfyUI through carefully scoped patches and runtime overrides. The current core integration points are: - -1) get_torch_device/text_encoder_device override (device selection) -- File: `__init__.py` -- Patch: - - `mm.get_torch_device = get_torch_device_patched` - - `mm.text_encoder_device = text_encoder_device_patched` -- Purpose: Respect user-selected devices handoff by MultiGPU wrappers and maintain ComfyUI alignment. - -2) soft_empty_cache (multi-device + CPU reset) -- File: `__init__.py` -- Patch: - - `mm.soft_empty_cache = soft_empty_cache_distorch2_patched` -- Behavior: - - Detects DisTorch2 activity, clears allocator caches across ALL devices via `soft_empty_cache_multigpu()` (from `device_utils.py`) - - Adaptive CPU memory reset (threshold-based), and optional forced `PromptExecutor.reset()` when `force=True` (Manager parity) - -3) unload_all_models (selective ejection) -- File: `model_management_mgpu.py` -- Patch: - - `mm.unload_all_models = _mgpu_patched_unload_all_models` -- Behavior: - - Splits `mm.current_loaded_models` into: - - `models_to_unload` where per-model `_mgpu_unload_distorch_model == True` - - `kept_models` for all others - - If flagged models exist: unload them only, then set `mm.current_loaded_models = kept_models` - - Current caveat: When none are flagged, the code delegates to the original unload (target is strict no-op; see System Patterns and Fix Plan) - -4) DisTorch2 load-time model flagging (per-model transient) -- File: `distorch_2.py` -- Where: - - In DisTorch2 wrappers (UNET/CLIP/VAE) within `override(...)` after original call: - - `out[0].model._mgpu_unload_distorch_model = (keep_loaded == False)` -- Rationale: - - Surgical per-model control enables selective ejection in patched unload without any global sentinel - -5) Manager parity helper -- File: `model_management_mgpu.py` -- Function: - - `force_full_system_cleanup(reason="manual", force=True)` -- Behavior: - - Sets both `unload_models=True` and `free_memory=True` on PromptQueue, matching Manager’s “Free model and node cache” button behavior - -## Selective Ejection Flow (Technical Overview) - -- Load time (DisTorch2 wrappers): - - Mark models for ejection if keep_loaded=False -- Free flow (Manager or programmatic parity): - - /free → prompt_worker picks flags → calls `mm.unload_all_models()` (selective) → `PromptExecutor.reset()` → GC → `mm.soft_empty_cache()` (multi-device) -- Intended properties: - - Models flagged for ejection are destroyed - - Retained models remain live after full flow (including reset/GC/soft_empty) - -Current caveat (to fix next): -- When no models are flagged, the patched unload delegates to the original unload, which unloads everything. The target is strict no-op in this branch. - -## Development Environment - -Supported OS -- Linux (primary) -- Windows 10/11 -- macOS (Apple Silicon via MPS) - -Tools -- IDE: VSCode -- VCS: Git (conventional commits encouraged) -- Testing: Manual validation across available hardware + community testing - -## Performance Characteristics - -Bandwidth hierarchy -1. NVLink (~50.8 GB/s) – near-native performance -2. PCIe 4.0 x16 (~27.2 GB/s) – excellent offloading -3. PCIe 3.0 x8 (~6.8 GB/s) -4. PCIe 3.0 x4 (~2.1 GB/s) - -Load-Patch-Distribute (LPD) -- Always load on compute device first -- Apply LoRAs at full precision -- Distribute blocks to assigned devices for final placement -- Ensures quality preservation and deterministic behavior - -## Configuration Management - -Expert allocation strings -- Bytes mode (recommended): - - `"cuda:0,2.5gb;cuda:1,3.0g;cpu,*"` -- Ratio mode: - - `"cuda:0,25%;cpu,75%"` -- Fraction mode (legacy): - - `0.8`, `0.5`, `0.95` - -## Debugging & Monitoring - -Logging -- `logger.mgpu_mm_log(...)` for structured memory/system logs -- `multigpu_memory_log(identifier, tag)` for timestamped CPU/VRAM snapshots - -Inspection -- `device_utils.comfyui_memory_load(tag)` for one-line current memory snapshot -- VRAM cache clearing logs around `soft_empty_cache_multigpu()` - -## Architectural Rationale (Updated) - -Per-model flag over global sentinel -- Granular control, no persistent global state -- Isolated to each loaded model, matches ComfyUI lifecycle - -Patched unload behavior (selective) -- Maintain `kept_models` across the full free path -- Only eject DisTorch2 models when explicitly requested via keep_loaded=False - -Patched soft empty (multi-device) -- Ensure cache clearing is not limited to the single `mm.get_torch_device()` device -- CPU memory behavior integrated with PromptExecutor.reset() semantics - -## Known Technical Work (Next) - -- Reinstate strict no-op in `_mgpu_patched_unload_all_models` when `models_to_unload` is empty (no delegation to original unload) -- Add instrumentation and assertions to guarantee no unintended ejection of retained models after `/free` flow -- Re-run verification matrix and capture logs in Memory Bank