comfygit-core 0.2.0__py3-none-any.whl

comfygit_core/configs/comfyui_models.py

@@ -0,0 +1,62 @@
+"""Default ComfyUI models configuration"""
+
+COMFYUI_MODELS_CONFIG = {
+    "version": "2024.1",
+    "default_extensions": [
+        ".ckpt",
+        ".pt",
+        ".pth",
+        ".pt2",
+        ".bin",
+        ".safetensors",
+        ".pkl",
+        ".sft",
+    ],
+    "standard_directories": [
+        "audio_encoders",
+        "checkpoints",
+        "clip",
+        "clip_vision",
+        "configs",
+        "controlnet",
+        "diffusers",
+        "diffusion_models",
+        "embeddings",
+        "gligen",
+        "hypernetworks",
+        "loras",
+        "model_patches",
+        "photomaker",
+        "style_models",
+        "text_encoders",
+        "unet",
+        "upscale_models",
+        "vae",
+        "vae_approx",
+    ],
+    "directory_overrides": {"configs": {"extensions": [".yaml", ".yml", ".json"]}},
+    "node_directory_mappings": {
+        "CheckpointLoaderSimple": ["checkpoints"],
+        "CheckpointLoader": ["checkpoints", "configs"],
+        "unCLIPCheckpointLoader": ["checkpoints"],
+        "ImageOnlyCheckpointLoader": ["checkpoints"],
+        "VAELoader": ["vae", "vae_approx"],
+        "LoraLoader": ["loras"],
+        "LoraLoaderModelOnly": ["loras"],
+        "CLIPLoader": ["clip"],
+        "DualCLIPLoader": ["clip"],
+        "TripleCLIPLoader": ["clip"],
+        "QuadrupleCLIPLoader": ["clip"],
+        "UNETLoader": ["diffusion_models"],
+        "CLIPVisionLoader": ["clip_vision"],
+        "ControlNetLoader": ["controlnet"],
+        "DiffControlNetLoader": ["controlnet"],
+        "StyleModelLoader": ["style_models"],
+        "UpscaleModelLoader": ["upscale_models"],
+        "GLIGENLoader": ["gligen"],
+        "HypernetworkLoader": ["hypernetworks"],
+        "PhotoMakerLoader": ["photomaker"],
+        "DiffusersLoader": ["diffusers"],
+    },
+    "node_widget_indices": {"CheckpointLoader": 1},
+}

comfygit_core/configs/model_config.py

@@ -0,0 +1,151 @@
+"""Model configuration loader and utilities."""
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from .comfyui_models import COMFYUI_MODELS_CONFIG
+from ..logging.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class ModelConfig:
+    """ComfyUI model configuration."""
+    version: str
+    default_extensions: list[str]
+    standard_directories: list[str]
+    directory_overrides: dict[str, dict[str, Any]]
+    node_directory_mappings: dict[str, list[str]]
+    node_widget_indices: dict[str, int]
+
+    @classmethod
+    def load(cls, config_path: Path | None = None) -> "ModelConfig":
+        """Load model configuration from file.
+
+        Args:
+            config_path: Path to config file, or None to use default
+
+        Returns:
+            ModelConfig instance
+        """
+        data = {}
+        if config_path is None:
+            # Load hardcoded config (fallback)
+            data = COMFYUI_MODELS_CONFIG
+        else:
+            if not config_path.exists():
+                raise FileNotFoundError(f"Model config file not found: {config_path}")
+            try:
+                with open(config_path, 'r', encoding='utf-8') as f:
+                    data = json.load(f)
+            except Exception as e:
+                logger.error(f"Failed to load model config from {config_path}: {e}")
+                raise
+
+        return cls(
+            version=data.get("version", "unknown"),
+            default_extensions=data.get("default_extensions", []),
+            standard_directories=data.get("standard_directories", []),
+            directory_overrides=data.get("directory_overrides", {}),
+            node_directory_mappings=data.get("node_directory_mappings", {}),
+            node_widget_indices=data.get("node_widget_indices", {})
+        )
+
+    def get_extensions_for_directory(self, directory: str) -> list[str]:
+        """Get file extensions for a specific directory.
+
+        Args:
+            directory: Directory name (e.g., "checkpoints")
+
+        Returns:
+            List of supported extensions for this directory
+        """
+        if directory in self.directory_overrides:
+            override = self.directory_overrides[directory]
+            if "extensions" in override:
+                return override["extensions"]
+
+        return self.default_extensions
+
+    def is_standard_directory(self, directory: str) -> bool:
+        """Check if a directory is a standard ComfyUI directory.
+
+        Args:
+            directory: Directory name to check
+
+        Returns:
+            True if it's a standard directory, False for custom
+        """
+        return directory in self.standard_directories
+
+    def is_model_file(self, file_path: Path) -> bool:
+        """Check if a file is a potential model file based on extension.
+
+        Args:
+            file_path: Path to file to check
+
+        Returns:
+            True if file has a model extension
+        """
+        extension = file_path.suffix.lower()
+
+        # Get directory from path to check for specific extensions
+        parts = file_path.parts
+        for part in parts:
+            if self.is_standard_directory(part):
+                return extension in self.get_extensions_for_directory(part)
+
+        # Default to checking against default extensions
+        return extension in self.default_extensions
+
+    def get_directories_for_node(self, node_type: str) -> list[str]:
+        """Get model directories for a node type.
+
+        Args:
+            node_type: ComfyUI node type (e.g., "LoraLoader")
+
+        Returns:
+            List of directories this node type loads from
+        """
+        return self.node_directory_mappings.get(node_type, [])
+
+    def get_widget_index_for_node(self, node_type: str) -> int:
+        """Get widget index containing model path for a node type.
+
+        Args:
+            node_type: ComfyUI node type
+
+        Returns:
+            Index in widgets_values containing the model path (defaults to 0)
+        """
+        return self.node_widget_indices.get(node_type, 0)
+
+    def is_model_loader_node(self, node_type: str) -> bool:
+        """Check if a node type is a known model loader.
+
+        Args:
+            node_type: ComfyUI node type
+
+        Returns:
+            True if node type loads models
+        """
+        return node_type in self.node_directory_mappings
+
+    def reconstruct_model_path(self, node_type: str, widget_value: str) -> list[str]:
+        """Reconstruct full model paths from node type and widget value.
+
+        Args:
+            node_type: ComfyUI node type
+            widget_value: Value from node's widgets_values
+
+        Returns:
+            List of possible full relative paths
+        """
+        directories = self.get_directories_for_node(node_type)
+        if not directories:
+            return []
+
+        return [f"{directory}/{widget_value}" for directory in directories]

