ComfyUI/app/subgraph_manager.py

from __future__ import annotations

from typing import TypedDict
import json
import logging
import os
import folder_paths
import glob
from aiohttp import web
import hashlib


class Source:
    custom_node = "custom_node"
    templates = "templates"

class SubgraphEntry(TypedDict):
    source: str
    """
    Source of subgraph - custom_nodes vs templates.
    """
    path: str
    """
    Relative path of the subgraph file.
    For custom nodes, will be the relative directory like <custom_node_dir>/subgraphs/<name>.json
    """
    name: str
    """
    Name of subgraph file.
    """
    info: CustomNodeSubgraphEntryInfo
    """
    Additional info about subgraph; in the case of custom_nodes, will contain nodepack name
    """
    data: str

class CustomNodeSubgraphEntryInfo(TypedDict):
    node_pack: str
    """Node pack name."""

class SubgraphManager:
    def __init__(self):
        self.cached_custom_node_subgraphs: dict[SubgraphEntry] | None = None
        self.cached_blueprint_subgraphs: dict[SubgraphEntry] | None = None
        self.distribution = os.environ.get("DISTRIBUTION", "localhost")

    def _create_entry(self, file: str, source: str, node_pack: str) -> tuple[str, SubgraphEntry]:
        """Create a subgraph entry from a file path. Expects normalized path (forward slashes)."""
        entry_id = hashlib.sha256(f"{source}{file}".encode()).hexdigest()
        entry: SubgraphEntry = {
            "source": source,
            "name": os.path.splitext(os.path.basename(file))[0],
            "path": file,
            "info": {"node_pack": node_pack},
        }
        return entry_id, entry

    async def load_entry_data(self, entry: SubgraphEntry):
        with open(entry['path'], 'r') as f:
            entry['data'] = f.read()
        return entry

    async def sanitize_entry(self, entry: SubgraphEntry | None, remove_data=False) -> SubgraphEntry | None:
        if entry is None:
            return None
        entry = entry.copy()
        entry.pop('path', None)
        if remove_data:
            entry.pop('data', None)
        return entry

    async def sanitize_entries(self, entries: dict[str, SubgraphEntry], remove_data=False) -> dict[str, SubgraphEntry]:
        entries = entries.copy()
        for key in list(entries.keys()):
            entries[key] = await self.sanitize_entry(entries[key], remove_data)
        return entries

    async def get_custom_node_subgraphs(self, loadedModules, force_reload=False):
        """Load subgraphs from custom nodes."""
        if not force_reload and self.cached_custom_node_subgraphs is not None:
            return self.cached_custom_node_subgraphs

        subgraphs_dict: dict[SubgraphEntry] = {}
        for folder in folder_paths.get_folder_paths("custom_nodes"):
            pattern = os.path.join(folder, "*/subgraphs/*.json")
            for file in glob.glob(pattern):
                file = file.replace('\\', '/')
                node_pack = "custom_nodes." + file.split('/')[-3]
                entry_id, entry = self._create_entry(file, Source.custom_node, node_pack)
                subgraphs_dict[entry_id] = entry

        self.cached_custom_node_subgraphs = subgraphs_dict
        return subgraphs_dict

    async def get_blueprint_subgraphs(self, force_reload=False):
        """Load subgraphs from the blueprints directory using index.json for discovery."""
        if not force_reload and self.cached_blueprint_subgraphs is not None:
            return self.cached_blueprint_subgraphs

        subgraphs_dict: dict[SubgraphEntry] = {}
        blueprints_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'blueprints')

        index_path = os.path.join(blueprints_dir, "index.json")
        if os.path.isfile(index_path):
            try:
                with open(index_path, "r", encoding="utf-8") as f:
                    categories = json.load(f)
            except (json.JSONDecodeError, OSError) as e:
                logging.error("Failed to load blueprint index %s: %s", index_path, e)
                categories = []

            if not isinstance(categories, list):
                logging.error("Blueprint index.json is not a list: %s", index_path)
                categories = []

            for category in categories:
                module_name = category.get("moduleName", "default")
                for blueprint in category.get("blueprints", []):
                    name = blueprint.get("name")
                    if not name:
                        logging.warning("Blueprint entry missing 'name' in category '%s', skipping", module_name)
                        continue

                    include_on = blueprint.get("includeOnDistributions")
                    if include_on is not None and self.distribution not in include_on:
                        continue

                    file_by_dist = blueprint.get("fileByDistribution", {})
                    filename = file_by_dist.get(self.distribution, f"{name}.json")
                    filepath = os.path.realpath(os.path.join(blueprints_dir, filename))
                    if not filepath.startswith(os.path.realpath(blueprints_dir) + os.sep):
                        logging.warning("Blueprint path escapes blueprints directory: %s", filepath)
                        continue

                    if not os.path.isfile(filepath):
                        logging.warning("Blueprint file not found: %s", filepath)
                        continue

                    entry_id, entry = self._create_entry(filepath, Source.templates, module_name)
                    subgraphs_dict[entry_id] = entry
        elif os.path.exists(blueprints_dir):
            logging.warning("No blueprint index.json found at %s, falling back to glob", index_path)
            for file in glob.glob(os.path.join(blueprints_dir, "*.json")):
                if os.path.basename(file) == "index.json":
                    continue
                file = file.replace('\\', '/')
                entry_id, entry = self._create_entry(file, Source.templates, "comfyui")
                subgraphs_dict[entry_id] = entry

        self.cached_blueprint_subgraphs = subgraphs_dict
        return subgraphs_dict

    async def get_all_subgraphs(self, loadedModules, force_reload=False):
        """Get all subgraphs from all sources (custom nodes and blueprints)."""
        custom_node_subgraphs = await self.get_custom_node_subgraphs(loadedModules, force_reload)
        blueprint_subgraphs = await self.get_blueprint_subgraphs(force_reload)
        return {**custom_node_subgraphs, **blueprint_subgraphs}

    async def get_subgraph(self, id: str, loadedModules):
        """Get a specific subgraph by ID from any source."""
        entry = (await self.get_all_subgraphs(loadedModules)).get(id)
        if entry is not None and entry.get('data') is None:
            await self.load_entry_data(entry)
        return entry

    def add_routes(self, routes, loadedModules):
        @routes.get("/global_subgraphs")
        async def get_global_subgraphs(request):
            subgraphs_dict = await self.get_all_subgraphs(loadedModules)
            return web.json_response(await self.sanitize_entries(subgraphs_dict, remove_data=True))

        @routes.get("/global_subgraphs/{id}")
        async def get_global_subgraph(request):
            id = request.match_info.get("id", None)
            subgraph = await self.get_subgraph(id, loadedModules)
            return web.json_response(await self.sanitize_entry(subgraph))