core_anim/wanvideo_wrapper_bridge.py

from __future__ import annotations

import importlib.util
import sys
import types
from pathlib import Path


_WRAPPER_ALIAS = "ComfyUI_WanVideoWrapper"
_WRAPPER_MODULE = None


def _repo_root() -> Path:
    return Path(__file__).resolve().parent.parent


def _wrapper_dir() -> Path:
    return _repo_root() / "core_anim" / "third_party" / "ComfyUI-WanVideoWrapper"


def _comfy_runtime_dir() -> Path:
    return _repo_root() / "core_anim" / "third_party" / "comfy_runtime"


def _install_server_stub() -> None:
    if "server" in sys.modules:
        return

    server = types.ModuleType("server")

    class _BinaryEventTypes:
        PREVIEW_IMAGE = 1

    class _DummyPromptServer:
        instance = None

        def __init__(self) -> None:
            self.last_node_id = "loop_runner"
            self.client_id = None

        def send_sync(self, *args, **kwargs):
            return None

        def send_progress_text(self, *args, **kwargs):
            return None

    _DummyPromptServer.instance = _DummyPromptServer()
    server.PromptServer = _DummyPromptServer
    server.BinaryEventTypes = _BinaryEventTypes
    sys.modules["server"] = server


def load_wrapper_module():
    global _WRAPPER_MODULE
    if _WRAPPER_MODULE is not None:
        return _WRAPPER_MODULE

    wrapper_dir = _wrapper_dir()
    comfy_runtime_dir = _comfy_runtime_dir()
    if not wrapper_dir.exists():
        raise FileNotFoundError(f"WanVideoWrapper non trovato: {wrapper_dir}")
    if not comfy_runtime_dir.exists():
        raise FileNotFoundError(f"comfy_runtime non trovato: {comfy_runtime_dir}")

    if str(comfy_runtime_dir) not in sys.path:
        sys.path.insert(0, str(comfy_runtime_dir))

    _install_server_stub()

    spec = importlib.util.spec_from_file_location(
        _WRAPPER_ALIAS,
        wrapper_dir / "__init__.py",
        submodule_search_locations=[str(wrapper_dir)],
    )
    if spec is None or spec.loader is None:
        raise ImportError(f"Impossibile creare spec per {_WRAPPER_ALIAS}")

    module = importlib.util.module_from_spec(spec)
    sys.modules[_WRAPPER_ALIAS] = module
    spec.loader.exec_module(module)
    _WRAPPER_MODULE = module
    return module


def get_wrapper_node_class(name: str):
    module = load_wrapper_module()
    try:
        return module.NODE_CLASS_MAPPINGS[name]
    except KeyError as exc:
        raise KeyError(f"Nodo WanVideoWrapper non trovato: {name}") from exc