add pbr glb create

Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/convert_utils.py

@@ -6,6 +6,9 @@ import base64
 import io
 import tempfile
 import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(__file__), "../../../"))
+from mesh_processor.fastmesh import FastMesh
 
 
 def combine_metallic_roughness(metallic_path, roughness_path, output_path):
@@ -82,7 +85,8 @@ def create_glb_with_pbr_materials(obj_path, textures_dict, output_path):
             with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_mr_file:
                 mr_combined_path = temp_mr_file.name
             temp_files.append(mr_combined_path)
-            combine_metallic_roughness(textures_dict["metallic"], textures_dict["roughness"], mr_combined_path)
+            # combine_metallic_roughness(textures_dict["metallic"], textures_dict["roughness"], mr_combined_path)
+            combine_metallic_roughness_fast(textures_dict["metallic"], textures_dict["roughness"], mr_combined_path)
             textures_dict["metallicRoughness"] = mr_combined_path
 
         # 6. 添加图像到GLTF
@@ -158,3 +162,17 @@ def create_glb_with_pbr_materials(obj_path, textures_dict, output_path):
                     os.remove(temp_file)
             except Exception as e:
                 print(f"Warning: Failed to remove temporary file {temp_file}: {e}")
+
+
+# FastMesh-based functions (new optimized versions)
+def combine_metallic_roughness_fast(metallic_path, roughness_path, output_path):
+    """Fast combine metallic and roughness maps using FastMesh"""
+    combined_array = FastMesh._combine_metallic_roughness(metallic_path, roughness_path)
+    combined = Image.fromarray(combined_array)
+    combined.save(output_path)
+    return output_path
+
+
+def create_glb_with_pbr_materials_fast(obj_path, textures_dict, output_path):
+    """Fast create GLB with PBR materials using FastMesh"""
+    FastMesh.create_glb_with_pbr_materials(obj_path, textures_dict, output_path)

Gen_3D_Modules/Hunyuan3D_2_1/hy3dpaint/textureGenPipeline.py

@@ -29,7 +29,7 @@ from .utils.pipeline_utils import ViewProcessor
 from .utils.image_super_utils import imageSuperNet
 from .utils.uvwrap_utils import mesh_uv_wrap
 from .DifferentiableRenderer.mesh_utils import convert_obj_to_glb
-from .convert_utils import create_glb_with_pbr_materials
+from .convert_utils import create_glb_with_pbr_materials, create_glb_with_pbr_materials_fast
 import warnings
 
 warnings.filterwarnings("ignore")
@@ -209,7 +209,9 @@ class Hunyuan3DPaintPipeline:
                 textures_dict['normal'] = normal_path
             
             try:
-                create_glb_with_pbr_materials(output_mesh_path, textures_dict, output_glb_path)
+                # create_glb_with_pbr_materials(output_mesh_path, textures_dict, output_glb_path)
+                print("Creating GLB with full PBR materials using FastMesh")
+                create_glb_with_pbr_materials_fast(output_mesh_path, textures_dict, output_glb_path)
                 print(f"Created GLB with full PBR materials: {output_glb_path}")
             except Exception as e:
                 print(f"Warning: Failed to create GLB with PBR materials: {e}")

mesh_processor/fastmesh.py

@@ -5,6 +5,7 @@ import torch
 from torch import Tensor
 import trimesh
 import numpy as np
+from PIL import Image
 
 from kiui.op import safe_normalize, dot
 from kiui.typing import *
@@ -57,9 +58,9 @@ class FastMesh:
             device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
         if path.endswith(".obj"):
-            mesh = cls._load_obj_fast(path, device=device, **kwargs)
+            mesh = cls._load_obj_fast(path, **{**kwargs, 'device': device})
         elif path.endswith((".glb", ".gltf")):
-            mesh = cls._load_gltf_fast(path, device=device, lazy_texture=lazy_texture, **kwargs)
+            mesh = cls._load_gltf_fast(path, lazy_texture=lazy_texture, **{**kwargs, 'device': device})
         else:
             standard_mesh = Mesh.load(path, resize=False, renormal=False, retex=False, clean=False, **kwargs)
             mesh = cls._from_standard_mesh(standard_mesh)
@@ -829,4 +830,61 @@ class FastMesh:
         )
         temp_mesh.ori_center = self.ori_center
         temp_mesh.ori_scale = self.ori_scale
-        return temp_mesh.write(path)
+        return temp_mesh.write(path)
+    
+    @classmethod
+    def create_glb_with_pbr_materials(cls, obj_path, textures_dict, output_path, device=None):
+        """Create GLB with PBR materials from separate texture files"""
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        try:
+            mesh = cls.load(obj_path, **{'device': device})
+            if mesh is None:
+                raise ValueError(f"Failed to load mesh from {obj_path}")
+            
+            # Combine metallic and roughness if both provided
+            if "metallic" in textures_dict and "roughness" in textures_dict:
+                mr_texture = cls._combine_metallic_roughness(
+                    textures_dict["metallic"], 
+                    textures_dict["roughness"]
+                )
+                mesh.metallicRoughness = torch.from_numpy(mr_texture.astype(np.float32) / 255.0).to(device).contiguous()
+            
+            # Load albedo texture
+            if "albedo" in textures_dict and os.path.exists(textures_dict["albedo"]):
+                albedo_img = np.array(Image.open(textures_dict["albedo"]).convert("RGB"))
+                mesh.albedo = torch.from_numpy(albedo_img.astype(np.float32) / 255.0).to(device).contiguous()
+            
+            # Generate UV if not present
+            if mesh.vt is None:
+                mesh.auto_uv()
+            
+            # Write GLB
+            mesh.write(output_path)
+            print(f"PBR GLB file saved: {output_path}")
+            
+        except Exception as e:
+            print(f"Error creating GLB with PBR materials: {e}")
+            raise e
+    
+    @staticmethod
+    def _combine_metallic_roughness(metallic_path, roughness_path):
+        """Combine metallic and roughness maps into single texture"""
+        metallic_img = Image.open(metallic_path).convert("L")
+        roughness_img = Image.open(roughness_path).convert("L")
+        
+        if metallic_img.size != roughness_img.size:
+            roughness_img = roughness_img.resize(metallic_img.size)
+        
+        width, height = metallic_img.size
+        metallic_array = np.array(metallic_img)
+        roughness_array = np.array(roughness_img)
+        
+        # Create combined array (R, G, B) = (AO, Roughness, Metallic)
+        combined_array = np.zeros((height, width, 3), dtype=np.uint8)
+        combined_array[:, :, 0] = 255  # R: AO (white if no AO map)
+        combined_array[:, :, 1] = roughness_array  # G: Roughness
+        combined_array[:, :, 2] = metallic_array  # B: Metallic
+        
+        return combined_array