custom_nodes/ComfyUI-ppm/src/compat/utils.py

from dataclasses import dataclass
from types import ModuleType

import torch

from comfy_api.latest import io


def from_zero(weights, base_emb):
    weight_tensor = torch.tensor(weights, dtype=base_emb.dtype, device=base_emb.device)
    weight_tensor = weight_tensor.reshape(1, -1, 1).expand(base_emb.shape)
    return base_emb * weight_tensor


def v3_schema_stub(module: ModuleType) -> list[type[io.ComfyNode]]:
    NODE_CLASS_MAPPINGS: dict[str, type[io.ComfyNode]] = module.NODE_CLASS_MAPPINGS
    NODE_DISPLAY_NAME_MAPPINGS: dict[str, str] = (
        module.NODE_DISPLAY_NAME_MAPPINGS if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") else {}
    )

    @dataclass
    class SchemaPPMStub:
        node_id: str
        display_name: str | None

    def inject_schema_stub(cls: type[io.ComfyNode], node_id: str, display_name: str | None = None):
        schema = SchemaPPMStub(node_id, display_name)

        if not hasattr(cls, "GET_SCHEMA"):
            setattr(cls, "GET_SCHEMA", lambda: schema)

        return cls

    return [inject_schema_stub(m[1], m[0], NODE_DISPLAY_NAME_MAPPINGS.get(m[0])) for m in NODE_CLASS_MAPPINGS.items()]