commit af8096b4b4abe4be8d4cfbbd01f260c2853e7daf Author: Scorpinaus Date: Sun Jun 2 14:51:22 2024 +0800 Initial Commit diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..8ee1a42 --- /dev/null +++ b/__init__.py @@ -0,0 +1,13 @@ +from .loader import DiffusersLoader, SDXLDiffusersLoader + +NODE_CLASS_MAPPINGS = { + "Diffusers_Loader": DiffusersLoader, + "SDXL_Diffusers_Loader": SDXLDiffusersLoader, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "Diffusers_Loader": "Diffusers Loader", + "SDXL_Diffusers_Loader": "SDXL Diffusers Loader", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/__pycache__/__init__.cpython-311.pyc b/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000..35da619 Binary files /dev/null and b/__pycache__/__init__.cpython-311.pyc differ diff --git a/__pycache__/folder_paths.cpython-311.pyc b/__pycache__/folder_paths.cpython-311.pyc new file mode 100644 index 0000000..3d73fe4 Binary files /dev/null and b/__pycache__/folder_paths.cpython-311.pyc differ diff --git a/__pycache__/loader.cpython-311.pyc b/__pycache__/loader.cpython-311.pyc new file mode 100644 index 0000000..03250e1 Binary files /dev/null and b/__pycache__/loader.cpython-311.pyc differ diff --git a/__pycache__/utils.cpython-310.pyc b/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000..d8be3a3 Binary files /dev/null and b/__pycache__/utils.cpython-310.pyc differ diff --git a/__pycache__/utils.cpython-311.pyc b/__pycache__/utils.cpython-311.pyc new file mode 100644 index 0000000..35cd8d8 Binary files /dev/null and b/__pycache__/utils.cpython-311.pyc differ diff --git a/folder_paths.py b/folder_paths.py new file mode 100644 index 0000000..6a465d7 --- /dev/null +++ b/folder_paths.py @@ -0,0 +1,20 @@ +import os + +def get_comfyui_base_path(): + # Assuming the script is in the standard directory structure + return os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + +def get_folder_paths(folder_type): + base_path = get_comfyui_base_path() + + if folder_type == "diffusers": + return [os.path.join(base_path, "models", "diffusers")] + if folder_type == "embeddings": + return [os.path.join(base_path, "models", "embeddings")] + + return [] + +# Example usage: +if __name__ == "__main__": + print(get_folder_paths("diffusers")) # Should print the path to the diffusers models + print(get_folder_paths("embeddings")) # Should print the path to the embeddings diff --git a/loader.py b/loader.py new file mode 100644 index 0000000..cbd69f5 --- /dev/null +++ b/loader.py @@ -0,0 +1,56 @@ +import os +from .utils import load_diffusers_model, load_diffusers_sdxl_model +from .folder_paths import get_folder_paths + +class DiffusersLoader: + @classmethod + def INPUT_TYPES(cls): + paths = [] + for search_path in get_folder_paths("diffusers"): + if os.path.exists(search_path): + for root, subdir, files in os.walk(search_path, followlinks=True): + if "model_index.json" in files: + paths.append(os.path.relpath(root, start=search_path)) + + return {"required": {"model_path": (paths,)}} + + RETURN_TYPES = ("MODEL", "CLIP", "VAE") + FUNCTION = "load_checkpoint" + CATEGORY = "advanced/loaders/diffusers" + + def load_checkpoint(self, model_path, output_vae=True, output_clip=True): + for search_path in get_folder_paths("diffusers"): + if os.path.exists(search_path): + path = os.path.join(search_path, model_path) + if os.path.exists(path): + model_path = path + break + + return load_diffusers_model(model_path, output_vae=output_vae, output_clip=output_clip) + + +class SDXLDiffusersLoader: + @classmethod + def INPUT_TYPES(cls): + paths = [] + for search_path in get_folder_paths("diffusers"): + if os.path.exists(search_path): + for root, subdir, files in os.walk(search_path, followlinks=True): + if "model_index.json" in files: + paths.append(os.path.relpath(root, start=search_path)) + + return {"required": {"model_path": (paths,)}} + + RETURN_TYPES = ("MODEL", "CLIP", "VAE", "TEXT_ENCODER_2") + FUNCTION = "load_checkpoint" + CATEGORY = "advanced/loaders/sdxl" + + def load_checkpoint(self, model_path, output_vae=True, output_clip=True): + for search_path in get_folder_paths("diffusers"): + if os.path.exists(search_path): + path = os.path.join(search_path, model_path) + if os.path.exists(path): + model_path = path + break + + return load_diffusers_sdxl_model(model_path, output_vae=output_vae, output_clip=output_clip) diff --git a/test.py b/test.py new file mode 100644 index 0000000..f2a1c2b --- /dev/null +++ b/test.py @@ -0,0 +1,34 @@ +import os +import torch + +def check_file_permissions(file_path): + try: + if os.access(file_path, os.R_OK | os.W_OK): + print(f"Read and write permissions are available for {file_path}") + else: + raise PermissionError(f"Permission issue with {file_path}") + except Exception as e: + print(f"Error checking permissions for {file_path}: {e}") + raise + +def test_loading_file(file_path): + try: + check_file_permissions(file_path) + print(f"Attempting to load file: {file_path}") + data = torch.load(file_path) + print(f"Successfully loaded file: {file_path}") + except PermissionError as e: + print(f"Permission denied when accessing {file_path}. Please check the file permissions.") + except Exception as e: + print(f"Error loading file {file_path}: {e}") + +# Paths to test +test_paths = [ + "C:\\Users\\Admin\\Desktop\\ComfyUI_windows_portable\\ComfyUI\\models\\diffusers\\raemumix_v90\\model_index.json", + "C:\\Users\\Admin\\Desktop\\ComfyUI_windows_portable\\ComfyUI\\models\\diffusers\\raemumix_v90\\vae\\diffusion_pytorch_model.safetensors", + "C:\\Users\\Admin\\Desktop\\ComfyUI_windows_portable\\ComfyUI\\models\\diffusers\\raemumix_v90\\unet\\diffusion_pytorch_model.safetensors", + "C:\\Users\\Admin\\Desktop\\ComfyUI_windows_portable\\ComfyUI\\models\\diffusers\\raemumix_v90\\text_encoder\\model.safetensors", +] + +for path in test_paths: + test_loading_file(path) diff --git a/test2.py b/test2.py new file mode 100644 index 0000000..2986b47 --- /dev/null +++ b/test2.py @@ -0,0 +1,39 @@ +import os +import utils + +# Define the base model path +base_model_path = "C:\\Users\\Admin\\Desktop\\ComfyUI_windows_portable\\ComfyUI\\models\\diffusers\\raemumix_v90" + +# Test loading the main model +try: + model = utils.load_model_from_path(os.path.join(base_model_path, "model_index.json")) + print("Main model loaded successfully.") +except Exception as e: + print(f"Error loading main model: {e}") + +# Test loading the VAE model +try: + vae = utils.load_vae_from_path(base_model_path) + print("VAE model loaded successfully.") +except Exception as e: + print(f"Error loading VAE model: {e}") + +# Test loading the CLIP model +try: + clip = utils.load_text_encoder_from_path(base_model_path, encoder_id=None) + print("CLIP model loaded successfully.") +except Exception as e: + print(f"Error loading CLIP model: {e}") + +# Test loading the text encoders for SDXL +try: + text_encoder_1 = utils.load_text_encoder_from_path(base_model_path, 1) + print("Text encoder 1 loaded successfully.") +except Exception as e: + print(f"Error loading text encoder 1: {e}") + +try: + text_encoder_2 = utils.load_text_encoder_from_path(base_model_path, 2) + print("Text encoder 2 loaded successfully.") +except Exception as e: + print(f"Error loading text encoder 2: {e}") diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..d3879bd --- /dev/null +++ b/utils.py @@ -0,0 +1,121 @@ +import os +import json +import torch +from safetensors import safe_open + +def load_diffusers_model(model_path, output_vae=True, output_clip=True): + try: + model = load_model_from_path(model_path) + vae = load_vae_from_path(model_path) if output_vae else None + clip = load_text_encoder_from_path(model_path, encoder_id=None) if output_clip else None + return model, vae, clip + except Exception as e: + print(f"Error loading diffusers model from {model_path}: {e}") + raise + +def load_diffusers_sdxl_model(model_path, output_vae=True, output_clip=True, embedding_directory=None): + try: + model = load_model_from_path(model_path) + vae = load_vae_from_path(model_path) if output_vae else None + clip = load_text_encoder_from_path(model_path, encoder_id=1) if output_clip else None + text_encoder_2 = load_text_encoder_from_path(model_path, encoder_id=2) + + if embedding_directory: + # Load or handle embeddings if applicable + pass + + return model, vae, clip, text_encoder_2 + except Exception as e: + print(f"Error loading SDXL diffusers model from {model_path}: {e}") + raise + +def load_model_from_path(model_path): + try: + print(f"Loading model from {model_path}") + check_file_permissions(model_path) + if model_path.endswith('.safetensors'): + model = load_safetensors(model_path) + else: + model = load_json(model_path) + return model + except PermissionError as e: + print(f"Permission denied when accessing {model_path}. Please check the file permissions.") + raise + except Exception as e: + print(f"Error loading model from {model_path}: {e}") + raise + +def load_vae_from_path(model_path): + try: + print(f"Loading VAE from {model_path}") + vae_path = os.path.join(model_path, "vae", "diffusion_pytorch_model.safetensors") + check_file_permissions(vae_path) + vae = load_safetensors(vae_path) + return vae + except PermissionError as e: + print(f"Permission denied when accessing {vae_path}. Please check the file permissions.") + raise + except Exception as e: + print(f"Error loading VAE from {vae_path}: {e}") + raise + +def load_clip_from_path(model_path): + try: + print(f"Loading CLIP model from {model_path}") + clip_path = os.path.join(model_path, "clip", "diffusion_pytorch_model.safetensors") + check_file_permissions(clip_path) + clip = load_safetensors(clip_path) + return clip + except PermissionError as e: + print(f"Permission denied when accessing {clip_path}. Please check the file permissions.") + raise + except Exception as e: + print(f"Error loading CLIP model from {clip_path}: {e}") + raise + +def load_text_encoder_from_path(model_path, encoder_id=None): + try: + if encoder_id is not None: + encoder_path = os.path.join(model_path, f"text_encoder_{encoder_id}", "model.safetensors") + print(f"Loading text encoder {encoder_id} from {encoder_path}") + else: + encoder_path = os.path.join(model_path, "text_encoder", "model.safetensors") + print(f"Loading text encoder from {encoder_path}") + + check_file_permissions(encoder_path) + text_encoder = load_safetensors(encoder_path) + return text_encoder + except PermissionError as e: + print(f"Permission denied when accessing {encoder_path}. Please check the file permissions.") + raise + except Exception as e: + print(f"Error loading text encoder from {encoder_path}: {e}") + raise + +def load_json(json_path): + try: + print(f"Loading JSON from {json_path}") + with open(json_path, 'r') as f: + return json.load(f) + except Exception as e: + print(f"Error loading JSON from {json_path}: {e}") + raise + +def load_safetensors(file_path): + try: + print(f"Loading safetensors from {file_path}") + with safe_open(file_path, framework="pt", device="cpu") as f: + return {k: f.get_tensor(k) for k in f.keys()} + except Exception as e: + print(f"Error loading safetensors from {file_path}: {e}") + raise + +def check_file_permissions(file_path): + try: + if os.access(file_path, os.R_OK | os.W_OK): + print(f"Read and write permissions are available for {file_path}") + else: + raise PermissionError(f"Permission issue with {file_path}") + except Exception as e: + print(f"Error checking permissions for {file_path}: {e}") + raise