Merge pull request #5 from lerignoux/multi_image

Adding multi-view support.

example_workflows/Stable3DGen Multiview.json

@@ -0,0 +1,284 @@
+{
+  "id": "a4bc9489-299f-499a-b427-ef84a35c34f6",
+  "revision": 0,
+  "last_node_id": 24,
+  "last_link_id": 44,
+  "nodes": [
+    {
+      "id": 19,
+      "type": "Preview3D",
+      "pos": [1478.064453125, -362.6379699707031],
+      "size": [400, 526],
+      "flags": {},
+      "order": 7,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "camera_info",
+          "shape": 7,
+          "type": "LOAD3D_CAMERA",
+          "link": null
+        },
+        {
+          "name": "model_file",
+          "type": "STRING",
+          "widget": {
+            "name": "model_file"
+          },
+          "link": 44
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "Preview3D",
+        "Last Time Model File": "2025-09-23-190112_mesh.glb",
+        "Camera Info": {
+          "position": {
+            "x": 1.0911383750728778,
+            "y": -0.5286630089592208,
+            "z": 13.277660350003037
+          },
+          "target": {
+            "x": 0,
+            "y": 2.5,
+            "z": 0
+          },
+          "zoom": 1,
+          "cameraType": "perspective"
+        }
+      },
+      "widgets_values": ["2025-09-23-190112_mesh.glb", ""]
+    },
+    {
+      "id": 17,
+      "type": "LoadImage",
+      "pos": [-170.2718048095703, -14.604534149169922],
+      "size": [278, 314],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [38]
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": ["example.png", "image"]
+    },
+    {
+      "id": 21,
+      "type": "LoadImage",
+      "pos": [-170.3321990966797, 358.9052429199219],
+      "size": [274.080078125, 314],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [39]
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": ["example.png", "image"]
+    },
+    {
+      "id": 20,
+      "type": "ImageBatch",
+      "pos": [382.28778076171875, 144.34243774414062],
+      "size": [140, 46],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "image1",
+          "type": "IMAGE",
+          "link": 38
+        },
+        {
+          "name": "image2",
+          "type": "IMAGE",
+          "link": 39
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [40]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ImageBatch"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 13,
+      "type": "PreviewImage",
+      "pos": [1032.534423828125, 5.965235710144043],
+      "size": [210, 246],
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 26
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 12,
+      "type": "Stable3DPreprocessImage",
+      "pos": [535.5403442382812, -187.32595825195312],
+      "size": [367.79998779296875, 86],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "hi3dgen_pipeline",
+          "type": "HI3DGEN_PIPELINE",
+          "link": 33
+        },
+        {
+          "name": "normal_predictor",
+          "type": "STABLE3D_NORMAL",
+          "link": 32
+        },
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 40
+        },
+        {
+          "name": "image",
+          "type": "IMAGE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "normal_images",
+          "type": "IMAGE",
+          "links": [26, 42]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "Stable3DPreprocessImage"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 16,
+      "type": "Stable3DLoadModels",
+      "pos": [122.61805725097656, -337.5390625],
+      "size": [315, 126],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "hi3dgen pipeline",
+          "type": "HI3DGEN_PIPELINE",
+          "links": [33, 43]
+        },
+        {
+          "name": "Normal predictor",
+          "type": "STABLE3D_NORMAL",
+          "links": [32]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "Stable3DLoadModels"
+      },
+      "widgets_values": [
+        "Stable-X/trellis-normal-v0-1",
+        "Stable-X/yoso-normal-v1-8-1",
+        "ZhengPeng7/BiRefNet"
+      ]
+    },
+    {
+      "id": 24,
+      "type": "Stable3DGenerate3D",
+      "pos": [999.1272583007812, -324.4736022949219],
+      "size": [288.08984375, 198],
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "hi3dgen_pipeline",
+          "type": "HI3DGEN_PIPELINE",
+          "link": 43
+        },
+        {
+          "name": "normal_images",
+          "type": "IMAGE",
+          "link": 42
+        }
+      ],
+      "outputs": [
+        {
+          "name": "mesh_file_path",
+          "type": "STRING",
+          "links": [44]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "Stable3DGenerate3D"
+      },
+      "widgets_values": [-1, "randomize", 3, 50, 3, 6]
+    }
+  ],
+  "links": [
+    [26, 12, 0, 13, 0, "IMAGE"],
+    [32, 16, 1, 12, 1, "STABLE3D_NORMAL"],
+    [33, 16, 0, 12, 0, "HI3DGEN_PIPELINE"],
+    [38, 17, 0, 20, 0, "IMAGE"],
+    [39, 21, 0, 20, 1, "IMAGE"],
+    [40, 20, 0, 12, 2, "IMAGE"],
+    [42, 12, 0, 24, 1, "IMAGE"],
+    [43, 16, 0, 24, 0, "HI3DGEN_PIPELINE"],
+    [44, 24, 0, 19, 1, "STRING"]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1.1,
+      "offset": [-86.75339658576321, 505.51848895776305]
+    },
+    "frontendVersion": "1.25.11",
+    "ue_links": []
+  },
+  "version": 0.4
+}

