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

Improved on model detection to include extra_model_paths.yaml

Browse Source
This commit is contained in:
Scorpinaus
2024-08-26 21:43:54 +08:00
parent 1a3da3a1a2
commit f9d8bb9592
6 changed files with 95 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+15 -5
View File
@@ -28,7 +28,7 @@ Follow these steps to install and set up the project:
- Add Nodes -> DiffusersLoader -> Combined -> CombinedDiffusersLoader
- ![image](https://github.com/user-attachments/assets/79a576b2-dc27-49e4-a7c7-ec0e01bc5bad)
3. Connect the DiffusersLoader
3. Connect the DiffusersLoader to your workflow
@@ -48,12 +48,16 @@ The combined loader work in the same manner as existing checkpoints loader as se
Select the name of the diffusers checkpoint folder from the dropdown list and connect the nodes.
Take note for SD3, AuraFlow and Flux:
- The first run through will be slightly slower due to creation of a combined_safetensors.safetensors
## Limitations & Future Improvements
- Add support for other compatible diffusers format checkpoints
- For AuraFlow and SD3, the loading and combining for the combined_model.safetensors will have a delay for low-vram or low-system ram machines
- Future model_types:
- HunyuanDiT
- PixArt
- Kolors
Note: This will need to raise a PR on main comfyUI repo
- For Flux, use of default weight_dtype would OOM (Use fp8-based d-types)
## Contributing
@@ -67,5 +71,11 @@ Hafiz Saffie - hafiz.safs@gmail.com
Project Link: https://github.com/Scorpinaus/ComfyUI-DiffusersLoader
## Version History
## Update History
26/8 - Improved on model detection to include extra_model_paths.yaml.
25/8 - Implemented multi-file loading support where the combined_safetensor.safetensor file method is used as a back-up
21/8 - Compatibility Update With ComfyUI, Addition of Flux Support
30/7 - Added repo to comfyUI registry
28/7 - Added AuraFlow support
20/6 - Added SD3 support
2/6 - Initial commit
clip_loader.py
+13 -27
View File
@@ -10,9 +10,10 @@ from .utils import DiffusersUtils
class DiffusersClipLoader(DiffusersLoaderBase):
@classmethod
def INPUT_TYPES(cls):
model_directories = DiffusersUtils.get_model_directories()
return {
"required": {
"sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),),
"sub_directory": (model_directories,),
"clip_type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "sdxl", "flux"],),
}
}
@@ -28,8 +29,12 @@ class DiffusersClipLoader(DiffusersLoaderBase):
@classmethod
def load_model(cls, sub_directory, clip_type="stable_diffusion"):
base_path = DiffusersUtils.get_base_path()
sub_dir_path = os.path.join(base_path, sub_directory)
base_paths = DiffusersUtils.get_base_path()
sub_dir_path = next((os.path.join(base_path, sub_directory) for base_path in base_paths if os.path.exists(os.path.join(base_path, sub_directory))), None)
if sub_dir_path is None:
raise ValueError(f"Subdirectory '{sub_directory}' not found in any of the diffusers paths.")
model_type = cls.detect_model_type(sub_dir_path)
clip_type_enum = cls.get_clip_type_enum(clip_type)
@@ -105,10 +110,9 @@ 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') or 'index.' in f and f.endswith('.json')]
index_files = [f for f in os.listdir(text_encoder_dir2) if f.endswith('index.json')]
if index_files:
index_file = os.path.join(text_encoder_dir2, index_files[0])
text_encoder_paths.append(cls.load_flux_text_encoder_2(index_file))
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))
@@ -118,10 +122,9 @@ 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. 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])
#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):
text_encoder_paths.append(cls.load_sd3_text_encoder_3(index_file))
else:
#If index file not found
@@ -141,23 +144,6 @@ 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)
combined_diffusers_loader.py
+7 -5
View File
@@ -9,8 +9,7 @@ import os
class CombinedDiffusersLoader:
@classmethod
def INPUT_TYPES(cls):
base_path = DiffusersUtils.get_base_path()
model_directories = DiffusersUtils.get_model_directories(base_path)
model_directories = DiffusersUtils.get_model_directories()
return {
"required": {
"sub_directory": (model_directories,),
@@ -28,12 +27,15 @@ class CombinedDiffusersLoader:
@classmethod
def load_models(cls, sub_directory, clip_type="stable_diffusion", transformer_parts="all", vae_type="default", weight_dtype="default"):
base_path = DiffusersUtils.get_base_path()
sub_dir_path = os.path.join(base_path, sub_directory)
base_paths = DiffusersUtils.get_base_path()
sub_dir_path = next((os.path.join(base_path, sub_directory) for base_path in base_paths if os.path.exists(os.path.join(base_path, sub_directory))), None)
if sub_dir_path is None:
raise ValueError(f"Subdirectory '{sub_directory}' not found in any of the diffusers paths.")
model_type = DiffusersLoaderBase.detect_model_type(sub_dir_path)
print(f"Detected model type: {model_type}")
unet_model = DiffusersUNETLoader.load_model(sub_directory, transformer_parts, weight_dtype)
unet_model, = DiffusersUNETLoader.load_model(sub_directory, transformer_parts, weight_dtype)
clip_model = DiffusersClipLoader.load_model(sub_directory, clip_type)
unet_loader.py
+38 -35
View File
@@ -10,16 +10,17 @@ import torch
class DiffusersUNETLoader(DiffusersLoaderBase):
@classmethod
def INPUT_TYPES(cls):
model_directories = DiffusersUtils.get_model_directories()
return {
"required": {
"sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),),
"sub_directory": (model_directories,),
"transformer_parts": (["all", "part_1", "part_2", "part_3"],),
"weight_dtype": (["default", "fp8_e4m3fn", "fp8_e5m2"],),
}
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "load_unet/transformer"
FUNCTION = "load_model"
CATEGORY = "DiffusersLoader"
@classmethod
@@ -28,8 +29,13 @@ class DiffusersUNETLoader(DiffusersLoaderBase):
@classmethod
def load_model(cls, sub_directory, transformer_parts="all", weight_dtype="default"):
base_path = DiffusersUtils.get_base_path()
sub_dir_path = os.path.join(base_path, sub_directory)
base_paths = DiffusersUtils.get_base_path()
sub_dir_path = next((os.path.join(base_path, sub_directory) for base_path in base_paths if os.path.exists(os.path.join(base_path, sub_directory))), None)
if sub_dir_path is None:
raise ValueError(f"Subdirectory '{sub_directory}' not found in any of the diffusers paths.")
model_type = cls.detect_model_type(sub_dir_path)
print(f"DiffusersUNETLoader: Detected model type: {model_type}")
@@ -46,23 +52,42 @@ 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 = cls.load_from_index(unet_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 = cls.load_from_index(unet_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)
else:
model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
print("UNET/Transformer model loaded successfully")
return model
#Return model as a tuple
return (model,)
except RuntimeError as e:
if "out of memory" in str(e):
@@ -71,18 +96,19 @@ 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 = cls.load_from_index(unet_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 = cls.load_from_index(unet_path)
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
#Return model as a tuple
return (model,)
except Exception as retry_e:
print(f"Error loading model after memory clear: {retry_e}")
raise
@@ -90,28 +116,6 @@ 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']
base_path = os.path.dirname(index_path)
state_dict = {}
for key, file_name in weight_map.items():
file_path = os.path.join(base_path, file_name)
if os.path.exists(file_path):
part_dict = comfy.utils.load_torch_file(file_path, safe_load=True)
state_dict.update({k: part_dict[k] for k in part_dict if k in weight_map})
else:
print(f"WARNING: File not found: {file_path}")
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...")
@@ -158,7 +162,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') or 'index.' in f and f.endswith('.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])
@@ -189,7 +193,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') or 'index.' in f and f.endswith('.json')]
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}")
@@ -202,7 +206,6 @@ 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])
utils.py
+15 -9
View File
@@ -2,6 +2,9 @@ import os
import safetensors.torch
import torch
import gc
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
import folder_paths
class DiffusersUtils:
_model_cache = {'clip': None, 'unet': None, 'vae': None}
@@ -43,19 +46,22 @@ class DiffusersUtils:
@staticmethod
def get_base_path():
base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'diffusers'))
if not os.path.exists(base_path):
base_path = folder_paths.get_folder_paths("diffusers")
#print(f"Base path: {base_path}")
if not base_path:
raise FileNotFoundError(f"The base path '{base_path}' does not exist.")
return base_path
@staticmethod
def get_model_directories(base_path):
model_dirs = []
for root, dirs, files in os.walk(base_path):
if "model_index.json" in files:
model_dirs.append(os.path.relpath(root, base_path))
dirs[:] = [d for d in dirs if d not in ["vae", "unet", "text_encoder", "text_encoder_2", "text_encoder_3", "transformer"]]
return model_dirs
def get_model_directories():
paths = []
base_paths = DiffusersUtils.get_base_path()
for base_path in base_paths:
if os.path.exists(base_path):
for root, _, files in os.walk(base_path, followlinks=True):
if "model_index.json" in files:
paths.append(os.path.relpath(root, start=base_path))
return paths
@staticmethod
def find_model_files(directory, file_parts=None):
vae_loader.py
+7 -3
View File
@@ -16,9 +16,10 @@ class DiffusersVAELoader(DiffusersLoaderBase):
}
@classmethod
def INPUT_TYPES(cls):
model_directories = DiffusersUtils.get_model_directories()
return {
"required": {
"sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),),
"sub_directory": (model_directories,),
"vae_type": (["default"] + list(cls.VAE_CONFIGS.keys()),),
}
}
@@ -40,8 +41,11 @@ class DiffusersVAELoader(DiffusersLoaderBase):
@classmethod
def load_default_vae(cls, sub_directory):
base_path = DiffusersUtils.get_base_path()
sub_dir_path = os.path.join(base_path, sub_directory)
base_paths = DiffusersUtils.get_base_path()
sub_dir_path = next((os.path.join(base_path, sub_directory) for base_path in base_paths if os.path.exists(os.path.join(base_path, sub_directory))), None)
if sub_dir_path is None:
raise ValueError(f"Subdirectory '{sub_directory}' not found in any of the diffusers paths.")
vae_folder = os.path.join(sub_dir_path, "vae")
vae_path = DiffusersUtils.find_model_file(vae_folder)
DiffusersUtils.check_and_clear_cache('vae', vae_path)