willmiao-ComfyUI-Lora-Manager/py/services/workflow_scanner.py

"""Scanner for ComfyUI workflow resources (.json documents and containers)."""

from __future__ import annotations

import logging
import os
from datetime import datetime
from typing import List, Optional

from ..config import config
from ..utils.file_utils import (
    calculate_sha256,
    find_preview_file,
    normalize_path,
)
from ..utils.metadata_manager import MetadataManager
from ..utils.models import WorkflowMetadata
from ..utils.sidecar_paths import get_preview_dir
from ..utils.workflow_analysis import analyze_workflow_file
from ..utils.workflow_extraction import (
    IMAGE_CONTAINER_EXTENSIONS,
    WORKFLOW_ARTIFACT_EXTENSIONS,
    WORKFLOW_DOCUMENT_SUFFIX,
    materialize_workflow_document,
)
from .model_hash_index import ModelHashIndex
from .model_scanner import ModelScanner

logger = logging.getLogger(__name__)

# Sidecar files share the .json extension; they must never enter the library.
_SIDECAR_SUFFIXES = (".metadata.json", WORKFLOW_DOCUMENT_SUFFIX)


class WorkflowScanner(ModelScanner):
    """Scan and manage ComfyUI workflow artifacts.

    Library entries are the original uploaded files, kept byte-identical so
    their SHA256 stays the CivitAI identity (by-hash lookup and version
    tracking depend on it). ``.json`` artifacts carry the graph directly;
    ``.zip``/``.png``/``.webp`` containers get the workflow document
    extracted into a derived ``.workflow.json`` sidecar that every consumer
    (graph analysis, send-to-ComfyUI) reads. Files from which no workflow
    document can be extracted never enter the library — this keeps settings
    JSON and plain preview images out.
    """

    def __init__(self):
        super().__init__(
            model_type="workflow",
            model_class=WorkflowMetadata,
            file_extensions=set(WORKFLOW_ARTIFACT_EXTENSIONS),
            hash_index=ModelHashIndex(),
        )

    def _is_indexable_filename(self, filename: str) -> bool:
        lowered = filename.lower()
        return not lowered.endswith(_SIDECAR_SUFFIXES)

    def get_model_roots(self) -> List[str]:
        """Workflow roots (primary + extra), empty while the feature is off."""
        try:
            from .settings_manager import get_settings_manager

            if not get_settings_manager().is_workflows_enabled():
                return []
        except Exception:  # pragma: no cover - defensive; settings unavailable
            return []
        roots: List[str] = []
        roots.extend(config.workflows_roots or [])
        roots.extend(config.extra_workflows_roots or [])
        seen: set[str] = set()
        unique_roots: List[str] = []
        for root in roots:
            if root and root not in seen:
                seen.add(root)
                unique_roots.append(root)
        return unique_roots

    @staticmethod
    def _is_image_artifact(file_path: str) -> bool:
        return os.path.splitext(file_path)[1].lower() in IMAGE_CONTAINER_EXTENSIONS

    def adjust_metadata(self, metadata, file_path, root_path):
        """Fill workflow analysis fields when they are missing.

        Runs on every scan and on the post-download path. The extraction
        result is cached in the derived ``.workflow.json`` sidecar, so the
        steady-state cost is one small JSON parse per entry whose analysis
        fields have not been populated yet.
        """
        if not isinstance(metadata, WorkflowMetadata) or metadata.workflow_format:
            return metadata
        try:
            document_path = materialize_workflow_document(file_path)
            analysis = (
                analyze_workflow_file(document_path) if document_path else None
            )
        except Exception as exc:
            logger.debug(
                "Workflow analysis refresh failed for %s: %s", file_path, exc
            )
            return metadata
        if analysis is None:
            return metadata
        metadata.workflow_format = analysis.workflow_format
        metadata.node_count = analysis.node_count
        metadata.node_classes = analysis.node_classes
        metadata.model_refs = analysis.model_refs
        metadata.civitai_refs = analysis.civitai_refs
        return metadata

    async def _create_default_metadata(
        self, file_path: str
    ) -> Optional[WorkflowMetadata]:
        """Create metadata by hashing the artifact and analyzing its workflow.

        Returns None for files that carry no ComfyUI workflow document, so
        they stay out of the library.
        """
        try:
            real_path = os.path.realpath(file_path)
            if not os.path.exists(real_path):
                logger.error(f"File not found: {file_path}")
                return None

            document_path = materialize_workflow_document(
                file_path, read_path=real_path
            )
            analysis = (
                analyze_workflow_file(document_path) if document_path else None
            )
            if analysis is None:
                logger.debug(f"Skipping artifact without a workflow: {file_path}")
                return None

            base_name = os.path.splitext(os.path.basename(file_path))[0]
            preview_url = find_preview_file(base_name, get_preview_dir(file_path))
            if not preview_url and self._is_image_artifact(file_path):
                # An image artifact doubles as its own preview.
                preview_url = file_path
            sha256 = await calculate_sha256(real_path)
            if not sha256:
                logger.warning(
                    f"Skipping metadata creation for empty workflow file: {file_path}"
                )
                return None

            metadata = WorkflowMetadata(
                file_name=base_name,
                model_name=base_name,
                file_path=normalize_path(file_path),
                size=os.path.getsize(real_path),
                modified=datetime.now().timestamp(),
                sha256=sha256,
                base_model="Unknown",
                preview_url=normalize_path(preview_url),
                tags=[],
                modelDescription="",
                from_civitai=False,
                workflow_format=analysis.workflow_format,
                node_count=analysis.node_count,
                node_classes=analysis.node_classes,
                model_refs=analysis.model_refs,
                civitai_refs=analysis.civitai_refs,
            )

            logger.info(f"Creating workflow metadata for {file_path}")
            await MetadataManager.save_metadata(file_path, metadata)

            return metadata

        except Exception as e:
            logger.error(
                f"Error creating default workflow metadata for {file_path}: {e}"
            )
            return None