From 2775e5c30a0a4f97225368ff7687991d2d9bce75 Mon Sep 17 00:00:00 2001 From: Scorpinaus Date: Sun, 25 Aug 2024 11:05:45 +0800 Subject: [PATCH] Replaced combined_safetensor.safetensor method with using config.json method for multiple-file TE or UNET/Transformer --- clip_loader.py | 29 ++++++++++++++++++++++++----- unet_loader.py | 50 ++++++++++++++++++++++++++------------------------ 2 files changed, 50 insertions(+), 29 deletions(-) diff --git a/clip_loader.py b/clip_loader.py index 315d9ea..bc2a5cd 100644 --- a/clip_loader.py +++ b/clip_loader.py @@ -105,9 +105,10 @@ class DiffusersClipLoader(DiffusersLoaderBase): if model_type == "Flux": # use the index file - index_files = [f for f in os.listdir(text_encoder_dir2) if f.endswith('index.json')] + index_files = [f for f in os.listdir(text_encoder_dir2) if f.endswith('index.json') or 'index.' in f and f.endswith('.json')] if index_files: - text_encoder_paths.append(os.path.join(text_encoder_dir2, index_files[0])) + index_file = os.path.join(text_encoder_dir2, index_files[0]) + text_encoder_paths.append(cls.load_flux_text_encoder_2(index_file)) 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)) @@ -117,9 +118,10 @@ class DiffusersClipLoader(DiffusersLoaderBase): if model_type == "SD3": text_encoder_dir3 = os.path.join(sub_dir_path, "text_encoder_3") - #For text_encoder_3, we will use the index file - index_file = os.path.join(text_encoder_dir3, "text_encoder_3_model.safetensors.index.fp16.json") - if os.path.exists(index_file): + #For text_encoder_3, we will use the index file. Look up for any file ending with 'index.json' or 'index.*.json' + index_files = [f for f in os.listdir(text_encoder_dir3) if f.endswith('index.json') or 'index.' in f and f.endswith('.json')] + if index_files: + index_file = os.path.join(text_encoder_dir3, index_files[0]) text_encoder_paths.append(cls.load_sd3_text_encoder_3(index_file)) else: #If index file not found @@ -139,6 +141,23 @@ class DiffusersClipLoader(DiffusersLoaderBase): weight_map = index_data['weight_map'] base_path = os.path.dirname(index_file) + sd = {} + for key, file_name in weight_map.items(): + file_path = os.path.join(base_path, file_name) + if os.path.exists(file_path): + part_sd = comfy.utils.load_torch_file(file_path, safe_load=True) + sd.update({key: part_sd[key] for key in part_sd if key in weight_map}) + + return sd + + @classmethod + def load_flux_text_encoder_2(cls, index_file): + with open(index_file, 'r') as f: + index_data = json.load(f) + + weight_map = index_data['weight_map'] + base_path = os.path.dirname(index_file) + sd = {} for key, file_name in weight_map.items(): file_path = os.path.join(base_path, file_name) diff --git a/unet_loader.py b/unet_loader.py index 9c6bac4..538b6a3 100644 --- a/unet_loader.py +++ b/unet_loader.py @@ -46,34 +46,16 @@ class DiffusersUNETLoader(DiffusersLoaderBase): 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) + state_dict = cls.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) + state_dict = cls.load_from_index(unet_path) print(f"State_dict obtained") model = cls.load_diffusion_model_from_state_dict(state_dict, model_options=model_options) @@ -89,12 +71,12 @@ class DiffusersUNETLoader(DiffusersLoaderBase): 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) + state_dict = cls.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) + state_dict = cls.load_from_index(unet_path) model = cls.load_diffusion_model_from_state_dict(state_dict, model_options=model_options) else: @@ -108,6 +90,25 @@ class DiffusersUNETLoader(DiffusersLoaderBase): print(f"DiffusersUNETLoader: Error loading UNET/Transformer model: {e}") raise + @classmethod + def load_from_index(cls, 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 + @classmethod def load_diffusion_model_from_state_dict(cls, state_dict, model_options): print("Running load_diffusion_model_from_state_dict function...") @@ -154,7 +155,7 @@ class DiffusersUNETLoader(DiffusersLoaderBase): @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')] + index_files = [f for f in os.listdir(unet_folder) if f.endswith('index.json') or 'index.' in f and f.endswith('.json')] print(f"Index files: {index_files}") if index_files: index_file = os.path.join(unet_folder, index_files[0]) @@ -185,7 +186,7 @@ class DiffusersUNETLoader(DiffusersLoaderBase): @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')] + index_files = [f for f in os.listdir(unet_folder) if f.endswith('index.json') or 'index.' in f and f.endswith('.json')] if index_files: index_file = os.path.join(unet_folder, index_files[0]) print(f"Found index file: {index_file}") @@ -198,6 +199,7 @@ class DiffusersUNETLoader(DiffusersLoaderBase): 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] + print(f"Part files: {part_files}") if part_files: return os.path.join(unet_folder, part_files[0])