comfygit_core/constants.py

@@ -0,0 +1,82 @@
+"""Constants and configuration for ComfyUI environment detection."""
+
+# PyTorch-related packages that need special handling
+PYTORCH_PACKAGE_NAMES = {
+    # Core PyTorch packages
+    'torch', 'torchvision', 'torchaudio',
+
+    # Triton packages (PyTorch's fused-kernel compiler)
+    'triton', 'triton-windows',
+
+    # NVIDIA CUDA packages
+    'nvidia-cublas-cu11', 'nvidia-cublas-cu12',
+    'nvidia-cuda-runtime-cu11', 'nvidia-cuda-runtime-cu12',
+    'nvidia-cuda-nvrtc-cu11', 'nvidia-cuda-nvrtc-cu12',
+    'nvidia-cudnn-cu11', 'nvidia-cudnn-cu12',
+    'nvidia-cufft-cu11', 'nvidia-cufft-cu12',
+    'nvidia-curand-cu11', 'nvidia-curand-cu12',
+    'nvidia-cusolver-cu11', 'nvidia-cusolver-cu12',
+    'nvidia-cusparse-cu11', 'nvidia-cusparse-cu12',
+    'nvidia-nccl-cu11', 'nvidia-nccl-cu12',
+    'nvidia-nvtx-cu11', 'nvidia-nvtx-cu12',
+
+    # New CUDA packages in PyTorch 2.6+
+    'nvidia-cuda-cupti-cu11', 'nvidia-cuda-cupti-cu12',  # CUPTI (profiler)
+    'nvidia-cufile-cu11', 'nvidia-cufile-cu12',  # cuFile (GPUDirect Storage)
+    'nvidia-cusparselt-cu11', 'nvidia-cusparselt-cu12',  # structured-sparse LT
+    'nvidia-nvjitlink-cu11', 'nvidia-nvjitlink-cu12',  # NVJitLink
+
+    # NOTE: nvidia-ml-py and nvidia-ml-py3 are NOT included
+    # These are optional NVML bindings for monitoring that PyTorch doesn't depend on
+}
+
+# Blacklist of directory names that should not be treated as custom nodes
+CUSTOM_NODES_BLACKLIST = {
+    '__pycache__',
+    '.pytest_cache',
+    '.mypy_cache',
+    '.tox',
+    'node_modules',
+    '.vscode',
+    '.idea',
+    'venv',
+    '.venv',
+    'env',
+    '.env',
+    'build',
+    'dist',
+    '.DS_Store',
+    'Thumbs.db',
+    '.coverage',
+    'htmlcov',
+    '.git',
+    '.svn',
+    '.hg',
+    'test',
+    'tests',
+    '__tests__',
+    'spec',
+    'specs',
+    'tmp',
+    'temp',
+    '.tmp',
+    '.temp'
+}
+
+# Default values
+DEFAULT_REGISTRY_URL = "https://api.comfy.org"
+DEFAULT_GITHUB_URL = "https://github.com"
+GITHUB_API_BASE = "https://api.github.com"
+
+GITHUB_NODE_MAPPINGS_URL = "https://raw.githubusercontent.com/ComfyDock/ComfyDock-Registry-Data/main/data/node_mappings.json"
+
+MAX_REGISTRY_DATA_AGE_HOURS = 24
+
+# Prevent infinite loops for optional group removal
+MAX_OPT_GROUP_RETRIES = 10  
+
+# PyTorch core packages
+PYTORCH_CORE_PACKAGES = ["torch", "torchvision", "torchaudio"]
+
+# PyTorch index base URL
+PYTORCH_INDEX_BASE_URL = "https://download.pytorch.org/whl"