110 lines
2.9 KiB
Python
110 lines
2.9 KiB
Python
"""
|
|
Sequential Loaders Base Class
|
|
Complete independent implementation - No dependency on FreeVRAM
|
|
"""
|
|
|
|
import time
|
|
from typing import Any, Dict, List, Tuple, Optional
|
|
|
|
# Type compatibility with FreeVRAM node
|
|
ANY_SOCKET_TYPES = "IMAGE,LATENT,MODEL,CONDITIONING,CLIP,VAE,MASK,INT,FLOAT,STRING,BOOLEAN"
|
|
|
|
|
|
class SequentialLoaderBase:
|
|
"""
|
|
Base class for all Sequential loaders
|
|
|
|
Features:
|
|
- trigger input/output management
|
|
- Memory tracking (independent implementation)
|
|
- Unified log output
|
|
- Error handling
|
|
"""
|
|
|
|
CATEGORY = "u5"
|
|
|
|
# Class variable: Memory tracker
|
|
_memory_tracker: List[Dict[str, Any]] = []
|
|
|
|
@classmethod
|
|
def add_trigger_input(cls, input_types: Dict) -> Dict:
|
|
"""
|
|
Add trigger input to INPUT_TYPES
|
|
|
|
Args:
|
|
input_types: Existing input definition
|
|
|
|
Returns:
|
|
Input definition with trigger input added
|
|
"""
|
|
if "optional" not in input_types:
|
|
input_types["optional"] = {}
|
|
|
|
input_types["optional"]["trigger"] = (ANY_SOCKET_TYPES, {
|
|
"tooltip": "Trigger(Any): Connect from previous node to control execution order"
|
|
})
|
|
return input_types
|
|
|
|
def with_trigger_passthrough(
|
|
self,
|
|
results: Any,
|
|
trigger: Optional[Any] = None
|
|
) -> Tuple:
|
|
"""
|
|
Add trigger to results for passthrough
|
|
|
|
Args:
|
|
results: Loader execution results
|
|
trigger: Input trigger
|
|
|
|
Returns:
|
|
Tuple of (results..., trigger)
|
|
"""
|
|
if isinstance(results, tuple):
|
|
return (*results, trigger)
|
|
return (results, trigger)
|
|
|
|
def log_execution(self, loader_type: str, params: Dict[str, Any]):
|
|
"""
|
|
Unified log output
|
|
|
|
Args:
|
|
loader_type: Loader type (e.g., "Checkpoint", "VAE")
|
|
params: Parameters to log
|
|
"""
|
|
print(f"[Sequential {loader_type}] Executing: {params}")
|
|
|
|
@classmethod
|
|
def track_memory(cls, node_name: str, memory_gb: float):
|
|
"""
|
|
Record memory usage
|
|
|
|
Args:
|
|
node_name: Node identifier
|
|
memory_gb: Memory usage in GB
|
|
"""
|
|
cls._memory_tracker.append({
|
|
"node": node_name,
|
|
"memory_gb": memory_gb,
|
|
"timestamp": time.time()
|
|
})
|
|
|
|
@classmethod
|
|
def get_memory_report(cls) -> Dict[str, Any]:
|
|
"""
|
|
Generate memory usage report
|
|
|
|
Returns:
|
|
Memory report dictionary
|
|
"""
|
|
total = sum(item["memory_gb"] for item in cls._memory_tracker)
|
|
return {
|
|
"total_gb": total,
|
|
"count": len(cls._memory_tracker),
|
|
"entries": cls._memory_tracker
|
|
}
|
|
|
|
@classmethod
|
|
def reset_memory_tracker(cls):
|
|
"""Reset memory tracker"""
|
|
cls._memory_tracker.clear() |