FLUX.1 Support and Intelligent Memory Management
Full FLUX Architecture Compatibility:
Added complete support for FLUX.1 models including FluxParams dataclass handling and custom MixedPrecisionOps layers
Smart Model Cloning:
Implements 3-tier fallback cloning system:
Fast deepcopy for standard models (SD/SDXL)
Config reconstruction for FLUX architectures
Manual recursive cloning for edge cases
Batch-Aware Parallelism:
Properly splits FLUX-specific kwargs tensors (guidance, y, etc.) across devices to prevent dimension mismatches
Intelligent Memory Management
Automatic Cleanup:
Added cleanup_parallel_model() function with weakref.finalize to automatically release GPU memory when models are deleted
OOM Prevention:
Aggressive memory clearing between clone operations using torch.cuda.empty_cache() and comfy.model_management.soft_empty_cache()
CPU Staging:
Temporarily moves base model to CPU during cloning to prevent VRAM spikes
Safe Model Switching:
Clears previous parallel replicas before creating new ones, preventing memory accumulation when switching between models
🔧 Technical Improvements
Robustness & Error Handling
Safe Attribute Access: Added safe_getattr() helper to handle custom layers (FLUX's Linear ops) that may lack standard attributes like bias
Cache Management: Automatically clears FLUX-specific cached tensors (img_ids, txt_ids, freqs_cis, etc.) on clone and between runs to support dynamic batch sizes
Storage Error Recovery:
Fixed "invalid python storage" and "can't pickle local object" errors through serialization-free cloning methods
Hardware Optimization
SM_80 Auto-Detection: Automatically disables Flash Attention and xFormers on GPUs with compute capability < 8.0 (pre-Ampere)
Multi-Device Support: Enhanced support for CUDA, XPU (Intel Arc), MPS (Apple Silicon), and DirectML with proper device context management
Data Handling
BFloat16 Compatibility: Proper nn.Parameter wrapping preserves FLUX's native BFloat16 precision during cloning
Dataclass Cloning: Added clone_dataclass_or_object() to properly handle nested configuration objects
Tensor Splitting: split_kwargs() function ensures all batch-dependent tensors (including guidance embeddings) are correctly partitioned across workers
🐛 Bug Fixes
Fixed: "The size of tensor a (X) must match the size of tensor b (Y)" errors when changing batch sizes between runs
Fixed: "'Linear' object has no attribute 'bias'" AttributeError in FLUX models
Fixed: "Attempted to call copy_() on an invalid python storage" errors during model cloning
Fixed: Memory not releasing when loading different models sequentially
Fixed: "FluxParams.init() missing required positional arguments" configuration errors