228 lines
10 KiB
Python
228 lines
10 KiB
Python
# unet_loader.py
|
|
import os
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import json
|
|
from .base_loader import DiffusersLoaderBase
|
|
from .utils import DiffusersUtils
|
|
import torch
|
|
|
|
class DiffusersUNETLoader(DiffusersLoaderBase):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"sub_directory": (DiffusersUtils.get_unique_display_names(DiffusersUtils.get_model_directories())[0],),
|
|
"transformer_parts": (["all", "part_1", "part_2", "part_3"],),
|
|
"weight_dtype": (["default", "fp8_e4m3fn", "fp8_e5m2"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "DiffusersLoader"
|
|
|
|
@classmethod
|
|
def load_unet(cls, sub_directory, transformer_parts="all", weight_dtype="default"):
|
|
return (cls.load_model(sub_directory, transformer_parts, weight_dtype),)
|
|
|
|
@classmethod
|
|
def load_model(cls, sub_directory, transformer_parts="all", weight_dtype="default"):
|
|
if os.path.exists(sub_directory):
|
|
full_path = sub_directory
|
|
else:
|
|
model_directories = DiffusersUtils.get_model_directories()
|
|
_, unique_names = DiffusersUtils.get_unique_display_names(model_directories)
|
|
|
|
if "(" in sub_directory:
|
|
dir_name, index = sub_directory.rsplit(" (", 1)
|
|
index = int(index[:-1]) - 1
|
|
full_path = unique_names[dir_name][index]
|
|
else:
|
|
full_path = unique_names[sub_directory][0]
|
|
|
|
if not os.path.exists(full_path):
|
|
raise ValueError(f"Selected directory does not exist: {full_path}")
|
|
|
|
model_type = cls.detect_model_type(full_path)
|
|
print(f"DiffusersUNETLoader: Detected model type: {model_type}")
|
|
|
|
unet_path = cls.get_unet_path(full_path, model_type, transformer_parts)
|
|
DiffusersUtils.check_and_clear_cache('unet', unet_path)
|
|
|
|
print(f"DiffusersUNETLoader: Attempting to load UNET model from: {unet_path}")
|
|
|
|
model_options = {}
|
|
if weight_dtype == "fp8_e4m3fn":
|
|
model_options["dtype"] = torch.float8_e4m3fn
|
|
elif weight_dtype == "fp8_e5m2":
|
|
model_options["dtype"] = torch.float8_e5m2
|
|
|
|
print(f"model_options: {model_options}")
|
|
|
|
def load_from_index(index_path):
|
|
print(f"Load_from_index: index_path: {index_path}")
|
|
|
|
with open(index_path, 'r') as f:
|
|
index_data = json.load(f)
|
|
weight_map = index_data['weight_map']
|
|
state_dict = {}
|
|
|
|
for key, file in weight_map.items():
|
|
file_path = os.path.join(os.path.dirname(index_path), file)
|
|
part_dict = comfy.utils.load_torch_file(file_path, safe_load=True)
|
|
state_dict.update(part_dict)
|
|
del part_dict
|
|
#torch.cuda.empty_cache()
|
|
|
|
print(f"Load from index function completed. Returning state_dict")
|
|
return state_dict
|
|
|
|
try:
|
|
if model_type == "AuraFlow" and transformer_parts == "all" and unet_path.endswith('index.json'):
|
|
|
|
state_dict = load_from_index(unet_path)
|
|
print(f"State_dict obtained")
|
|
model = cls.load_diffusion_model_from_state_dict(state_dict, model_options=model_options)
|
|
|
|
elif model_type == "Flux" and transformer_parts == "all" and unet_path.endswith('index.json'):
|
|
|
|
state_dict = load_from_index(unet_path)
|
|
print(f"State_dict obtained")
|
|
model = cls.load_diffusion_model_from_state_dict(state_dict, model_options=model_options)
|
|
|
|
else:
|
|
model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
|
|
print("UNET/Transformer model loaded successfully")
|
|
#Return model as a tuple
|
|
return (model,)
|
|
|
|
except RuntimeError as e:
|
|
if "out of memory" in str(e):
|
|
print("Out of memory error. Attempting to clear memory and retry...")
|
|
DiffusersUtils.clear_memory()
|
|
try:
|
|
if model_type == "AuraFlow" and transformer_parts == "all" and unet_path.endswith('index.json'):
|
|
# Load model using the index file (retry)
|
|
state_dict = load_from_index(unet_path)
|
|
model = cls.load_diffusion_model_from_state_dict(state_dict, model_options=model_options)
|
|
|
|
elif model_type == "Flux" and transformer_parts == "all" and unet_path.endswith('index.json'):
|
|
|
|
state_dict = load_from_index(unet_path)
|
|
model = cls.load_diffusion_model_from_state_dict(state_dict, model_options=model_options)
|
|
|
|
else:
|
|
model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
|
|
print("UNET/Transformer model loaded successfully after memory clear")
|
|
#Return model as a tuple
|
|
return (model,)
|
|
except Exception as retry_e:
|
|
print(f"Error loading model after memory clear: {retry_e}")
|
|
raise
|
|
else:
|
|
print(f"DiffusersUNETLoader: Error loading UNET/Transformer model: {e}")
|
|
raise
|
|
|
|
@classmethod
|
|
def load_diffusion_model_from_state_dict(cls, state_dict, model_options):
|
|
print("Running load_diffusion_model_from_state_dict function...")
|
|
model = comfy.sd.load_diffusion_model_state_dict(state_dict, model_options=model_options)
|
|
return model
|
|
|
|
@classmethod
|
|
def get_unet_path(cls, full_path, model_type, transformer_parts):
|
|
if model_type in ["AuraFlow", "Flux", "SD3"]:
|
|
unet_folder = os.path.join(full_path, 'transformer')
|
|
print(f"Using transformer folder for {model_type}: {unet_folder}")
|
|
if model_type == "SD3":
|
|
return cls.handle_sd3_transformer(unet_folder)
|
|
if model_type == "AuraFlow":
|
|
return cls.handle_auraflow_transformer(unet_folder, transformer_parts)
|
|
elif model_type == "Flux":
|
|
return cls.handle_flux_transformer(unet_folder, transformer_parts)
|
|
else: # SD15, SDXL, etc.
|
|
unet_folder = os.path.join(full_path, 'unet')
|
|
if not os.path.exists(unet_folder):
|
|
#Fallback to transformer folder if unet folder does not exist
|
|
unet_folder = os.path.join(full_path, 'transformer')
|
|
if os.path.exists(unet_folder):
|
|
print(f"Unet folder not found, using transformer folder: {unet_folder}")
|
|
return cls.handle_flux_transformer(unet_folder, transformer_parts)
|
|
return DiffusersUtils.find_model_file(unet_folder)
|
|
|
|
@classmethod
|
|
def handle_transformer(cls, unet_folder, transformer_parts, num_parts):
|
|
if transformer_parts == "all":
|
|
combined_file_path = os.path.join(unet_folder, "combined_transformer.safetensors")
|
|
if not os.path.exists(combined_file_path):
|
|
print(f"Combined file not found, creating at: {combined_file_path}")
|
|
try:
|
|
combined_file_path = DiffusersUtils.combine_safetensor_files(unet_folder, os.path.dirname(unet_folder), num_parts=num_parts)
|
|
except Exception as e:
|
|
print(f"Error combining safetensor files: {e}")
|
|
combined_file_path = DiffusersUtils.find_model_files(unet_folder, num_parts=1)[0]
|
|
return combined_file_path
|
|
else:
|
|
part_num = int(transformer_parts.split('_')[1])
|
|
return DiffusersUtils.find_model_files(unet_folder, num_parts=part_num)[-1]
|
|
|
|
@classmethod
|
|
def handle_auraflow_transformer(cls, unet_folder, transformer_parts):
|
|
#Find file ending with 'index.json'
|
|
index_files = [f for f in os.listdir(unet_folder) if f.endswith('index.json')]
|
|
print(f"Index files: {index_files}")
|
|
if index_files:
|
|
index_file = os.path.join(unet_folder, index_files[0])
|
|
print(f"Found index file: {index_file}")
|
|
|
|
if transformer_parts == "all":
|
|
return index_file
|
|
|
|
else:
|
|
with open(index_file, 'r') as f:
|
|
index_data = json.load(f)
|
|
|
|
part_num = int(transformer_parts.split('_')[1])
|
|
part_files = [file for file in index_data['weight_map'].values() if f"-0000{part_num}-of-" in file]
|
|
|
|
if part_files:
|
|
return os.path.join(unet_folder, part_files[0])
|
|
else:
|
|
print(f"No file found for part {part_num}")
|
|
return None
|
|
else:
|
|
print(f"No index file found in {unet_folder}. Checking for combined transformer step")
|
|
return cls.handle_transformer(unet_folder, transformer_parts, num_parts=2)
|
|
|
|
@classmethod
|
|
def handle_sd3_transformer(cls, unet_folder):
|
|
return DiffusersUtils.find_model_file(unet_folder)
|
|
|
|
@classmethod
|
|
def handle_flux_transformer(cls, unet_folder, transformer_parts):
|
|
index_files = [f for f in os.listdir(unet_folder) if f.endswith('index.json')]
|
|
if index_files:
|
|
index_file = os.path.join(unet_folder, index_files[0])
|
|
print(f"Found index file: {index_file}")
|
|
|
|
if transformer_parts == "all":
|
|
return index_file
|
|
else:
|
|
with open(index_file, 'r') as f:
|
|
index_data = json.load(f)
|
|
|
|
part_num = int(transformer_parts.split('_')[1])
|
|
part_files = [file for file in index_data['weight_map'].values() if f"-0000{part_num}-of-" in file]
|
|
|
|
if part_files:
|
|
return os.path.join(unet_folder, part_files[0])
|
|
else:
|
|
print(f"No file found for part {part_num}")
|
|
return None
|
|
|
|
else:
|
|
print(f"No index file found in {unet_folder}. Checking for combined transformer step")
|
|
return cls.handle_transformer(unet_folder, transformer_parts, num_parts=3)
|
|
|