example_workflows/Stable3DGen.json

@@ -7,14 +7,8 @@
     {
       "id": 16,
       "type": "Stable3DLoadModels",
-      "pos": [
-        122.61805725097656,
-        -337.5390625
-      ],
-      "size": [
-        315,
-        126
-      ],
+      "pos": [122.61805725097656, -337.5390625],
+      "size": [315, 126],
       "flags": {},
       "order": 0,
       "mode": 0,
@@ -23,17 +17,12 @@
         {
           "name": "hi3dgen pipeline",
           "type": "HI3DGEN_PIPELINE",
-          "links": [
-            31,
-            33
-          ]
+          "links": [31, 33]
         },
         {
           "name": "Normal predictor",
           "type": "STABLE3D_NORMAL",
-          "links": [
-            32
-          ]
+          "links": [32]
         }
       ],
       "properties": {
@@ -49,14 +38,8 @@
     {
       "id": 13,
       "type": "PreviewImage",
-      "pos": [
-        989.544677734375,
-        -49.72162628173828
-      ],
-      "size": [
-        210,
-        222
-      ],
+      "pos": [989.544677734375, -49.72162628173828],
+      "size": [210, 222],
       "flags": {},
       "order": 3,
       "mode": 0,
@@ -71,21 +54,13 @@
       "properties": {
         "Node name for S&R": "PreviewImage"
       },
-      "widgets_values": [
-        null
-      ]
+      "widgets_values": [null]
     },
     {
       "id": 12,
       "type": "Stable3DPreprocessImage",
-      "pos": [
-        535.5403442382812,
-        -187.32595825195312
-      ],
-      "size": [
-        367.79998779296875,
-        74
-      ],
+      "pos": [535.5403442382812, -187.32595825195312],
+      "size": [367.79998779296875, 74],
       "flags": {},
       "order": 2,
       "mode": 0,
@@ -101,7 +76,7 @@
           "link": 32
         },
         {
-          "name": "image",
+          "name": "images",
           "type": "IMAGE",
           "link": 34
         }
@@ -110,30 +85,19 @@
         {
           "name": "normal_images",
           "type": "IMAGE",
-          "links": [
-            26,
-            27
-          ]
+          "links": [26, 27]
         }
       ],
       "properties": {
         "Node name for S&R": "Stable3DPreprocessImage"
       },
