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

Replaced combined_safetensor.safetensor method with using config.json method for multiple-file TE or UNET/Transformer

Browse Source
This commit is contained in:
Scorpinaus
2024-08-25 11:05:45 +08:00
parent 090854a521
commit 2775e5c30a
2 changed files with 50 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
clip_loader.py
+24 -5
View File
@@ -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)
unet_loader.py
+26 -24
View File
@@ -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])