pymss-project-comfy-mss/web/comfy_mss.js

import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";

import { applyLinkColorPatch, applyTypeColors, colorLink, colorNodeSlots, colorSlot } from "./comfy_mss/colors.js";
import {
  AUDIO_ENSEMBLE_NODE_TYPE,
  FIXED_200_NODE_TYPES,
  FIXED_200_NODE_WIDTH,
  FIXED_420_NODE_TYPES,
  FIXED_420_NODE_WIDTH,
  FIXED_260_NODE_TYPES,
  LOAD_AUDIO_NODE_TYPE,
  SAVE_AUDIO_NODE_TYPE,
  SEPARATE_LIST_NODE_TYPES,
  SEPARATE_NODE_TYPES,
} from "./comfy_mss/constants.js";
import { registerAudioEnsembleNode } from "./comfy_mss/ensemble.js";
import { registerLoadAudioNode } from "./comfy_mss/load_audio.js";
import { registerSaveAudioNode } from "./comfy_mss/save_audio.js";
import { registerSeparateNode } from "./comfy_mss/separate.js";
import { registerFixedWidthNode, registerFixedWidthOnlyNode } from "./comfy_mss/sizing.js";
import {
  loadTranslations,
  onTranslationsLoaded,
  t,
  translateNodeLabels,
  withTranslatedWidgetNames,
} from "./comfy_mss/i18n.js";

function applyColorSetup() {
  applyTypeColors(app);
  applyLinkColorPatch(app);
}

const COMFY_MSS_NODE_TYPES = new Set([
  ...SEPARATE_NODE_TYPES,
  ...SEPARATE_LIST_NODE_TYPES,
  ...FIXED_200_NODE_TYPES,
  ...FIXED_260_NODE_TYPES,
  ...FIXED_420_NODE_TYPES,
  LOAD_AUDIO_NODE_TYPE,
  AUDIO_ENSEMBLE_NODE_TYPE,
  SAVE_AUDIO_NODE_TYPE,
]);

function refreshGraphLanguage() {
  for (const node of app.graph?._nodes ?? []) {
    const nodeName = node?.comfyClass ?? node?.type;
    if (!COMFY_MSS_NODE_TYPES.has(nodeName)) {
      continue;
    }
    translateNodeLabels(node, nodeName);
  }
  app.graph?.setDirtyCanvas?.(true, true);
}

app.registerExtension({
  name: "comfy-mss.ui",

  setup() {
    applyColorSetup();
    loadTranslations(api);
    onTranslationsLoaded(refreshGraphLanguage);
  },

  async beforeRegisterNodeDef(nodeType, nodeData) {
    applyColorSetup();
    loadTranslations(api);
    function wrapOnNodeCreated(extra) {
      const previousOnNodeCreated = nodeType.prototype.onNodeCreated;
      nodeType.prototype.onNodeCreated = function () {
        const result = previousOnNodeCreated?.apply(this, arguments);
        colorNodeSlots(this);
        extra?.call(this);
        translateNodeLabels(this, nodeData.name);
        return result;
      };
    }

    if (!COMFY_MSS_NODE_TYPES.has(nodeData.name)) {
      return;
    }

    const originalOnConnectionsChange = nodeType.prototype.onConnectionsChange;
    nodeType.prototype.onConnectionsChange = function (...args) {
      const result = originalOnConnectionsChange?.apply(this, args);
      const link = args[3];
      colorSlot(args[4]);
      colorLink(this.graph, link);
      colorNodeSlots(this);
      this.graph?.setDirtyCanvas?.(true, true);
      return result;
    };

    const originalOnDrawForeground = nodeType.prototype.onDrawForeground;
    nodeType.prototype.onDrawForeground = function (...args) {
      translateNodeLabels(this, nodeData.name);
      return withTranslatedWidgetNames(this, () => originalOnDrawForeground?.apply(this, args));
    };

    const originalOnDrawBackground = nodeType.prototype.onDrawBackground;
    nodeType.prototype.onDrawBackground = function (...args) {
      translateNodeLabels(this, nodeData.name);
      return withTranslatedWidgetNames(this, () => originalOnDrawBackground?.apply(this, args));
    };

    if (nodeData.name === LOAD_AUDIO_NODE_TYPE) {
      registerLoadAudioNode(nodeType, wrapOnNodeCreated, api);
      return;
    }

    if (nodeData.name === AUDIO_ENSEMBLE_NODE_TYPE) {
      registerAudioEnsembleNode(nodeType, wrapOnNodeCreated);
      return;
    }

    if (SEPARATE_NODE_TYPES.has(nodeData.name) || SEPARATE_LIST_NODE_TYPES.has(nodeData.name)) {
      registerSeparateNode(nodeType, wrapOnNodeCreated, api);
      if (FIXED_420_NODE_TYPES.has(nodeData.name)) {
        registerFixedWidthOnlyNode(wrapOnNodeCreated, FIXED_420_NODE_WIDTH);
      }
      return;
    }

    if (nodeData.name === SAVE_AUDIO_NODE_TYPE) {
      registerSaveAudioNode(nodeType, wrapOnNodeCreated);
      return;
    }

    if (FIXED_260_NODE_TYPES.has(nodeData.name)) {
      registerFixedWidthNode(wrapOnNodeCreated);
      return;
    }

    if (FIXED_200_NODE_TYPES.has(nodeData.name)) {
      registerFixedWidthNode(wrapOnNodeCreated, FIXED_200_NODE_WIDTH);
      return;
    }

    if (FIXED_420_NODE_TYPES.has(nodeData.name)) {
      registerFixedWidthOnlyNode(wrapOnNodeCreated, FIXED_420_NODE_WIDTH);
      return;
    }

  },
});