Fix primitive resize (#683)

* Fix primitive resize on load
Fixes #676

* Add test

* Add playwright test

* Update test expectations [skip ci]

---------

Co-authored-by: huchenlei <chenlei.hu@mail.utoronto.ca>
Co-authored-by: github-actions <github-actions@github.com>

browser_tests/assets/primitive_node.json

@@ -0,0 +1,145 @@
+{
+  "last_node_id": 2,
+  "last_link_id": 1,
+  "nodes": [
+    {
+      "id": 2,
+      "type": "KSampler",
+      "pos": {
+        "0": 521.0906982421875,
+        "1": 40.999996185302734,
+        "2": 0,
+        "3": 0,
+        "4": 0,
+        "5": 0,
+        "6": 0,
+        "7": 0,
+        "8": 0,
+        "9": 0
+      },
+      "size": {
+        "0": 315,
+        "1": 262
+      },
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "model",
+          "type": "MODEL",
+          "link": null
+        },
+        {
+          "name": "positive",
+          "type": "CONDITIONING",
+          "link": null
+        },
+        {
+          "name": "negative",
+          "type": "CONDITIONING",
+          "link": null
+        },
+        {
+          "name": "latent_image",
+          "type": "LATENT",
+          "link": null
+        },
+        {
+          "name": "steps",
+          "type": "INT",
+          "link": 1,
+          "widget": {
+            "name": "steps"
+          }
+        }
+      ],
+      "outputs": [
+        {
+          "name": "LATENT",
+          "type": "LATENT",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "KSampler"
+      },
+      "widgets_values": [
+        0,
+        "randomize",
+        20,
+        8,
+        "euler",
+        "normal",
+        1
+      ]
+    },
+    {
+      "id": 1,
+      "type": "PrimitiveNode",
+      "pos": {
+        "0": 15,
+        "1": 46,
+        "2": 0,
+        "3": 0,
+        "4": 0,
+        "5": 0,
+        "6": 0,
+        "7": 0,
+        "8": 0,
+        "9": 0
+      },
+      "size": [
+        446.96645387135936,
+        108.34243389566905
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "INT",
+          "type": "INT",
+          "links": [
+            1
+          ],
+          "slot_index": 0,
+          "widget": {
+            "name": "steps"
+          }
+        }
+      ],
+      "properties": {
+        "Run widget replace on values": false
+      },
+      "widgets_values": [
+        20,
+        "fixed"
+      ]
+    }
+  ],
+  "links": [
+    [
+      1,
+      1,
+      0,
+      2,
+      4,
+      "INT"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1,
+      "offset": [
+        0,
+        0
+      ]
+    }
+  },
+  "version": 0.4
+}

browser_tests/primitiveNode.spec.ts

@@ -0,0 +1,9 @@
+import { expect } from '@playwright/test'
+import { comfyPageFixture as test } from './ComfyPage'
+
+test.describe('Primitive Node', () => {
+  test('Can load with correct size', async ({ comfyPage }) => {
+    await comfyPage.loadWorkflow('primitive_node')
+    await expect(comfyPage.canvas).toHaveScreenshot('primitive_node.png')
+  })
+})

src/extensions/core/widgetInputs.ts

@@ -203,6 +203,8 @@ class PrimitiveNode extends LGraphNode {
       type = 'COMBO'
     }
 
+    // Store current size as addWidget resizes the node
+    const size = this.size
     let widget
     if (type in ComfyWidgets) {
       widget = (ComfyWidgets[type](this, 'value', inputData, app) || {}).widget
@@ -261,8 +263,14 @@ class PrimitiveNode extends LGraphNode {
       return r
     }
 
+    // Use the biggest dimensions in case the widgets caused the node to grow
+    this.size = [
+      Math.max(this.size[0], size[0]),
+      Math.max(this.size[1], size[1])
+    ]
+
     if (!recreating) {
-      // Grow our node if required
+      // Grow our node more if required
       const sz = this.computeSize()
       if (this.size[0] < sz[0]) {
         this.size[0] = sz[0]

tests-ui/tests/widgetInputs.test.ts

@@ -622,5 +622,21 @@ describe('widget inputs', () => {
         })
       })
     })
+    it('primitive maintains size on reload', async () => {
+      const { ez, graph } = await start()
+      const primitive = ez.PrimitiveNode()
+      const image = ez.EmptyImage()
+
+      image.widgets.width.convertToInput()
+      primitive.outputs[0].connectTo(image.inputs.width)
+
+      primitive.node.size = [999, 999]
+
+      await checkBeforeAndAfterReload(graph, async (r) => {
+        const { node } = graph.find(primitive)
+        expect(node.size[0]).toBe(999)
+        expect(node.size[1]).toBe(999)
+      })
+    })
   })
 })