Merge branch 'main' of https://huggingface.co/spaces/Stable-X/Hi3DGen

README.md

@@ -1,5 +1,5 @@
 ---
-title: TRELLIS
+title: Hi3DGen
 emoji: 🏢
 colorFrom: indigo
 colorTo: blue
@@ -8,9 +8,9 @@ sdk_version: 4.44.1
 app_file: app.py
 pinned: false
 license: mit
-short_description: Scalable and Versatile 3D Generation from images
+short_description: High-fidelity 3D Geometry Generation from images
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
-Paper: https://huggingface.co/papers/2412.01506
+Paper: 

requirements.txt

@@ -1,9 +1,10 @@
 --extra-index-url https://download.pytorch.org/whl/cu121
-#huggingface_hub==0.25.0
-#diffusers==0.28.0
-#accelerate==1.2.1
-#kornia==0.8.0
-#timm==0.6.7
+huggingface_hub==0.25.0
+diffusers==0.28.0
+accelerate==1.2.1
+kornia==0.8.0
+timm==0.6.7
+
 torch==2.4.0
 torchvision==0.19.0
 pillow==10.4.0

trellis/renderers/mesh_renderer.py

@@ -64,7 +64,7 @@ class MeshRenderer:
             mesh : MeshExtractResult,
             extrinsics: torch.Tensor,
             intrinsics: torch.Tensor,
-            return_types = ["mask", "normal", "depth"]
+            return_types = ["mask", "normal", "depth", "color"]
         ) -> edict:
         """
         Render the mesh.

trellis/representations/mesh/cube2mesh.py

@@ -10,7 +10,8 @@ from ...modules.sparse import SparseTensor
 from easydict import EasyDict as edict
 from .utils_cube import *
 from .flexicube import FlexiCubes
-
+import numpy as np
+import trimesh
 
 class MeshExtractResult:
     def __init__(self,

trellis/utils/render_utils.py

@@ -83,18 +83,28 @@ def render_frames(sample, extrinsics, intrinsics, options={}, colors_overwrite=N
         else:
             res = renderer.render(sample, extr, intr)
             if 'normal' not in rets: rets['normal'] = []
+            if 'color' not in rets: rets['color'] = []
+            rets['color'].append(np.clip(res['color'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8))
             rets['normal'].append(np.clip(res['normal'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8))
     return rets
 
-
-def render_video(sample, resolution=512, bg_color=(0, 0, 0), num_frames=300, r=2, fov=40, **kwargs):
-    yaws = torch.linspace(0, 2 * 3.1415, num_frames)
-    pitch = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames))
+def render_video(sample, resolution=512, ssaa=4, bg_color=(0, 0, 0), num_frames=300, r=2, fov=40, 
+                 inverse_direction=False, pitch=-1, device="cuda", **kwargs):
+    if inverse_direction:
+        yaws = torch.linspace(3.1415, -3.1415, num_frames)
+    else:
+        yaws = torch.linspace(0, 2 * 3.1415, num_frames)
+    if pitch != -1:
+        pitch = pitch * torch.ones(num_frames)
+    else:
+        pitch = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames))
     yaws = yaws.tolist()
     pitch = pitch.tolist()
     extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitch, r, fov)
-    return render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs)
-
+    
+    res = render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color, 'ssaa': ssaa}, device=device, **kwargs)
+    res.update({'extrinsics': extrinsics, 'intrinsics': intrinsics})
+    return res
 
 def render_multiview(sample, resolution=512, nviews=30):
     r = 2