-      "widgets_values": [
-        null
-      ]
+      "widgets_values": [null]
     },
     {
       "id": 17,
       "type": "LoadImage",
-      "pos": [
-        127.95042419433594,
-        -40.27117919921875
-      ],
-      "size": [
-        278,
-        102
-      ],
+      "pos": [127.95042419433594, -40.27117919921875],
+      "size": [278, 102],
       "flags": {},
       "order": 1,
       "mode": 0,
@@ -142,9 +106,7 @@
         {
           "name": "IMAGE",
           "type": "IMAGE",
-          "links": [
-            34
-          ]
+          "links": [34]
         },
         {
           "name": "MASK",
@@ -155,23 +117,13 @@
       "properties": {
         "Node name for S&R": "LoadImage"
       },
-      "widgets_values": [
-        "example.png",
-        "image",
-        null
-      ]
+      "widgets_values": ["example.png", "image", null]
     },
     {
       "id": 19,
       "type": "Preview3D",
-      "pos": [
-        1478.064453125,
-        -362.6379699707031
-      ],
-      "size": [
-        400,
-        526
-      ],
+      "pos": [1478.064453125, -362.6379699707031],
+      "size": [400, 526],
       "flags": {},
       "order": 5,
       "mode": 0,
@@ -195,23 +147,13 @@
       "properties": {
         "Node name for S&R": "Preview3D"
       },
-      "widgets_values": [
-        "",
-        "",
-        null
-      ]
+      "widgets_values": ["", "", null]
     },
     {
       "id": 11,
       "type": "Stable3DGenerate3D",
-      "pos": [
-        978.2960205078125,
-        -371.08160400390625
-      ],
-      "size": [
-        393,
-        198
-      ],
+      "pos": [978.2960205078125, -371.08160400390625],
+      "size": [393, 198],
       "flags": {},
       "order": 4,
       "mode": 0,
@@ -231,94 +173,32 @@
         {
           "name": "mesh_file_path",
           "type": "STRING",
-          "links": [
-            35
-          ]
+          "links": [35]
         }
       ],
       "properties": {
         "Node name for S&R": "Stable3DGenerate3D"
       },
-      "widgets_values": [
-        1028886841,
-        "randomize",
-        3,
-        50,
-        3,
-        6,
-        null
-      ]
+      "widgets_values": [1028886841, "randomize", 3, 50, 3, 6, null]
     }
   ],
   "links": [
-    [
-      26,
-      12,
-      0,
-      13,
-      0,
-      "IMAGE"
-    ],
-    [
-      27,
-      12,
-      0,
-      11,
-      1,
-      "IMAGE"
-    ],
-    [
-      31,
-      16,
-      0,
-      11,
-      0,
-      "HI3DGEN_PIPELINE"
-    ],
-    [
-      32,
-      16,
-      1,
-      12,
-      1,
-      "STABLE3D_NORMAL"
-    ],
-    [
-      33,
-      16,
-      0,
-      12,
-      0,
-      "HI3DGEN_PIPELINE"
-    ],
-    [
-      34,
-      17,
-      0,
-      12,
-      2,
-      "IMAGE"
-    ],
-    [
-      35,
-      11,
-      0,
-      19,
-      1,
-      "STRING"
-    ]
+    [26, 12, 0, 13, 0, "IMAGE"],
+    [27, 12, 0, 11, 1, "IMAGE"],
+    [31, 16, 0, 11, 0, "HI3DGEN_PIPELINE"],
+    [32, 16, 1, 12, 1, "STABLE3D_NORMAL"],
+    [33, 16, 0, 12, 0, "HI3DGEN_PIPELINE"],
+    [34, 17, 0, 12, 2, "IMAGE"],
+    [35, 11, 0, 19, 1, "STRING"]
   ],
   "groups": [],
   "config": {},
   "extra": {
     "ds": {
       "scale": 1.3310000000000004,
-      "offset": [
-        -48.53984360644931,
-        424.3690096481302
-      ]
+      "offset": [-48.53984360644931, 424.3690096481302]
     },
     "frontendVersion": "1.23.4"
   },
   "version": 0.4
-}
+}

pyproject.toml

@@ -13,10 +13,9 @@ dependencies = [
     "triton=~3.1",
     "kornia~=0.8",
     "timm~=0.6",
-    "transformers~=4.46"
+    "transformers~=4.46",
 ]
 
-
 name = "ComfyUI-Stable3DGen"
 description = "A ComfyUI custom node to generate 3D assets using Stable3D"
 version = "1.0.0"

stable_3d.py

@@ -8,6 +8,7 @@ from huggingface_hub import snapshot_download
 from PIL import Image
 from transformers import AutoModelForImageSegmentation
 
+from comfy.utils import common_upscale
 import folder_paths
 
 sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'Stable3DGen'))
@@ -126,8 +127,7 @@ class Stable3DLoadModels:
         # download dinov2 feature detection model and library to ~/.cache/torch/hub/
         torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14_reg', pretrained=True)
 
-        trellis_folder = folder_paths.get_folder_paths("trellis")[0]
-        hi3dgen_pipeline = Hi3DGenPipeline.from_pretrained(os.path.join(trellis_folder, trellis_model))
+        hi3dgen_pipeline = Hi3DGenPipeline.from_pretrained(os.path.join(self.models_path, trellis_model))
         hi3dgen_pipeline.cuda()
 
         self.load_birefnet_model(hi3dgen_pipeline, birefnet_model)
@@ -154,7 +154,7 @@ class Stable3DPreprocessImage:
                     "STABLE3D_NORMAL",
                     {"tooltip": "The normal predictor model to generate the image normal."}
                 ),
