Merge pull request #5 from lerignoux/multi_image

Adding multi-view support.

.github/workflows/publish_action.yml

@@ -0,0 +1,24 @@
+name: Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+    paths:
+      - "pyproject.toml"
+
+permissions:
+  issues: write
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    if: ${{ github.repository_owner == 'lerignoux' }}
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Publish Custom Node
+        uses: Comfy-Org/publish-node-action@v1
+        with:
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} ## Add your own personal access token to your Github Repository secrets and reference it here.

.github/workflows/python-lint.yml

@@ -0,0 +1,38 @@
+name: Python Lint Check
+
+on:
+  push:
+    branches: [main, master, develop]
+  pull_request:
+    branches: [main, master, develop]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    container:
+      image: python:3.9-alpine
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Install dependencies
+        run: |
+          # Update package manager and install required packages
+          apk update
+          apk add --no-cache findutils
+          python -m pip install --upgrade pip
+          pip install pycodestyle
+
+      - name: Show pycodestyle configuration
+        run: |
+          echo "Running pycodestyle with setup.cfg configuration:"
+          echo "- Max line length: 120"
+          echo "- Excluded directories: Stable3DGen/, __pycache__/"
+          echo "- Files being checked:"
+          find . -name "*.py" -not -path "./Stable3DGen/*" -not -path "./__pycache__/*"
+
+      - name: Run pycodestyle
+        run: |
+          # Run pycodestyle using setup.cfg configuration
+          pycodestyle .

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

@@ -0,0 +1,204 @@
+{
+  "id": "a4bc9489-299f-499a-b427-ef84a35c34f6",
+  "revision": 0,
+  "last_node_id": 19,
+  "last_link_id": 35,
+  "nodes": [
+    {
+      "id": 16,
+      "type": "Stable3DLoadModels",
+      "pos": [122.61805725097656, -337.5390625],
+      "size": [315, 126],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "hi3dgen pipeline",
+          "type": "HI3DGEN_PIPELINE",
+          "links": [31, 33]
+        },
+        {
+          "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",
+        null
+      ]
+    },
+    {
+      "id": 13,
+      "type": "PreviewImage",
+      "pos": [989.544677734375, -49.72162628173828],
+      "size": [210, 222],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 26
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      },
+      "widgets_values": [null]
+    },
+    {
+      "id": 12,
+      "type": "Stable3DPreprocessImage",
+      "pos": [535.5403442382812, -187.32595825195312],
+      "size": [367.79998779296875, 74],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "hi3dgen_pipeline",
+          "type": "HI3DGEN_PIPELINE",
+          "link": 33
+        },
+        {
+          "name": "normal_predictor",
+          "type": "STABLE3D_NORMAL",
+          "link": 32
+        },
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 34
+        }
+      ],
+      "outputs": [
+        {
+          "name": "normal_images",
+          "type": "IMAGE",
+          "links": [26, 27]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "Stable3DPreprocessImage"
+      },
+      "widgets_values": [null]
+    },
+    {
+      "id": 17,
+      "type": "LoadImage",
+      "pos": [127.95042419433594, -40.27117919921875],
+      "size": [278, 102],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [34]
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": ["example.png", "image", null]
+    },
+    {
+      "id": 19,
+      "type": "Preview3D",
+      "pos": [1478.064453125, -362.6379699707031],
+      "size": [400, 526],
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "camera_info",
+          "shape": 7,
+          "type": "LOAD3D_CAMERA",
+          "link": null
+        },
+        {
+          "name": "model_file",
+          "type": "STRING",
+          "widget": {
+            "name": "model_file"
+          },
+          "link": 35
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "Preview3D"
+      },
+      "widgets_values": ["", "", null]
+    },
+    {
+      "id": 11,
+      "type": "Stable3DGenerate3D",
+      "pos": [978.2960205078125, -371.08160400390625],
+      "size": [393, 198],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "hi3dgen_pipeline",
+          "type": "HI3DGEN_PIPELINE",
+          "link": 31
+        },
+        {
+          "name": "normal_images",
+          "type": "IMAGE",
+          "link": 27
+        }
+      ],
+      "outputs": [
+        {
+          "name": "mesh_file_path",
+          "type": "STRING",
+          "links": [35]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "Stable3DGenerate3D"
+      },
+      "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"]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1.3310000000000004,
+      "offset": [-48.53984360644931, 424.3690096481302]
+    },
+    "frontendVersion": "1.23.4"
+  },
+  "version": 0.4
+}

