mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-22 05:32:12 -03:00
b99d78bda6
- Allow direct file paths in addition to registered LoRA names - Add graceful handling for missing LoRA files with warning logs - Add comprehensive unit tests for missing LoRA file handling - Ensure backward compatibility with existing LoRA loading behavior
134 lines
5.0 KiB
Python
134 lines
5.0 KiB
Python
class AnyType(str):
|
|
"""A special class that is always equal in not equal comparisons. Credit to pythongosssss"""
|
|
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
# Credit to Regis Gaughan, III (rgthree)
|
|
class FlexibleOptionalInputType(dict):
|
|
"""A special class to make flexible nodes that pass data to our python handlers.
|
|
|
|
Enables both flexible/dynamic input types (like for Any Switch) or a dynamic number of inputs
|
|
(like for Any Switch, Context Switch, Context Merge, Power Lora Loader, etc).
|
|
|
|
Note, for ComfyUI, all that's needed is the `__contains__` override below, which tells ComfyUI
|
|
that our node will handle the input, regardless of what it is.
|
|
|
|
However, with https://github.com/comfyanonymous/ComfyUI/pull/2666 a large change would occur
|
|
requiring more details on the input itself. There, we need to return a list/tuple where the first
|
|
item is the type. This can be a real type, or use the AnyType for additional flexibility.
|
|
|
|
This should be forwards compatible unless more changes occur in the PR.
|
|
"""
|
|
def __init__(self, type):
|
|
self.type = type
|
|
|
|
def __getitem__(self, key):
|
|
return (self.type, )
|
|
|
|
def __contains__(self, key):
|
|
return True
|
|
|
|
|
|
any_type = AnyType("*")
|
|
|
|
# Common methods extracted from lora_loader.py and lora_stacker.py
|
|
import os
|
|
import logging
|
|
import copy
|
|
import folder_paths
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def extract_lora_name(lora_path):
|
|
"""Extract the lora name from a lora path (e.g., 'IL\\aorunIllstrious.safetensors' -> 'aorunIllstrious')"""
|
|
# Get the basename without extension
|
|
basename = os.path.basename(lora_path)
|
|
return os.path.splitext(basename)[0]
|
|
|
|
def get_loras_list(kwargs):
|
|
"""Helper to extract loras list from either old or new kwargs format"""
|
|
if 'loras' not in kwargs:
|
|
return []
|
|
|
|
loras_data = kwargs['loras']
|
|
# Handle new format: {'loras': {'__value__': [...]}}
|
|
if isinstance(loras_data, dict) and '__value__' in loras_data:
|
|
return loras_data['__value__']
|
|
# Handle old format: {'loras': [...]}
|
|
elif isinstance(loras_data, list):
|
|
return loras_data
|
|
# Unexpected format
|
|
else:
|
|
logger.warning(f"Unexpected loras format: {type(loras_data)}")
|
|
return []
|
|
|
|
def load_state_dict_in_safetensors(path, device="cpu", filter_prefix=""):
|
|
"""Simplified version of load_state_dict_in_safetensors that just loads from a local path"""
|
|
import safetensors.torch
|
|
|
|
state_dict = {}
|
|
with safetensors.torch.safe_open(path, framework="pt", device=device) as f:
|
|
for k in f.keys():
|
|
if filter_prefix and not k.startswith(filter_prefix):
|
|
continue
|
|
state_dict[k.removeprefix(filter_prefix)] = f.get_tensor(k)
|
|
return state_dict
|
|
|
|
def to_diffusers(input_lora):
|
|
"""Simplified version of to_diffusers for Flux LoRA conversion"""
|
|
import torch
|
|
from diffusers.utils.state_dict_utils import convert_unet_state_dict_to_peft
|
|
from diffusers.loaders import FluxLoraLoaderMixin
|
|
|
|
if isinstance(input_lora, str):
|
|
tensors = load_state_dict_in_safetensors(input_lora, device="cpu")
|
|
else:
|
|
tensors = {k: v for k, v in input_lora.items()}
|
|
|
|
# Convert FP8 tensors to BF16
|
|
for k, v in tensors.items():
|
|
if v.dtype not in [torch.float64, torch.float32, torch.bfloat16, torch.float16]:
|
|
tensors[k] = v.to(torch.bfloat16)
|
|
|
|
new_tensors = FluxLoraLoaderMixin.lora_state_dict(tensors)
|
|
new_tensors = convert_unet_state_dict_to_peft(new_tensors)
|
|
|
|
return new_tensors
|
|
|
|
def nunchaku_load_lora(model, lora_name, lora_strength):
|
|
"""Load a Flux LoRA for Nunchaku model"""
|
|
model_wrapper = model.model.diffusion_model
|
|
transformer = model_wrapper.model
|
|
|
|
# Save the transformer temporarily
|
|
model_wrapper.model = None
|
|
ret_model = copy.deepcopy(model) # copy everything except the model
|
|
ret_model_wrapper = ret_model.model.diffusion_model
|
|
|
|
# Restore the model and set it for the copy
|
|
model_wrapper.model = transformer
|
|
ret_model_wrapper.model = transformer
|
|
|
|
# Get full path to the LoRA file. Allow both direct paths and registered LoRA names.
|
|
lora_path = lora_name if os.path.isfile(lora_name) else folder_paths.get_full_path("loras", lora_name)
|
|
if not lora_path or not os.path.isfile(lora_path):
|
|
logger.warning("Skipping LoRA '%s' because it could not be found", lora_name)
|
|
return model
|
|
|
|
ret_model_wrapper.loras.append((lora_path, lora_strength))
|
|
|
|
# Convert the LoRA to diffusers format
|
|
sd = to_diffusers(lora_path)
|
|
|
|
# Handle embedding adjustment if needed
|
|
if "transformer.x_embedder.lora_A.weight" in sd:
|
|
new_in_channels = sd["transformer.x_embedder.lora_A.weight"].shape[1]
|
|
assert new_in_channels % 4 == 0
|
|
new_in_channels = new_in_channels // 4
|
|
|
|
old_in_channels = ret_model.model.model_config.unet_config["in_channels"]
|
|
if old_in_channels < new_in_channels:
|
|
ret_model.model.model_config.unet_config["in_channels"] = new_in_channels
|
|
|
|
return ret_model |