63 lines
2.1 KiB
Python
63 lines
2.1 KiB
Python
"""Script to create properly structured node files from extracted classes"""
|
|
import re
|
|
import os
|
|
|
|
# Mapping of old function names to new import paths
|
|
IMPORT_REPLACEMENTS = {
|
|
r'\bgentle_memory_cleanup\(\)': 'from ..utils.memory import gentle_memory_cleanup',
|
|
r'\bcheck_memory_safety\(': 'from ..utils.memory import check_memory_safety',
|
|
r'\bemergency_memory_cleanup\(\)': 'from ..utils.memory import emergency_memory_cleanup',
|
|
r'\bget_gpu_memory_info\(\)': 'from ..utils.memory import get_gpu_memory_info',
|
|
r'\bDEBUG_MODE\b': 'from ..utils.debug import DEBUG_MODE',
|
|
r'\blog_debug\(': 'from ..utils.debug import log_debug',
|
|
r'\bsave_debug_data\(': 'from ..utils.debug import save_debug_data',
|
|
r'\bcapture_timing\(': 'from ..utils.debug import capture_timing',
|
|
r'\bcapture_memory_usage\(': 'from ..utils.debug import capture_memory_usage',
|
|
}
|
|
|
|
# Base imports needed for all node files
|
|
BASE_IMPORTS = """import time
|
|
import logging
|
|
from typing import Dict, Any, Tuple, Optional
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import comfy.model_management as model_management
|
|
import comfy.utils
|
|
import comfy.sample
|
|
import comfy.samplers
|
|
import latent_preview
|
|
import folder_paths
|
|
|
|
# Import utilities from refactored modules
|
|
from ..utils.memory import (
|
|
gentle_memory_cleanup,
|
|
check_memory_safety,
|
|
emergency_memory_cleanup,
|
|
get_gpu_memory_info,
|
|
)
|
|
from ..utils.debug import (
|
|
DEBUG_MODE,
|
|
log_debug,
|
|
save_debug_data,
|
|
capture_timing,
|
|
capture_memory_usage,
|
|
)
|
|
from ..utils.quantization import detect_model_quantization
|
|
from ..constants import DEFAULT_TILE_SIZE, DEFAULT_TILE_OVERLAP
|
|
|
|
"""
|
|
|
|
# Read original imports for reference
|
|
with open('rocm_nodes.py', 'r', encoding='utf-8') as f:
|
|
original_imports = []
|
|
for line in f:
|
|
if line.strip().startswith('import ') or line.strip().startswith('from '):
|
|
original_imports.append(line)
|
|
elif line.strip() and not line.strip().startswith('#'):
|
|
break
|
|
|
|
print("Creating node files...")
|
|
print("(This script demonstrates the structure - actual files will be created manually)")
|
|
|