439 lines
17 KiB
Python
439 lines
17 KiB
Python
import os
|
|
from pathlib import Path
|
|
import folder_paths
|
|
import torch
|
|
from typing import Dict, List, Optional, Tuple
|
|
import yaml
|
|
import logging
|
|
import json
|
|
|
|
# Setup logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger('HunyuanLoraLoader')
|
|
|
|
class BaseHunyuanLoraLoader:
|
|
"""Base class for Hunyuan LoRA loading functionality"""
|
|
|
|
def __init__(self):
|
|
self.loaded_lora: Optional[Tuple[str, Dict[str, torch.Tensor]]] = None
|
|
|
|
@classmethod
|
|
def get_cache_dir(cls) -> str:
|
|
"""Get or create the cache directory for block settings"""
|
|
try:
|
|
from folder_paths import base_path, folder_names_and_paths
|
|
cache_dir = Path(folder_names_and_paths["custom_nodes"][0][0]) / "ComfyUI_PixelFun" / "cache"
|
|
if not os.path.exists(cache_dir):
|
|
os.makedirs(cache_dir)
|
|
return cache_dir
|
|
except Exception as e:
|
|
logger.error(f"Failed to create or access cache directory: {str(e)}")
|
|
raise
|
|
|
|
def get_settings_filename(self, lora_name: str) -> str:
|
|
"""Generate the settings filename for a given LoRA"""
|
|
base_name = os.path.splitext(lora_name)[0]
|
|
return os.path.join(self.get_cache_dir(), f"{base_name}_blocks.yaml")
|
|
|
|
def get_block_settings(self, lora_name: str, use_block_cache: bool = True, include_single_blocks: bool = False) -> dict:
|
|
"""Load block settings from cache or return defaults"""
|
|
# Initialize with all double blocks enabled and single blocks based on parameter
|
|
default_settings = {
|
|
**{f"double_blocks.{i}.": True for i in range(20)},
|
|
**{f"single_blocks.{i}.": include_single_blocks for i in range(40)}
|
|
}
|
|
|
|
if not use_block_cache:
|
|
return default_settings
|
|
|
|
try:
|
|
settings_file = self.get_settings_filename(lora_name)
|
|
if os.path.exists(settings_file):
|
|
cached_settings = yaml.safe_load(open(settings_file, 'r'))
|
|
# Merge cached settings with default single block settings
|
|
return {
|
|
**default_settings,
|
|
**cached_settings,
|
|
**{f"single_blocks.{i}.": include_single_blocks for i in range(40)} # Override single blocks
|
|
}
|
|
return default_settings
|
|
except Exception as e:
|
|
logger.error(f"Failed to load block settings for {lora_name}: {str(e)}")
|
|
return default_settings
|
|
|
|
def save_block_settings(self, lora_name: str, block_settings: dict):
|
|
"""Save block settings to cache"""
|
|
try:
|
|
settings_file = self.get_settings_filename(lora_name)
|
|
# Ensure directory exists
|
|
os.makedirs(os.path.dirname(settings_file), exist_ok=True)
|
|
save_settings = {k: v for k, v in block_settings.items() if k.startswith('double_blocks.')}
|
|
with open(settings_file, 'w') as f:
|
|
yaml.safe_dump(save_settings, f)
|
|
except Exception as e:
|
|
logger.error(f"Failed to save block settings for {lora_name}: {str(e)}")
|
|
|
|
def filter_lora_keys(self, lora: Dict[str, torch.Tensor], block_settings: dict) -> Dict[str, torch.Tensor]:
|
|
"""Filter LoRA keys based on block settings"""
|
|
filtered_blocks = {k: v for k, v in block_settings.items() if v is True}
|
|
return {key: value for key, value in lora.items()
|
|
if any(block in key for block in filtered_blocks)}
|
|
|
|
def load_lora_file(self, lora_name: str) -> Dict[str, torch.Tensor]:
|
|
"""Load LoRA file and cache it"""
|
|
from comfy.utils import load_torch_file
|
|
|
|
lora_path = folder_paths.get_full_path("loras", lora_name)
|
|
if not os.path.exists(lora_path):
|
|
raise Exception(f"LoRA {lora_name} not found at {lora_path}")
|
|
|
|
if self.loaded_lora is not None and self.loaded_lora[0] == lora_path:
|
|
return self.loaded_lora[1]
|
|
|
|
lora = load_torch_file(lora_path)
|
|
self.loaded_lora = (lora_path, lora)
|
|
return lora
|
|
|
|
def get_file_mtime(self, filepath: str) -> str:
|
|
"""Get modification time of file as string"""
|
|
try:
|
|
return str(os.path.getmtime(filepath))
|
|
except:
|
|
return "0"
|
|
|
|
def get_lora_mtime(self, lora_name: str) -> str:
|
|
"""Get modification time of LoRA file"""
|
|
try:
|
|
lora_path = folder_paths.get_full_path("loras", lora_name)
|
|
return self.get_file_mtime(lora_path)
|
|
except:
|
|
return "0"
|
|
|
|
def get_cache_mtime(self, lora_name: str) -> str:
|
|
"""Get modification time of cache file"""
|
|
try:
|
|
cache_file = self.get_settings_filename(lora_name)
|
|
return self.get_file_mtime(cache_file)
|
|
except:
|
|
return "0"
|
|
|
|
def apply_lora(self, model, lora_name: str, strength: float, block_settings: Optional[dict] = None) -> torch.nn.Module:
|
|
"""Apply LoRA to model with given settings"""
|
|
from comfy.sd import load_lora_for_models
|
|
|
|
if not lora_name:
|
|
return model
|
|
|
|
try:
|
|
lora = self.load_lora_file(lora_name)
|
|
if block_settings is None:
|
|
block_settings = self.get_block_settings(lora_name, True) # Always use cache for direct loading
|
|
|
|
filtered_lora = self.filter_lora_keys(lora, block_settings)
|
|
new_model, _ = load_lora_for_models(model, None, filtered_lora, strength, 0)
|
|
return new_model if new_model is not None else model
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error applying LoRA {lora_name}: {str(e)}")
|
|
return model
|
|
|
|
class HunyuanLoadAndEditLoraBlocks(BaseHunyuanLoraLoader):
|
|
"""Interactive LoRA block editor"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
arg_dict = {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"lora_name": (folder_paths.get_filename_list("loras"),),
|
|
"strength": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": -10.0,
|
|
"max": 10.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"save_settings": ("BOOLEAN", {"default": True}),
|
|
"use_single_blocks": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
for i in range(20):
|
|
arg_dict["required"][f"double_blocks.{i}."] = ("BOOLEAN", {"default": True})
|
|
|
|
return arg_dict
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_lora"
|
|
CATEGORY = "loaders/hunyuan"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, model, lora_name: str, strength: float, save_settings: bool, use_single_blocks: bool, **kwargs):
|
|
instance = s()
|
|
lora_mtime = instance.get_lora_mtime(lora_name)
|
|
return f"{lora_name}_{strength}_{lora_mtime}"
|
|
|
|
def load_lora(self, model, lora_name: str, strength: float, save_settings: bool, use_single_blocks: bool, **kwargs):
|
|
if not lora_name:
|
|
return (model,)
|
|
|
|
# Add single blocks settings based on the parameter
|
|
block_settings = {
|
|
**kwargs,
|
|
**{f"single_blocks.{i}.": use_single_blocks for i in range(40)}
|
|
}
|
|
|
|
if save_settings:
|
|
self.save_block_settings(lora_name, block_settings)
|
|
return (self.apply_lora(model, lora_name, strength, block_settings),)
|
|
|
|
class HunyuanLoadFromBlockCache(BaseHunyuanLoraLoader):
|
|
"""Load LoRA using cached block settings"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"lora_name": (folder_paths.get_filename_list("loras"),),
|
|
"strength": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": -10.0,
|
|
"max": 10.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"use_single_blocks": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_lora"
|
|
CATEGORY = "loaders/hunyuan"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, model, lora_name: str, strength: float, use_single_blocks: bool):
|
|
instance = s()
|
|
lora_mtime = instance.get_lora_mtime(lora_name)
|
|
cache_mtime = instance.get_cache_mtime(lora_name)
|
|
return f"{lora_name}_{strength}_{lora_mtime}_{cache_mtime}"
|
|
|
|
def load_lora(self, model, lora_name: str, strength: float, use_single_blocks: bool):
|
|
block_settings = self.get_block_settings(lora_name, True, use_single_blocks)
|
|
return (self.apply_lora(model, lora_name, strength, block_settings),)
|
|
|
|
class HunyuanLoraFromJson(BaseHunyuanLoraLoader):
|
|
"""Load multiple LoRAs from JSON configuration"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"json_data": ("STRING", {"multiline": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "process_json"
|
|
CATEGORY = "loaders/hunyuan"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, model, json_data: str):
|
|
try:
|
|
instance = s()
|
|
lora_configs = json.loads(json_data)
|
|
if not isinstance(lora_configs, list):
|
|
return json_data
|
|
|
|
mtimes = []
|
|
for config in lora_configs:
|
|
if isinstance(config, dict) and 'filename' in config and 'strength' in config:
|
|
lora_mtime = instance.get_lora_mtime(config['filename'])
|
|
use_cache = config.get('use_block_cache', True)
|
|
cache_mtime = instance.get_cache_mtime(config['filename']) if use_cache else "0"
|
|
mtimes.append(f"{config['filename']}_{config['strength']}_{lora_mtime}_{cache_mtime}")
|
|
|
|
return "_".join(mtimes) if mtimes else json_data
|
|
except:
|
|
return json_data
|
|
|
|
def process_json(self, model, json_data: str):
|
|
try:
|
|
lora_configs = json.loads(json_data)
|
|
if not isinstance(lora_configs, list):
|
|
logger.error("JSON data must be an array of LoRA configurations")
|
|
return (model,)
|
|
|
|
current_model = model
|
|
for config in lora_configs:
|
|
if not isinstance(config, dict) or 'filename' not in config or 'strength' not in config:
|
|
logger.error(f"Invalid LoRA config: {config}")
|
|
continue
|
|
|
|
logger.info(f"Applying LoRA: {config['filename']}")
|
|
use_cache = config.get('use_block_cache', True)
|
|
use_single = config.get('use_single_blocks', False)
|
|
block_settings = self.get_block_settings(config['filename'], use_cache, use_single)
|
|
current_model = self.apply_lora(
|
|
current_model,
|
|
config['filename'],
|
|
float(config['strength']),
|
|
block_settings
|
|
)
|
|
|
|
return (current_model,)
|
|
|
|
except json.JSONDecodeError as e:
|
|
logger.error(f"Invalid JSON data: {str(e)}")
|
|
return (model,)
|
|
except Exception as e:
|
|
logger.error(f"Error processing LoRAs: {str(e)}")
|
|
return (model,)
|
|
|
|
class HunyuanLoraFromPrompt(BaseHunyuanLoraLoader):
|
|
"""Load LoRAs based on trigger phrases in the prompt"""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._default_config = "lora_triggers.yaml"
|
|
self._config_file = None
|
|
|
|
@property
|
|
def config_file(self):
|
|
"""Get the full path to the config file"""
|
|
if not self._config_file:
|
|
return os.path.join(self.get_cache_dir(), self._default_config)
|
|
return os.path.join(self.get_cache_dir(), self._config_file)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"prompt": ("STRING", {"multiline": True}),
|
|
"config_file": ("STRING", {
|
|
"default": "lora_triggers.yaml",
|
|
"multiline": False
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "process_prompt"
|
|
CATEGORY = "loaders/hunyuan"
|
|
|
|
def create_default_config(self):
|
|
"""Create a default configuration file with example"""
|
|
default_config = """# LoRA Trigger Configuration
|
|
# Each LoRA entry should have:
|
|
# - filename: The name of the LoRA file
|
|
# - strength: The strength to apply the LoRA (typically between 0.0 and 1.0)
|
|
# - use_block_cache: Whether to use cached block settings (default: true)
|
|
# - use_single_blocks: Whether to enable single block layers (default: false)
|
|
# - keywords: List of trigger words or phrases that will activate this LoRA
|
|
# (phrases can contain multiple words like "red hair" or "wearing glasses")
|
|
|
|
loras:
|
|
# Example LoRA configuration (uncomment and modify to use):
|
|
# - filename: example_lora.safetensors
|
|
# strength: 0.75
|
|
# use_block_cache: true
|
|
# use_single_blocks: true
|
|
# keywords:
|
|
# - red hair
|
|
# - wearing glasses
|
|
# - simple keyword
|
|
"""
|
|
try:
|
|
with open(self.config_file, 'w') as f:
|
|
f.write(default_config)
|
|
except Exception as e:
|
|
logger.error(f"Failed to create default config: {str(e)}")
|
|
|
|
def load_config(self) -> List[dict]:
|
|
"""Load trigger configuration file"""
|
|
try:
|
|
if not os.path.exists(self.config_file):
|
|
logger.info("Creating default LoRA triggers config file")
|
|
self.create_default_config()
|
|
return []
|
|
|
|
with open(self.config_file, 'r') as f:
|
|
config = yaml.safe_load(f)
|
|
|
|
if not config or 'loras' not in config:
|
|
logger.warning("Config file exists but contains no LoRA configurations")
|
|
return []
|
|
|
|
return config['loras']
|
|
except Exception as e:
|
|
logger.error(f"Failed to load config file: {str(e)}")
|
|
return []
|
|
|
|
def get_matching_loras(self, prompt: str, config: List[dict]) -> List[dict]:
|
|
"""Find LoRAs with matching trigger phrases in the prompt"""
|
|
prompt = prompt.lower()
|
|
matching_loras = []
|
|
|
|
for lora_config in config:
|
|
if not all(key in lora_config for key in ['filename', 'strength', 'keywords']):
|
|
logger.warning(f"Skipping invalid LoRA config: {lora_config}")
|
|
continue
|
|
|
|
for keyword in lora_config['keywords']:
|
|
if keyword.lower() in prompt:
|
|
matching_loras.append(lora_config)
|
|
logger.info(f"Applying LoRA: {lora_config['filename']}")
|
|
break
|
|
|
|
return matching_loras
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, model, prompt: str, config_file: str):
|
|
instance = s()
|
|
instance._config_file = config_file
|
|
# Get trigger file modification time
|
|
trigger_mtime = instance.get_file_mtime(instance.config_file)
|
|
|
|
# Get matching LoRAs and their mtimes
|
|
config = instance.load_config()
|
|
matching_loras = instance.get_matching_loras(prompt, config)
|
|
|
|
mtimes = []
|
|
for lora_config in matching_loras:
|
|
lora_mtime = instance.get_lora_mtime(lora_config['filename'])
|
|
use_cache = lora_config.get('use_block_cache', True)
|
|
cache_mtime = instance.get_cache_mtime(lora_config['filename']) if use_cache else "0"
|
|
mtimes.append(f"{lora_config['filename']}_{lora_config['strength']}_{lora_mtime}_{cache_mtime}")
|
|
|
|
return f"{prompt}_{trigger_mtime}_{'_'.join(mtimes)}"
|
|
|
|
def process_prompt(self, model, prompt: str, config_file: str):
|
|
self._config_file = config_file
|
|
config = self.load_config()
|
|
if not config:
|
|
return (model,)
|
|
|
|
matching_loras = self.get_matching_loras(prompt, config)
|
|
if not matching_loras:
|
|
return (model,)
|
|
|
|
current_model = model
|
|
for lora_config in matching_loras:
|
|
use_cache = lora_config.get('use_block_cache', True)
|
|
use_single = lora_config.get('use_single_blocks', True)
|
|
block_settings = self.get_block_settings(lora_config['filename'], use_cache, use_single)
|
|
current_model = self.apply_lora(
|
|
current_model,
|
|
lora_config['filename'],
|
|
float(lora_config['strength']),
|
|
block_settings
|
|
)
|
|
|
|
return (current_model,)
|