benfinklea-ComfyUI-ModelSync/nodes.py

"""
Optional ComfyUI nodes for ModelSync.

These nodes are not required for ModelSync functionality but provide
convenient workflow integration for specific use cases.
"""

from typing import Tuple, Any


class ModelInfoNode:
    """
    Node that outputs model metadata from the ModelSync database.

    Useful for debugging or displaying model information in workflows.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_path": ("STRING", {"default": "", "multiline": False}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING", "INT", "BOOLEAN")
    RETURN_NAMES = ("source_url", "source_type", "vram_mb", "update_available")
    FUNCTION = "get_info"
    CATEGORY = "ModelSync"

    def get_info(self, model_path: str) -> Tuple[str, str, int, bool]:
        """Get model info from the database."""
        try:
            from .server import get_api

            api = get_api()
            if api is None:
                return ("", "unknown", 0, False)

            model = api.store.get_model_by_path(model_path)
            if model is None:
                return ("", "unknown", 0, False)

            return (
                model.get("source_url") or "",
                model.get("source_type") or "unknown",
                model.get("vram_mb") or 0,
                bool(model.get("update_available")),
            )

        except Exception as e:
            print(f"[ModelSync] Error getting model info: {e}")
            return ("", "unknown", 0, False)


class CheckModelUpdateNode:
    """
    Node that checks if a model has an update available.

    Can be used in workflows to conditionally notify about updates.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_path": ("STRING", {"default": "", "multiline": False}),
                "check_remote": ("BOOLEAN", {"default": False}),
            },
        }

    RETURN_TYPES = ("BOOLEAN", "STRING")
    RETURN_NAMES = ("update_available", "message")
    FUNCTION = "check_update"
    CATEGORY = "ModelSync"

    def check_update(self, model_path: str, check_remote: bool) -> Tuple[bool, str]:
        """Check if model has an update available."""
        try:
            from .server import get_api
            import asyncio

            api = get_api()
            if api is None:
                return (False, "ModelSync API not initialized")

            model = api.store.get_model_by_path(model_path)
            if model is None:
                return (False, f"Model not found: {model_path}")

            if check_remote and model.get("source_url"):
                # Check remote for updates
                loop = asyncio.new_event_loop()
                try:
                    remote_info = loop.run_until_complete(
                        api.tracker.check_version(model["source_url"])
                    )
                finally:
                    loop.close()

                if remote_info:
                    is_update = api.tracker.is_update_available(
                        local_hash=model.get("file_hash"),
                        local_size=model.get("file_size"),
                        remote_hash=remote_info.get("hash"),
                        remote_size=remote_info.get("size"),
                    )
                    if is_update:
                        return (True, f"Update available for {model['filename']}")
                    return (False, f"{model['filename']} is up to date")

            # Use cached status
            if model.get("update_available"):
                return (True, f"Update available for {model['filename']}")
            return (False, f"{model['filename']} is up to date")

        except Exception as e:
            return (False, f"Error: {e}")


class SetModelSourceNode:
    """
    Node to set the source URL for a model.

    Useful for associating downloaded models with their source.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_path": ("STRING", {"default": "", "multiline": False}),
                "source_url": ("STRING", {"default": "", "multiline": False}),
            },
            "optional": {
                "auto_update": ("BOOLEAN", {"default": False}),
            },
        }

    RETURN_TYPES = ("BOOLEAN",)
    RETURN_NAMES = ("success",)
    FUNCTION = "set_source"
    CATEGORY = "ModelSync"
    OUTPUT_NODE = True

    def set_source(
        self, model_path: str, source_url: str, auto_update: bool = False
    ) -> Tuple[bool]:
        """Set the source URL for a model."""
        try:
            from .server import get_api

            api = get_api()
            if api is None:
                print("[ModelSync] API not initialized")
                return (False,)

            model = api.store.get_model_by_path(model_path)
            if model is None:
                print(f"[ModelSync] Model not found: {model_path}")
                return (False,)

            # Detect source type
            source_type = api.tracker.detect_source_type(source_url)

            api.store.update_model(model["id"], {
                "source_url": source_url,
                "source_type": source_type,
                "auto_update": 1 if auto_update else 0,
            })

            print(f"[ModelSync] Set source for {model['filename']}: {source_url}")
            return (True,)

        except Exception as e:
            print(f"[ModelSync] Error setting source: {e}")
            return (False,)


# Node registration
NODE_CLASS_MAPPINGS = {
    "ModelSyncInfo": ModelInfoNode,
    "ModelSyncCheckUpdate": CheckModelUpdateNode,
    "ModelSyncSetSource": SetModelSourceNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ModelSyncInfo": "Model Info (ModelSync)",
    "ModelSyncCheckUpdate": "Check Model Update (ModelSync)",
    "ModelSyncSetSource": "Set Model Source (ModelSync)",
}