pyproject.toml

@@ -1,16 +1,21 @@
 [project]
 dependencies = [
+    "scikit-image~=0.25",
+    "spconv-cu120==2.3.6",
+    "torch==2.5.1",
+    "torchaudio==2.5.1",
+    "torchsde==0.2.6",
+    "torchvision==0.20.1",
+    "trimesh~=4.7",
+    "xformers~=0.0",
     "diffusers~=0.28",
     "accelerate~=1.9",
-    "triton~=3.2",
+    "triton=~3.1",
     "kornia~=0.8",
     "timm~=0.6",
     "transformers~=4.46",
-    "trimesh~=4.7",
-    "scikit-image~=0.25",
 ]
 
-
 name = "ComfyUI-Stable3DGen"
 description = "A ComfyUI custom node to generate 3D assets using Stable3D"
 version = "1.0.0"

requirements.txt

@@ -1,20 +1,19 @@
 --extra-index-url https://download.pytorch.org/whl/cu124
+scikit-image~=0.25
+spconv-cu120==2.3.6
+torch==2.5.1
+torchaudio==2.5.1
+torchsde==0.2.6
+torchvision==0.20.1
+trimesh~=4.7
+xformers~=0.0
+
 # For StableNormal
 diffusers~=0.28
 accelerate~=1.9
-triton
+triton~=3.1
 
 # For BirefNet
 kornia~=0.8
 timm~=0.6
 transformers~=4.46
-
-trimesh~=4.7
-scikit-image~=0.25
-xformers~=0.0
-
-torch==2.5.1
-torchaudio==2.5.1
-torchvision==0.20.1
-torchsde
-spconv-cu120==2.3.6

setup.cfg

@@ -0,0 +1,5 @@
+[pycodestyle]
+max-line-length = 120
+exclude = Stable3DGen/,__pycache__/
+statistics = True
+count = True

stable_3d.py

@@ -8,10 +8,11 @@ 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'))
-from hi3dgen.pipelines import Hi3DGenPipeline
+from hi3dgen.pipelines import Hi3DGenPipeline  # noqa E402
 
 log = logging.getLogger(__name__)
 MAX_SEED = numpy.iinfo(numpy.int32).max
@@ -95,7 +96,11 @@ class Stable3DLoadModels:
             log.info(f"Already cached at: {local_path}")
             return local_path
         log.info(f"Downloading and caching model: {model_id} to trellis model folder")
-        local_path = snapshot_download(repo_id=model_id, local_dir=os.path.join(self.models_path, model_id), force_download=False)
+        local_path = snapshot_download(
+            repo_id=model_id,
+            local_dir=os.path.join(self.models_path, model_id),
+            force_download=False
+        )
         return local_path
 
     @set_spconv_algo
@@ -112,12 +117,17 @@ class Stable3DLoadModels:
             loaded_models.append(self.download_model(model_id))
         # Loads yoso pedictor model to ~/models/trellis and StableNormal_turbo predictor library to ~/.cache/torch/hub/
         yozo_model_folder = os.path.join(self.models_path, "Stable-X")
-        normal_predictor = torch.hub.load("hugoycj/StableNormal", "StableNormal_turbo", trust_repo=True, yoso_version='yoso-normal-v1-8-1', local_cache_dir=yozo_model_folder)
+        normal_predictor = torch.hub.load(
+            "hugoycj/StableNormal",
+            "StableNormal_turbo",
+            trust_repo=True,
+            yoso_version='yoso-normal-v1-8-1',
+            local_cache_dir=yozo_model_folder
+        )
         # 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)
@@ -144,7 +154,7 @@ class Stable3DPreprocessImage:
                     "STABLE3D_NORMAL",
                     {"tooltip": "The normal predictor model to generate the image normal."}
                 ),
-                "image": ("IMAGE",)
+                "images": ("IMAGE",)
             },
         }
 
@@ -170,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:
@@ -293,32 +320,44 @@ 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,)
 
     @classmethod
-    def IS_CHANGED(s, images, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps):
+    def IS_CHANGED(
+        s, images, seed,
+        ss_guidance_strength, ss_sampling_steps,
+        slat_guidance_strength, slat_sampling_steps
+    ):
         # FIXME We should properly handle re-generation depending on the input parameters.
         return float("NaN")