Add torch 2.6.0 wheel for custom_rasterizer

.gitignore

@@ -9,3 +9,6 @@ logs/
 .idea
 tools/
 .vscode/
+build/
+dist/
+*.egg-info*

hy3dgen/texgen/custom_rasterizer/custom_rasterizer.egg-info/PKG-INFO

@@ -1,3 +0,0 @@
-Metadata-Version: 2.1
-Name: custom_rasterizer
-Version: 0.1

hy3dgen/texgen/custom_rasterizer/custom_rasterizer.egg-info/SOURCES.txt

@@ -1,12 +0,0 @@
-setup.py
-./custom_rasterizer/__init__.py
-./custom_rasterizer/io_glb.py
-./custom_rasterizer/io_obj.py
-./custom_rasterizer/render.py
-custom_rasterizer.egg-info/PKG-INFO
-custom_rasterizer.egg-info/SOURCES.txt
-custom_rasterizer.egg-info/dependency_links.txt
-custom_rasterizer.egg-info/top_level.txt
-lib/custom_rasterizer_kernel/grid_neighbor.cpp
-lib/custom_rasterizer_kernel/rasterizer.cpp
-lib/custom_rasterizer_kernel/rasterizer_gpu.cu

hy3dgen/texgen/custom_rasterizer/custom_rasterizer.egg-info/dependency_links.txt

@@ -1 +0,0 @@
-

hy3dgen/texgen/custom_rasterizer/custom_rasterizer.egg-info/top_level.txt

@@ -1,2 +0,0 @@
-custom_rasterizer
-custom_rasterizer_kernel

hy3dgen/texgen/custom_rasterizer/setup.py

@@ -1,6 +1,11 @@
 from setuptools import setup, find_packages
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+import torch
 
+torch_version = torch.__version__.split('+')[0].replace('.', '')
+cuda_version = torch.version.cuda.replace('.', '')
+
+version = f"0.1.0+torch{torch_version}.cuda{cuda_version}"
 # build custom rasterizer
 # build with `python setup.py install`
 # nvcc is needed
@@ -13,7 +18,7 @@ custom_rasterizer_module = CUDAExtension('custom_rasterizer_kernel', [
 
 setup(
     packages=find_packages(),
-    version='0.1',
+    version=version,
     name='custom_rasterizer',
     include_package_data=True,
     package_dir={'': '.'},
@@ -22,5 +27,5 @@ setup(
     ],
     cmdclass={
         'build_ext': BuildExtension
-    }
+    },   
 )

nodes.py

@@ -143,6 +143,9 @@ class DownloadAndLoadHy3DDelightModel:
             "required": {
                 "model": (["hunyuan3d-delight-v2-0"],),
             },
+            "optional": {
+                "compile_args": ("HY3DCOMPILEARGS", {"tooltip": "torch.compile settings, when connected to the model loader, torch.compile of the selected models is attempted. Requires Triton and torch 2.5.0 is recommended"}),
+            }
         }
 
     RETURN_TYPES = ("HY3DDIFFUSERSPIPE",)
@@ -150,9 +153,8 @@ class DownloadAndLoadHy3DDelightModel:
     FUNCTION = "loadmodel"
     CATEGORY = "Hunyuan3DWrapper"
 
-    def loadmodel(self, model):
+    def loadmodel(self, model, compile_args=None):
         device = mm.get_torch_device()
-        offload_device = mm.unet_offload_device()
 
         download_path = os.path.join(folder_paths.models_dir,"diffusers")
         model_path = os.path.join(download_path, model)
@@ -176,8 +178,17 @@ class DownloadAndLoadHy3DDelightModel:
         )
         delight_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(delight_pipe.scheduler.config)
         delight_pipe = delight_pipe.to(device, torch.float16)
+
         delight_pipe.enable_model_cpu_offload()
 
+        if compile_args is not None:
+            torch._dynamo.config.cache_size_limit = compile_args["dynamo_cache_size_limit"]
+            if compile_args["compile_transformer"]:
+                delight_pipe.unet = torch.compile(delight_pipe.unet)
+            if compile_args["compile_vae"]:
+                delight_pipe.vae = torch.compile(delight_pipe.vae)
+        
+        
         return (delight_pipe,)
         
 class Hy3DDelightImage:
@@ -242,6 +253,9 @@ class DownloadAndLoadHy3DPaintModel:
             "required": {
                 "model": (["hunyuan3d-paint-v2-0"],),
             },
+            "optional": {
+                "compile_args": ("HY3DCOMPILEARGS", {"tooltip": "torch.compile settings, when connected to the model loader, torch.compile of the selected models is attempted. Requires Triton and torch 2.5.0 is recommended"}),
+            }
         }
 
     RETURN_TYPES = ("HY3DDIFFUSERSPIPE",)
@@ -249,7 +263,7 @@ class DownloadAndLoadHy3DPaintModel:
     FUNCTION = "loadmodel"
     CATEGORY = "Hunyuan3DWrapper"
 
-    def loadmodel(self, model):
+    def loadmodel(self, model, compile_args=None):
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
 
@@ -311,6 +325,13 @@ class DownloadAndLoadHy3DPaintModel:
             feature_extractor=feature_extractor,
             )
         
+        if compile_args is not None:
+            torch._dynamo.config.cache_size_limit = compile_args["dynamo_cache_size_limit"]
+            if compile_args["compile_transformer"]:
+                pipeline.unet = torch.compile(pipeline.unet)
+            if compile_args["compile_vae"]:
+                pipeline.vae = torch.compile(pipeline.vae)
+
         pipeline.enable_model_cpu_offload()
         return (pipeline,)
 

readme.md

@@ -23,13 +23,18 @@ or with portable:
 
 For the texturegen part compilation is needed, I have included my compilations as a wheel for the rasterizer, and compiled .pyd for the mesh_processor (already in place), these are compiled for:
 
-- Windows 11 python 3.12 cu126 (works with torch build on 124)
+**Windows 11 python 3.12 cu126 (works with torch build on 124)**
 
 You would do `pip install wheels\custom_rasterizer-0.1-cp312-cp312-win_amd64.whl`
 
 or with portable (in `ComfyUI_windows_portable` -folder):
 `python_embeded\python.exe -m pip install ComfyUI\custom_nodes\ComfyUI-Hunyuan3DWrapper\wheels\custom_rasterizer-0.1-cp312-cp312-win_amd64.whl`
 
+**Windows 11 python 3.12 torch 2.6.0 + cu126**
+Current latest portable was updated to use pytorch 2.6.0, for this you should use new wheel:
+`python_embeded\python.exe -m pip install ComfyUI\custom_nodes\ComfyUI-Hunyuan3DWrapper\wheels\custom_rasterizer-0.1.0+torch260.cuda126-cp312-cp312-win_amd64.whl`
+
+
 This was tested to work on latest ComfyUI portable install
 
 If this doesn't work, you need to compile yourself: