From 090854a52100cab2dd30da8634763e1f382d0bfc Mon Sep 17 00:00:00 2001 From: Scorpinaus Date: Sun, 25 Aug 2024 10:20:30 +0800 Subject: [PATCH] Removed combined_text_encoder/transformer step and used index.json method instead --- clip_loader.py | 46 +++++++++++++++++++++------- unet_loader.py | 83 ++++++++++++++++++++++++++++++++++++++------------ utils.py | 11 +++---- 3 files changed, 104 insertions(+), 36 deletions(-) diff --git a/clip_loader.py b/clip_loader.py index 45c27b6..315d9ea 100644 --- a/clip_loader.py +++ b/clip_loader.py @@ -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") diff --git a/unet_loader.py b/unet_loader.py index f612370..9c6bac4 100644 --- a/unet_loader.py +++ b/unet_loader.py @@ -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) diff --git a/utils.py b/utils.py index 0168a46..170e0cf 100644 --- a/utils.py +++ b/utils.py @@ -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}")