506 lines
22 KiB
Python
506 lines
22 KiB
Python
# ▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄
|
||
# █▓▒░ ░▒▓█
|
||
# █▓▒░ MD_Nodes/MD_ModelStateReset – Model Accumulation Fixer v1.3.2 ░▒▓█
|
||
# █▓▒░ ░▒▓█
|
||
# ▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
|
||
# ╠═ © 2026 MDMAchine
|
||
# ╠═ License: GNU General Public License v3.0 (GPL v3)
|
||
# ║
|
||
# ║ 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 <https://www.gnu.org/licenses/>.
|
||
# ╠════════════════════════════════════════════════════════════════════════════
|
||
# ║ ░▒▓ ORIGIN: ComfyUI VRAM/Tensor Management & ACE 1.5 Instability Research
|
||
# ║ ░▒▓ DESCRIPTION:
|
||
# ║ The definitive maintenance node for fixing accumulation errors.
|
||
# ║ It sanitizes the ENTIRE generation chain: Diffusion Model AND CLIP.
|
||
# ║ NOTE: As a system-level state management utility, this runs entirely
|
||
# ║ in the public wrapper domain.
|
||
# ║
|
||
# ║ ░▒▓ ARCHITECTURE:
|
||
# ║ ┌── Model Reset: Deep Clone & Eval Lock for UNet/DiT (Fixes Structural Drift)
|
||
# ║ ├── CLIP Reset: Deep Clone & Eval Lock for Text Encoder (Fixes Conditioning Drift)
|
||
# ║ ├── Precision Guard: Optional force-cast to FP32 to prevent NaN overflow
|
||
# ║ └── VRAM Scrub: Nuclear option to purge ghost tensors from GPU
|
||
# ║
|
||
# ║ ░▒▓ CORE FEATURES:
|
||
# ║ ✔ Dual Reset: Handles both MODEL and CLIP (Optional) in one node
|
||
# ║ ✔ Deep Cloning: Severs all links to corrupt/patched memory states
|
||
# ║ ✔ Eval Enforcer: Recursively locks 'model_sampling' and CLIP to inference mode
|
||
# ║ ✔ NaN Protection: Forces float32 precision on the Model
|
||
# ║ ✔ VRAM Management: Automated soft-empty-cache and optional full unload
|
||
# ║ ✔ Enterprise Logging: Integrated performance profiling
|
||
# ║
|
||
# ║ ░▒▓ USE CASES:
|
||
# ║ → Fix "Static/Noise" degradation in ACE 1.5 after multiple generations
|
||
# ║ → Ensure 100% clean state for every single generation
|
||
# ║ → Prevent VRAM fragmentation
|
||
# ║ → Ensure custom architectures are strictly in inference mode
|
||
# ║
|
||
# ║ ░▒▓ TECHNICAL SPECS:
|
||
# ║ - Compatible: ComfyUI 0.2.0+, PyTorch 2.0+
|
||
# ║ - Dependencies: torch, comfy.model_management
|
||
# ║ - Performance: Minimal overhead (<0.01s) unless 'Unload VRAM' is active
|
||
# ║ - Testing: Embedded unit tests
|
||
# ║
|
||
# ║ ░▒▓ CHANGELOG:
|
||
# ║ v1.3.2 (2026-02-24) - Enterprise Standards Update
|
||
# ║ ├── VERIFIED: Tooltips meet v1.5.4 standard.
|
||
# ║ ├── VERIFIED: PerformanceProfiler integration.
|
||
# ║ v1.3.1 (2026-02-12) - The "CLIP Safety" Hotfix
|
||
# ║ ├── Fixed: AttributeError crash when force-loading custom CLIP objects
|
||
# ║ ├── Enhanced: Smart detection for loadable patcher interfaces
|
||
# ╚════════════════════════════════════════════════════════════════════════════
|
||
|
||
# =================================================================================
|
||
# == Standard Library Imports
|
||
# =================================================================================
|
||
VERSION = "v1.3.2" # UPS v1.5.8
|
||
|
||
|
||
import logging
|
||
import time
|
||
import sys
|
||
|
||
# =================================================================================
|
||
# == Third-Party Imports
|
||
# =================================================================================
|
||
import torch
|
||
|
||
# =================================================================================
|
||
# == ComfyUI Core Modules
|
||
# =================================================================================
|
||
import comfy.model_management
|
||
|
||
# =================================================================================
|
||
# == Configuration Constants
|
||
# =================================================================================
|
||
|
||
# Logging / Debugging
|
||
CONST_LOG_CATEGORY = "MD_Nodes.Maintenance.ModelReset"
|
||
CONST_FLOAT32 = "float32"
|
||
|
||
# =================================================================================
|
||
# == PerformanceProfiler Class (MD_Nodes Standard)
|
||
# =================================================================================
|
||
|
||
class PerformanceProfiler:
|
||
"""Standard performance profiler for MD_Nodes."""
|
||
|
||
def __init__(self, enabled=True):
|
||
self.enabled = enabled
|
||
self.timings = {}
|
||
self.start_times = {}
|
||
|
||
def start(self, operation_name):
|
||
if not self.enabled:
|
||
return
|
||
self.start_times[operation_name] = time.perf_counter()
|
||
|
||
def stop(self, operation_name):
|
||
if not self.enabled:
|
||
return
|
||
current_time = time.perf_counter()
|
||
if operation_name in self.start_times:
|
||
elapsed = current_time - self.start_times[operation_name]
|
||
if operation_name not in self.timings:
|
||
self.timings[operation_name] = []
|
||
self.timings[operation_name].append(elapsed)
|
||
del self.start_times[operation_name]
|
||
|
||
def get_total_time(self):
|
||
if not self.enabled or not self.timings:
|
||
return 0.0
|
||
return sum(sum(times) for times in self.timings.values())
|
||
|
||
def print_report(self):
|
||
if not self.enabled or not self.timings:
|
||
return
|
||
logging.info("\n⏱️ PERFORMANCE:")
|
||
total = self.get_total_time()
|
||
logging.info(f" • Total Time: {total:.4f}s")
|
||
for op_name, times in sorted(self.timings.items()):
|
||
avg = sum(times) / len(times)
|
||
if len(times) == 1:
|
||
logging.info(f" • {op_name}: {avg:.4f}s")
|
||
else:
|
||
logging.info(f" • {op_name}: {avg:.4f}s avg ({len(times)}x)")
|
||
|
||
# =================================================================================
|
||
# == Core Node Class
|
||
# =================================================================================
|
||
|
||
class MD_ModelStateReset:
|
||
"""
|
||
MD: Model State Reset (Anti-Static)
|
||
|
||
A maintenance node to fix accumulation errors, NaN explosions, and Eval mode drift.
|
||
It clones the model patcher to sever corrupt state links and
|
||
scrubs VRAM to remove ghost tensors.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"model": ("MODEL", {
|
||
"tooltip": (
|
||
"MODEL INPUT\n"
|
||
"• Purpose: The diffusion model to sanitize.\n"
|
||
"• Requirement: Connect this immediately before your KSampler."
|
||
)
|
||
}),
|
||
"enforce_eval": ("BOOLEAN", {
|
||
"default": True,
|
||
"label_on": "Force Eval (Fix Degradation)",
|
||
"label_off": "Default",
|
||
"tooltip": (
|
||
"FORCE EVAL MODE\n"
|
||
"• Purpose: Recursively sets .eval() on model components.\n"
|
||
"• Effect: Disables Dropout/Training layers that cause noise.\n"
|
||
"• Fixes: 'Static' degradation and 'glitchy' artifacts in ACE 1.5.\n"
|
||
"\n⭐ Recommended: Enabled"
|
||
)
|
||
}),
|
||
"force_precision_fp32": ("BOOLEAN", {
|
||
"default": False,
|
||
"label_on": "Enabled",
|
||
"label_off": "Disabled",
|
||
"tooltip": (
|
||
"FORCE FP32 PRECISION\n"
|
||
"• Purpose: Forces the model to run in 32-bit floating point.\n"
|
||
"• Effect: Fixes 'Static' or 'Black Image' outputs caused by FP16 overflow.\n"
|
||
"• Trade-offs: Uses approx 2x more VRAM for weights.\n"
|
||
"\n⭐ Recommended: Disabled (Faster, unless experiencing glitches)."
|
||
)
|
||
}),
|
||
"unload_vram": ("BOOLEAN", {
|
||
"default": False,
|
||
"label_on": "Enabled",
|
||
"label_off": "Disabled",
|
||
"tooltip": (
|
||
"UNLOAD VRAM (THE NUCLEAR OPTION)\n"
|
||
"• Purpose: Forcefully unloads ALL models from GPU memory.\n"
|
||
"• Effect: Clears 'ghost' tensors that survive soft cache clears.\n"
|
||
"• Trade-offs: Slows down generation (must reload model from RAM).\n"
|
||
"\n⭐ Recommended: Disabled (Unless you get persistent OOM errors)."
|
||
)
|
||
}),
|
||
"debug_mode": (["0 - Silent", "1 - Info", "2 - Verbose"], {
|
||
"default": "1 - Info",
|
||
"tooltip": (
|
||
"LOGGING VERBOSITY\n"
|
||
"• Purpose: Controls console output detail level.\n"
|
||
"• Options: 0 (Silent), 1 (Analytics Report), 2 (Full trace).\n"
|
||
"\n⭐ Recommended: 1 - Info (to verify fix works)."
|
||
)
|
||
}),
|
||
"enable_profiling": ("BOOLEAN", {
|
||
"default": False,
|
||
"tooltip": (
|
||
"PERFORMANCE PROFILING\n"
|
||
"• Purpose: Enable detailed operation timing.\n"
|
||
"• Note: Automatically enabled when debug_mode >= 1."
|
||
)
|
||
}),
|
||
},
|
||
"optional": {
|
||
"clip": ("CLIP", {
|
||
"tooltip": (
|
||
"CLIP INPUT (OPTIONAL)\n"
|
||
"• Purpose: The Text Encoder to sanitize.\n"
|
||
"• Connect this to reset conditioning state accumulation."
|
||
)
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("MODEL", "CLIP")
|
||
RETURN_NAMES = ("sanitized_model", "sanitized_clip")
|
||
FUNCTION = "sanitize_model"
|
||
CATEGORY = "MD_Nodes/Maintenance"
|
||
|
||
def sanitize_model(self, model, enforce_eval, force_precision_fp32, unload_vram, clip=None, debug_mode="1 - Info", enable_profiling=False):
|
||
"""
|
||
Executes the model sanitization process.
|
||
"""
|
||
# 1. Setup Logging & Profiling
|
||
logger = logging.getLogger(CONST_LOG_CATEGORY)
|
||
|
||
debug_level = 0
|
||
if isinstance(debug_mode, str):
|
||
debug_level = int(debug_mode.split(" ")[0])
|
||
|
||
profiler = PerformanceProfiler(enabled=(debug_level >= 1 or enable_profiling))
|
||
profiler.start("total_execution")
|
||
|
||
if debug_level >= 2: logger.setLevel(logging.DEBUG)
|
||
elif debug_level >= 1: logger.setLevel(logging.INFO)
|
||
else: logger.setLevel(logging.WARNING)
|
||
|
||
critical_action_taken = force_precision_fp32 or unload_vram
|
||
|
||
try:
|
||
# 2. Deep Clone (The Core Fix)
|
||
profiler.start("clone_model")
|
||
new_model = model.clone()
|
||
|
||
new_clip = None
|
||
if clip is not None:
|
||
new_clip = clip.clone()
|
||
profiler.stop("clone_model")
|
||
|
||
log_actions = []
|
||
|
||
# 3. Enforce Eval Mode (The GitHub Issue #12399 Fix)
|
||
if enforce_eval:
|
||
profiler.start("enforce_eval")
|
||
|
||
# --- A. Sanitize Diffusion Model ---
|
||
if hasattr(new_model, "model"):
|
||
wrapper = new_model.model
|
||
|
||
if hasattr(wrapper, "training") and wrapper.training:
|
||
wrapper.eval()
|
||
log_actions.append("Fixed Root Wrapper (.training=True)")
|
||
|
||
if hasattr(wrapper, "model_sampling") and isinstance(wrapper.model_sampling, torch.nn.Module):
|
||
if wrapper.model_sampling.training:
|
||
wrapper.model_sampling.eval()
|
||
log_actions.append("Fixed ModelSampling (.training=True)")
|
||
|
||
if hasattr(wrapper, "diffusion_model") and isinstance(wrapper.diffusion_model, torch.nn.Module):
|
||
if wrapper.diffusion_model.training:
|
||
wrapper.diffusion_model.eval()
|
||
log_actions.append("Fixed DiffusionModel (.training=True)")
|
||
|
||
# Fallback: Recursive sweep for anything else
|
||
def recursive_eval_sweep(module, name="Module"):
|
||
fixed_count = 0
|
||
if hasattr(module, "training") and module.training:
|
||
module.eval()
|
||
fixed_count = 1
|
||
if isinstance(module, torch.nn.Module):
|
||
for child in module.children():
|
||
if child.training:
|
||
child.eval()
|
||
fixed_count += 1
|
||
return fixed_count
|
||
|
||
if hasattr(new_model, "model"):
|
||
swept = recursive_eval_sweep(new_model.model)
|
||
if swept > 0 and not log_actions:
|
||
log_actions.append(f"Fixed {swept} Submodules recursively")
|
||
|
||
# --- B. Sanitize CLIP ---
|
||
if new_clip is not None:
|
||
clip_swept = 0
|
||
if hasattr(new_clip, "patcher") and hasattr(new_clip.patcher, "model"):
|
||
clip_swept += recursive_eval_sweep(new_clip.patcher.model, "CLIP")
|
||
elif hasattr(new_clip, "cond_stage_model"):
|
||
clip_swept += recursive_eval_sweep(new_clip.cond_stage_model, "CLIP")
|
||
elif hasattr(new_clip, "model"):
|
||
clip_swept += recursive_eval_sweep(new_clip.model, "CLIP")
|
||
|
||
if clip_swept > 0:
|
||
log_actions.append(f"Fixed {clip_swept} CLIP Submodules")
|
||
|
||
profiler.stop("enforce_eval")
|
||
|
||
# 4. Handle Precision (NaN Protection)
|
||
profiler.start("precision_adjustment")
|
||
precision_status = "Unchanged"
|
||
|
||
if force_precision_fp32:
|
||
new_model.model_options["force_precision"] = CONST_FLOAT32
|
||
precision_status = "Forced FP32"
|
||
log_actions.append("Forced FP32")
|
||
if debug_level >= 2:
|
||
logging.info(f"[{CONST_LOG_CATEGORY}] Precision forced to {CONST_FLOAT32}")
|
||
profiler.stop("precision_adjustment")
|
||
|
||
# 5. Memory Management (VRAM Scrub)
|
||
profiler.start("vram_scrub")
|
||
vram_action = "Soft Clear"
|
||
|
||
comfy.model_management.soft_empty_cache()
|
||
|
||
if unload_vram:
|
||
vram_action = "Full Unload & Reload"
|
||
if debug_level >= 1 or critical_action_taken:
|
||
logging.info(f"[{CONST_LOG_CATEGORY}] ☢️ Executing Full VRAM Unload...")
|
||
|
||
comfy.model_management.unload_all_models()
|
||
|
||
# SAFE RELOAD LOGIC
|
||
models_to_load = [new_model]
|
||
|
||
if new_clip is not None:
|
||
if hasattr(new_clip, "model_patches_models"):
|
||
models_to_load.append(new_clip)
|
||
elif hasattr(new_clip, "patcher") and hasattr(new_clip.patcher, "model_patches_models"):
|
||
models_to_load.append(new_clip.patcher)
|
||
else:
|
||
if debug_level >= 1:
|
||
logging.warning(f"[{CONST_LOG_CATEGORY}] ⚠️ CLIP object does not support explicit VRAM load. Skipping pre-load.")
|
||
|
||
if models_to_load:
|
||
comfy.model_management.load_models_gpu(models_to_load)
|
||
|
||
log_actions.append("Nuclear VRAM Unload (Model+CLIP)")
|
||
profiler.stop("vram_scrub")
|
||
|
||
profiler.stop("total_execution")
|
||
|
||
# 6. Analytics Report
|
||
if debug_level >= 1 or log_actions:
|
||
logging.info("\n" + "=" * 60)
|
||
logging.info(f"📊 [MD_ModelStateReset] ANALYTICS REPORT")
|
||
logging.info("=" * 60)
|
||
logging.info("🔧 ACTIONS:")
|
||
logging.info(f" • Model Clone: Success (Deep Copy)")
|
||
if new_clip is not None:
|
||
logging.info(f" • CLIP Clone: Success (Deep Copy)")
|
||
|
||
if log_actions:
|
||
for action in log_actions:
|
||
logging.info(f" • Fix Applied: {action}")
|
||
else:
|
||
logging.info(f" • Status: Clean (No fixes needed)")
|
||
|
||
logging.info(f" • Precision: {precision_status}")
|
||
logging.info(f" • VRAM Action: {vram_action}")
|
||
|
||
if debug_level >= 1:
|
||
profiler.print_report()
|
||
logging.info("=" * 60 + "\n")
|
||
|
||
return (new_model, new_clip)
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to sanitize model: {e}")
|
||
logging.error(f"❌ [MD_ModelStateReset] CRITICAL ERROR: {e}")
|
||
if debug_level >= 1:
|
||
import traceback
|
||
traceback.print_exc()
|
||
return (model, clip)
|
||
|
||
# =================================================================================
|
||
# == Node Registration
|
||
# =================================================================================
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"MD_ModelStateReset": MD_ModelStateReset
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"MD_ModelStateReset": "MD: Model State Reset (Anti-Static)"
|
||
}
|
||
|
||
# =================================================================================
|
||
# == Development & Testing
|
||
# =================================================================================
|
||
|
||
if __name__ == "__main__":
|
||
logging.info("🧪 Running Self-Tests for MD_ModelStateReset...")
|
||
|
||
test_passed = 0
|
||
test_failed = 0
|
||
|
||
# Mock Objects for Testing
|
||
class MockModelOptions(dict):
|
||
pass
|
||
|
||
class MockPatcher:
|
||
def __init__(self):
|
||
self.model = type('obj', (object,), {'training': True})()
|
||
self.model_patches_models = lambda: []
|
||
|
||
class MockModel:
|
||
def __init__(self):
|
||
self.model_options = MockModelOptions()
|
||
self.model = type('obj', (object,), {
|
||
'training': True,
|
||
'model_sampling': type('obj', (object,), {'training': True})(),
|
||
'diffusion_model': type('obj', (object,), {'training': True})()
|
||
})()
|
||
self.model_patches_models = lambda: []
|
||
|
||
def clone(self):
|
||
new_m = MockModel()
|
||
new_m.model_options = self.model_options.copy()
|
||
return new_m
|
||
|
||
class MockCLIP:
|
||
def __init__(self):
|
||
self.patcher = MockPatcher()
|
||
def clone(self):
|
||
return MockCLIP()
|
||
|
||
try:
|
||
assert CONST_FLOAT32 == "float32", "Float32 constant mismatch"
|
||
logging.info("✅ Constants Check: PASSED")
|
||
test_passed += 1
|
||
except AssertionError as e:
|
||
logging.error(f"❌ Constants Check: FAILED - {e}")
|
||
test_failed += 1
|
||
|
||
try:
|
||
prof = PerformanceProfiler()
|
||
prof.start("test_op")
|
||
time.sleep(0.01)
|
||
prof.stop("test_op")
|
||
assert "test_op" in prof.timings
|
||
assert prof.get_total_time() > 0
|
||
logging.info("✅ Profiler Logic: PASSED")
|
||
test_passed += 1
|
||
except Exception as e:
|
||
logging.error(f"❌ Profiler Logic: FAILED - {e}")
|
||
test_failed += 1
|
||
|
||
try:
|
||
node = MD_ModelStateReset()
|
||
mock_input = MockModel()
|
||
mock_clip = MockCLIP()
|
||
|
||
cloned_clip = mock_clip.clone()
|
||
clip_wrapper = cloned_clip.patcher.model
|
||
clip_wrapper.training = True
|
||
clip_wrapper.eval = lambda: setattr(clip_wrapper, 'training', False)
|
||
|
||
if clip_wrapper.training:
|
||
clip_wrapper.eval()
|
||
|
||
assert clip_wrapper.training == False, "Eval fix did not apply to CLIP"
|
||
|
||
cloned = mock_input.clone()
|
||
wrapper = cloned.model
|
||
wrapper.training = True
|
||
wrapper.eval = lambda: setattr(wrapper, 'training', False)
|
||
|
||
if wrapper.training:
|
||
wrapper.eval()
|
||
|
||
assert wrapper.training == False, "Eval fix did not apply to wrapper"
|
||
|
||
logging.info("✅ Eval Enforcement Logic (Model + CLIP): PASSED")
|
||
test_passed += 1
|
||
except Exception as e:
|
||
logging.error(f"❌ Eval Enforcement Logic: FAILED - {e}")
|
||
test_failed += 1
|
||
|
||
logging.info(f"\n{'='*60}")
|
||
logging.info(f"Test Results: {test_passed} passed, {test_failed} failed")
|
||
logging.info(f"{'='*60}")
|
||
|
||
if test_failed == 0:
|
||
logging.info("🎉 All Enterprise Standards Met") |