nodes/Scorpinaus-ComfyUI-DiffusersLoader
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Removed combined_text_encoder/transformer step and used index.json method instead

Browse Source
This commit is contained in:
Scorpinaus
2024-08-25 10:20:30 +08:00
parent 6934773cab
commit 090854a521
3 changed files with 104 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
clip_loader.py
+35 -11
View File
@@ -1,6 +1,7 @@
# clip_loader.py
import os
import json
import torch
import comfy.sd
import comfy.utils
from .base_loader import DiffusersLoaderBase
@@ -37,11 +38,30 @@ class DiffusersClipLoader(DiffusersLoaderBase):
clip_data = []
for path in text_encoder_paths:
if isinstance(path, dict):
clip_data.append(path)
else:
DiffusersUtils.check_and_clear_cache('clip', path)
clip_data.append(comfy.utils.load_torch_file(path, safe_load=True))
try:
if isinstance(path, dict):
clip_data.append(path)
elif path.endswith('index.json'):
with open(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(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()
clip_data.append(state_dict)
else:
DiffusersUtils.check_and_clear_cache('clip', path)
clip_data.append(comfy.utils.load_torch_file(path, safe_load=True))
except Exception as e:
print(f"Error loading clip model part from {path}: {e}")
raise
print(f"DiffusersClipLoader: Loading CLIP model(s) from: {text_encoder_paths}")
@@ -54,6 +74,9 @@ class DiffusersClipLoader(DiffusersLoaderBase):
if not hasattr(clip_model, 'tokenize'):
raise AttributeError("DiffusersClipLoader: Loaded clip model does not have 'tokenize' method.")
del clip_data
torch.cuda.empty_cache()
return clip_model
@@ -81,14 +104,15 @@ class DiffusersClipLoader(DiffusersLoaderBase):
text_encoder_paths.append(DiffusersUtils.find_model_file(text_encoder_dir1))
if model_type == "Flux":
# Combine split files for Flux models
combined_file_path = os.path.join(text_encoder_dir2, "combined_text_encoder.safetensors")
if not os.path.exists(combined_file_path):
combined_file_path = DiffusersUtils.combine_safetensor_files(text_encoder_dir2, sub_dir_path, num_parts=2)
text_encoder_paths.append(combined_file_path)
# use the index file
index_files = [f for f in os.listdir(text_encoder_dir2) if f.endswith('index.json')]
if index_files:
text_encoder_paths.append(os.path.join(text_encoder_dir2, index_files[0]))
else:
print(f"No index file found in {text_encoder_dir2}. Checking for combined text encoder step")
text_encoder_paths.append(DiffusersUtils.find_model_file(text_encoder_dir2))
else:
text_encoder_paths.append(DiffusersUtils.find_model_file(text_encoder_dir2))
if model_type == "SD3":
text_encoder_dir3 = os.path.join(sub_dir_path, "text_encoder_3")
unet_loader.py
+64 -19
View File
@@ -43,38 +43,60 @@ class DiffusersUNETLoader(DiffusersLoaderBase):
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('.json'):
#Load model using index_file
with open(unet_path, 'r') as f:
index_data = json.load(f)
weight_map = index_data['weight_map']
state_dict = {}
if model_type == "AuraFlow" and transformer_parts == "all" and unet_path.endswith('index.json'):
for key, file in weight_map.items():
file_path = os.path.join(os.path.dirname(unet_path), file)
state_dict.update(comfy.utils.load_torch_file(file_path))
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
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('.json'):
if model_type == "AuraFlow" and transformer_parts == "all" and unet_path.endswith('index.json'):
# Load model using the index file (retry)
with open(unet_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(unet_path), file)
state_dict.update(comfy.utils.load_torch_file(file_path))
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")
@@ -88,6 +110,7 @@ class DiffusersUNETLoader(DiffusersLoaderBase):
@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
@@ -162,5 +185,27 @@ class DiffusersUNETLoader(DiffusersLoaderBase):
@classmethod
def handle_flux_transformer(cls, unet_folder, transformer_parts):
return cls.handle_transformer(unet_folder, transformer_parts, num_parts=3)
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)
utils.py
+5 -6
View File
@@ -58,19 +58,17 @@ class DiffusersUtils:
return model_dirs
@staticmethod
def find_model_files(directory, file_parts=None, num_parts=None):
def find_model_files(directory, file_parts=None):
if not os.path.exists(directory):
raise FileNotFoundError(f"The directory '{directory}' does not exist.")
files = [os.path.join(directory, f) for f in os.listdir(directory)
if f.endswith((".safetensors", ".bin"))]
if f.endswith((".safetensors", ".bin", "index.json"))]
if not files:
raise FileNotFoundError(f"No .safetensors or .bin file found in {directory}")
raise FileNotFoundError(f"No .safetensors or .bin or index.json file found in {directory}")
if file_parts:
filtered_files = [file for file in files if any(part in file for part in file_parts)]
elif num_parts:
filtered_files = [file for file in files if any(f"0000{i+1}-of-" in file for i in range(num_parts))]
else:
filtered_files = files
print("Filtered files:", filtered_files)
@@ -82,7 +80,7 @@ class DiffusersUtils:
files = DiffusersUtils.find_model_files(directory)
if files:
return files[0]
raise FileNotFoundError(f"No .safetensors or .bin file found in {directory}")
raise FileNotFoundError(f"No .safetensors or .bin pr index.json file found in {directory}")
@staticmethod
def load_safetensor_paths(file_paths):
@@ -115,6 +113,7 @@ class DiffusersUtils:
print("Combining Safetensors completed")
return combined_tensors
#This function is not used. Will be removed in the subsequent release
@staticmethod
def combine_safetensor_files(directory, base_path, num_parts):
print(f"Running combine_safetensor_files for directory: {directory}")