-                "image": ("IMAGE",)
+                "images": ("IMAGE",)
             },
         }
 
@@ -180,19 +180,36 @@ class Stable3DPreprocessImage:
         log.debug(f"normal_image saved as {path}")
         return path
 
-    def preprocess_image(self, hi3dgen_pipeline, normal_predictor, image):
-        # FIXME We should support properly batch mode here.
-        numpy_image = image.squeeze(0).cpu().numpy()
-        numpy_image_scaled = numpy.clip(numpy_image * 255, 0, 255).astype(numpy.uint8)
-        image = Image.fromarray(numpy_image_scaled)
+    @staticmethod
+    def uniformize_images(images):
+        for index in range(1, len(images)):
+            if images[index].shape[1:] != images[0].shape[1:]:
+                images[index] = common_upscale(
+                    images[index].movedim(-1, 1),
+                    images[0].shape[2],
+                    images[0].shape[1],
+                    "bilinear",
+                    "center"
+                ).movedim(1, -1)
+        return images
 
-        # FIXME We should properly handle batch here.
-        image = hi3dgen_pipeline.preprocess_image(image, resolution=1024)
-        normal_image = normal_predictor(image, resolution=768, match_input_resolution=True, data_type='object')
+    def preprocess_image(self, hi3dgen_pipeline, normal_predictor, images):
+        normal_images = []
+        for (_, img) in enumerate(images):
+            numpy_image = 255. * img.cpu().numpy()
+            numpy_image_scaled = Image.fromarray(numpy.clip(numpy_image, 0, 255).astype(numpy.uint8))
 
-        self.save_normal_image(normal_image)
+            image = hi3dgen_pipeline.preprocess_image(numpy_image_scaled, resolution=1024)
+            normal_image = normal_predictor(image, resolution=768, match_input_resolution=True, data_type='object')
 
-        return (self.pil_to_tensor(normal_image),)
+            self.save_normal_image(normal_image)
+            normal_images.append(self.pil_to_tensor(normal_image))
+
+        if len(normal_images) > 1:
+            output_image = torch.cat(self.uniformize_images(normal_images), dim=0)
+        else:
+            output_image = normal_images[0]
+        return (output_image,)
 
 
 class Stable3DGenerate3D:
@@ -303,28 +320,36 @@ class Stable3DGenerate3D:
             seed = numpy.random.randint(0, MAX_SEED)
 
         log.info("Starting 3d mesh generation.")
-        for (batch_number, image) in enumerate(normal_images):
+        pil_input = []
+        for (_, image) in enumerate(normal_images):
             numpy_image = 255. * image.cpu().numpy()
-            pil_image = Image.fromarray(numpy.clip(numpy_image, 0, 255).astype(numpy.uint8))
+            pil_input.append(Image.fromarray(numpy.clip(numpy_image, 0, 255).astype(numpy.uint8)))
 
-            outputs = hi3dgen_pipeline.run(
-                pil_image,
-                seed=seed,
-                formats=["mesh",],
-                preprocess_image=False,
-                sparse_structure_sampler_params={
-                    "steps": ss_sampling_steps,
-                    "cfg_strength": ss_guidance_strength,
-                },
-                slat_sampler_params={
-                    "steps": slat_sampling_steps,
-                    "cfg_strength": slat_guidance_strength,
-                },
-            )
-            generated_mesh = outputs['mesh'][0]
-            saved_path = self.save_3d_asset(generated_mesh)
+        method = 'run'
+        if len(pil_input) == 1:
+            pil_input = pil_input[0]
+        elif len(pil_input) > 1:
+            method = 'run_multi_image'
+        else:
+            raise ValueError("No image provided to run 3d pipeline.")
+
+        outputs = getattr(hi3dgen_pipeline, method)(
+            pil_input,
+            seed=seed,
+            formats=["mesh",],
+            preprocess_image=False,
+            sparse_structure_sampler_params={
+                "steps": ss_sampling_steps,
+                "cfg_strength": ss_guidance_strength,
+            },
+            slat_sampler_params={
+                "steps": slat_sampling_steps,
+                "cfg_strength": slat_guidance_strength,
+            },
+        )
+        generated_mesh = outputs['mesh'][0]
+        saved_path = self.save_3d_asset(generated_mesh)
 
-        # FiXME we should return all the files in case of batch.
         filename = saved_path.split("output/")[1]
         return (filename,)