From 3e25e08b10e7dda5af8d951d03c0520dd9b07f73 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Wed, 26 Mar 2025 16:00:54 -0400 Subject: [PATCH] [Bug] Fix broken output link id during reroute migration (#3244) --- src/utils/migration/migrateReroute.ts | 31 ++++++- .../workflows/reroute/native/branching.json | 90 +++++++++---------- .../reroute/native/single_connected.json | 88 +++++++++--------- 3 files changed, 116 insertions(+), 93 deletions(-) diff --git a/src/utils/migration/migrateReroute.ts b/src/utils/migration/migrateReroute.ts index 69252bcd1..95146d6e4 100644 --- a/src/utils/migration/migrateReroute.ts +++ b/src/utils/migration/migrateReroute.ts @@ -1,3 +1,5 @@ +import _ from 'lodash' + import type { ComfyLink, ComfyNode, @@ -67,10 +69,15 @@ export function createNewLinks( workflow: WorkflowJSON04, rerouteMap: Map ): { + nodes: ComfyNode[] reroutes: Reroute[] links: ComfyLink[] linkExtensions: LinkExtension[] } { + const nodeById = _.keyBy( + workflow.nodes.filter((node) => node.type !== 'Reroute').map(_.cloneDeep), + 'id' + ) const links: ComfyLink[] = [] const linkExtensions: LinkExtension[] = [] @@ -153,6 +160,23 @@ export function createNewLinks( validLinkExtensions.push(linkExtension) chainedReroutes.forEach((reroute) => validReroutes.add(reroute)) + + // Update source node's output slot's link ids to point to the new link. + const sourceNode = nodeById[sourceNodeId] + if (!sourceNode) { + throw new Error( + `Corrupted workflow: Source node ${sourceNodeId} not found` + ) + } + const outputSlot = sourceNode.outputs?.[sourceSlot] + if (!outputSlot) { + throw new Error( + `Corrupted workflow: Output slot ${sourceSlot} not found` + ) + } + outputSlot.links = outputSlot.links?.map((l) => + l === rerouteInputLink[0] ? linkId : l + ) } entry = undefined } @@ -160,6 +184,7 @@ export function createNewLinks( } return { + nodes: Object.values(nodeById), links, linkExtensions: validLinkExtensions, reroutes: Array.from(validReroutes) @@ -192,16 +217,14 @@ export const migrateLegacyRerouteNodes = ( const rerouteMap = createReroutePoints(legacyRerouteNodes) // Create new links and link extensions - const { links, linkExtensions, reroutes } = createNewLinks( + const { nodes, links, linkExtensions, reroutes } = createNewLinks( workflow, rerouteMap ) // Update the workflow newWorkflow.links = links - newWorkflow.nodes = newWorkflow.nodes.filter( - (node) => node.type !== 'Reroute' - ) + newWorkflow.nodes = nodes newWorkflow.extra.reroutes = reroutes newWorkflow.extra.linkExtensions = linkExtensions diff --git a/tests-ui/tests/utils/migration/workflows/reroute/native/branching.json b/tests-ui/tests/utils/migration/workflows/reroute/native/branching.json index 70784b576..49d159959 100644 --- a/tests-ui/tests/utils/migration/workflows/reroute/native/branching.json +++ b/tests-ui/tests/utils/migration/workflows/reroute/native/branching.json @@ -2,6 +2,51 @@ "last_node_id": 27, "last_link_id": 34, "nodes": [ + { + "id": 4, + "type": "CheckpointLoaderSimple", + "pos": [ + 47.948699951171875, + 239.2628173828125 + ], + "size": [ + 315, + 98 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "slot_index": 0, + "links": [] + }, + { + "name": "CLIP", + "type": "CLIP", + "slot_index": 1, + "links": [] + }, + { + "name": "VAE", + "type": "VAE", + "slot_index": 2, + "links": [ + 13, + 21 + ] + } + ], + "properties": { + "Node name for S&R": "CheckpointLoaderSimple" + }, + "widgets_values": [ + "v1-5-pruned-emaonly.safetensors" + ] + }, { "id": 12, "type": "VAEDecode", @@ -40,51 +85,6 @@ }, "widgets_values": [] }, - { - "id": 4, - "type": "CheckpointLoaderSimple", - "pos": [ - 47.948699951171875, - 239.2628173828125 - ], - "size": [ - 315, - 98 - ], - "flags": {}, - "order": 0, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "MODEL", - "type": "MODEL", - "slot_index": 0, - "links": [] - }, - { - "name": "CLIP", - "type": "CLIP", - "slot_index": 1, - "links": [] - }, - { - "name": "VAE", - "type": "VAE", - "slot_index": 2, - "links": [ - 13, - 31 - ] - } - ], - "properties": { - "Node name for S&R": "CheckpointLoaderSimple" - }, - "widgets_values": [ - "v1-5-pruned-emaonly.safetensors" - ] - }, { "id": 26, "type": "VAEDecode", diff --git a/tests-ui/tests/utils/migration/workflows/reroute/native/single_connected.json b/tests-ui/tests/utils/migration/workflows/reroute/native/single_connected.json index 95cb451f2..d3d1b3244 100644 --- a/tests-ui/tests/utils/migration/workflows/reroute/native/single_connected.json +++ b/tests-ui/tests/utils/migration/workflows/reroute/native/single_connected.json @@ -2,6 +2,50 @@ "last_node_id": 24, "last_link_id": 30, "nodes": [ + { + "id": 4, + "type": "CheckpointLoaderSimple", + "pos": [ + 160, + 240 + ], + "size": [ + 315, + 98 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "slot_index": 0, + "links": [] + }, + { + "name": "CLIP", + "type": "CLIP", + "slot_index": 1, + "links": [] + }, + { + "name": "VAE", + "type": "VAE", + "slot_index": 2, + "links": [ + 21 + ] + } + ], + "properties": { + "Node name for S&R": "CheckpointLoaderSimple" + }, + "widgets_values": [ + "v1-5-pruned-emaonly.safetensors" + ] + }, { "id": 12, "type": "VAEDecode", @@ -39,50 +83,6 @@ "Node name for S&R": "VAEDecode" }, "widgets_values": [] - }, - { - "id": 4, - "type": "CheckpointLoaderSimple", - "pos": [ - 160, - 240 - ], - "size": [ - 315, - 98 - ], - "flags": {}, - "order": 0, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "MODEL", - "type": "MODEL", - "slot_index": 0, - "links": [] - }, - { - "name": "CLIP", - "type": "CLIP", - "slot_index": 1, - "links": [] - }, - { - "name": "VAE", - "type": "VAE", - "slot_index": 2, - "links": [ - 13 - ] - } - ], - "properties": { - "Node name for S&R": "CheckpointLoaderSimple" - }, - "widgets_values": [ - "v1-5-pruned-emaonly.safetensors" - ] } ], "links": [