From cfed768528e5f8347df8a692efeb1b0ed7d2b06f Mon Sep 17 00:00:00 2001
From: Robert Wojciechowski <hello@robertvoy.com>
Date: Sun, 1 Mar 2026 23:31:56 +0000
Subject: [PATCH] Apply dynamic num_branches sockets to DistributedJoin UI

---
 web/branch_outputs.js | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/web/branch_outputs.js b/web/branch_outputs.js
index 7409f21..bd983a1 100644
--- a/web/branch_outputs.js
+++ b/web/branch_outputs.js
@@ -1,6 +1,6 @@
 import { app } from "/scripts/app.js";
 
-const NODE_CLASS = "DistributedBranch";
+const NODE_CLASSES = new Set(["DistributedBranch", "DistributedJoin"]);
 const MIN_BRANCHES = 2;
 const MAX_BRANCHES = 10;
 
@@ -13,7 +13,7 @@ function clampBranchCount(value) {
 app.registerExtension({
     name: "Distributed.BranchOutputs",
     async nodeCreated(node) {
-        if (node.comfyClass !== NODE_CLASS) return;
+        if (!NODE_CLASSES.has(node.comfyClass)) return;
 
         const updateOutputs = () => {
             if (!node.widgets) return;
@@ -77,7 +77,7 @@ app.registerExtension({
     },
 
     nodeBeforeRemove(node) {
-        if (node.comfyClass === NODE_CLASS && node._branchOutputsCleanup) {
+        if (NODE_CLASSES.has(node.comfyClass) && node._branchOutputsCleanup) {
             node._branchOutputsCleanup();
         }
     },