update

app.py

@@ -1,6 +1,8 @@
 import os
 
 os.environ["GRADIO_APP"] = "imageto3d"
+from glob import glob
+
 import gradio as gr
 from common import (
     MAX_SEED,
@@ -198,12 +200,9 @@ with gr.Blocks(
             with gr.Row() as single_image_example:
                 examples = gr.Examples(
                     label="Image Gallery",
-                    examples=[
-                        [f"assets/example_image/{image}"]
-                        for image in os.listdir(
-                            "assets/example_image"
-                        )
-                    ],
+                    examples=sorted(
+                        glob("assets/example_image/*")
+                    ),
                     inputs=[image_prompt],
                     fn=preprocess_image_fn,
                     outputs=[image_prompt, raw_image_cache],
@@ -214,12 +213,9 @@ with gr.Blocks(
             with gr.Row(visible=False) as single_sam_image_example:
                 examples = gr.Examples(
                     label="Image Gallery",
-                    examples=[
-                        f"assets/example_image/{image}"
-                        for image in os.listdir(
-                            "assets/example_image"
-                        )
-                    ],
+                    examples=sorted(
+                        glob("assets/example_image/*")
+                    ),
                     inputs=[image_prompt_sam],
                     fn=preprocess_sam_image_fn,
                     outputs=[image_prompt_sam, raw_image_cache],

asset3d_gen/data/backproject_v2.py

@@ -2,10 +2,11 @@ import argparse
 import logging
 import math
 import os
-import spaces
+
 import cv2
 import numpy as np
 import nvdiffrast.torch as dr
+import spaces
 import torch
 import torch.nn.functional as F
 import trimesh
@@ -220,7 +221,7 @@ class TextureBacker:
         bake_angle_thresh: int = 75,
         mask_thresh: float = 0.5,
     ):
-        
+
         self.camera_params = camera_params
         self.renderer = None
         self.view_weights = view_weights
@@ -419,17 +420,17 @@ class TextureBacker:
         return texture
 
     @spaces.GPU
-    def cuda_forward(
+    def compute_texture(
         self,
         colors: list[Image.Image],
         mesh: trimesh.Trimesh,
     ) -> trimesh.Trimesh:
         self._lazy_init_render(self.camera_params, self.mask_thresh)
-        
+
         vertices = torch.from_numpy(mesh.vertices).to(self.device).float()
         faces = torch.from_numpy(mesh.faces).to(self.device).to(torch.int)
         uv_map = torch.from_numpy(mesh.visual.uv).to(self.device).float()
-        
+
         rendered_depth, masks = self.renderer.render_depth(vertices, faces)
         norm_deps = self.renderer.normalize_map_by_mask(rendered_depth, masks)
         render_uvs, _ = self.renderer.render_uv(vertices, faces, uv_map)
@@ -454,7 +455,7 @@ class TextureBacker:
 
         texture_np = texture.cpu().numpy()
         mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)
-        
+
         return texture_np, mask_np
 
     def __call__(
@@ -464,8 +465,8 @@ class TextureBacker:
         output_path: str,
     ) -> trimesh.Trimesh:
         mesh = self.load_mesh(mesh)
-        texture_np, mask_np = self.cuda_forward(colors, mesh)
-        
+        texture_np, mask_np = self.compute_texture(colors, mesh)
+
         texture_np = self.uv_inpaint(mesh, texture_np, mask_np)
         texture_np = post_process_texture(texture_np)
         vertices, faces, uv_map = self.get_mesh_np_attrs(

asset3d_gen/data/mesh_operator.py

@@ -1,9 +1,10 @@
 import logging
 from typing import Tuple, Union
-import spaces
+
 import igraph
 import numpy as np
 import pyvista as pv
+import spaces
 import torch
 import utils3d
 from pymeshfix import _meshfix

asset3d_gen/models/delight_model.py

@@ -1,8 +1,9 @@
 import os
 from typing import Union
-import spaces
+
 import cv2
 import numpy as np
+import spaces
 import torch
 from diffusers import (
     EulerAncestralDiscreteScheduler,
@@ -28,7 +29,6 @@ class DelightingModel(object):
         device: str = "cuda",
         seed: int = 0,
     ) -> None:
-        
         self.image_guide_scale = image_guide_scale
         self.text_guide_scale = text_guide_scale
         self.num_infer_step = num_infer_step
@@ -38,8 +38,8 @@ class DelightingModel(object):
         )
         self.seed = seed
         self.device = device
-        self.pipeline = None # lazy load model adapt to @spaces.GPU
-        
+        self.pipeline = None  # lazy load model adapt to @spaces.GPU
+
         if model_path is None:
             suffix = "hunyuan3d-delight-v2-0"
             model_path = snapshot_download(
@@ -114,7 +114,7 @@ class DelightingModel(object):
         target_wh: tuple[int, int] = None,
     ) -> Image.Image:
         self._lazy_init_pipeline()
-        
+
         if isinstance(image, str):
             image = Image.open(image)
         elif isinstance(image, np.ndarray):
@@ -158,7 +158,7 @@ if __name__ == "__main__":
     delighting_model = DelightingModel(
         # model_path="/horizon-bucket/robot_lab/users/xinjie.wang/weights/hunyuan3d-delight-v2-0"  # noqa
     )
-    image_path = "/home/users/xinjie.wang/xinjie/asset3d-gen/scripts/apps/assets/example_image/room_bottle_002.jpeg"
+    image_path = "scripts/apps/assets/example_image/room_bottle_002.jpeg"
     image = delighting_model(
         image_path, preprocess=True, target_wh=(512, 512)
     )  # noqa

asset3d_gen/models/gs_model.py

@@ -428,7 +428,7 @@ class GaussianOperator(GaussianBase):
             sh_degree=self.sh_degree,
             device=self.device,
         )
-        print("self.device", self.device)
+
         return GaussianOperator(**gs_dict)
 
     def rescale(self, scale: float):
@@ -459,7 +459,9 @@ class GaussianOperator(GaussianBase):
         sh_degree: int = 0,
         device: str = "cuda",
     ) -> None:
-        gs_model = GaussianOperator.load_from_ply(in_ply, sh_degree, device=device)
+        gs_model = GaussianOperator.load_from_ply(
+            in_ply, sh_degree, device=device
+        )
 
         if instance_pose is not None:
             gs_model = gs_model.get_gaussians(instance_pose=instance_pose)

asset3d_gen/models/sr_model.py

@@ -1,8 +1,9 @@
 import logging
 import os
 from typing import Union
-import spaces
+
 import numpy as np
+import spaces
 import torch
 from huggingface_hub import snapshot_download
 from PIL import Image
@@ -66,11 +67,16 @@ class ImageRealESRGAN:
         if version.parse(torchvision.__version__) > version.parse("0.16"):
             import sys
             import types
+
             import torchvision.transforms.functional as TF
 
-            functional_tensor = types.ModuleType("torchvision.transforms.functional_tensor")
+            functional_tensor = types.ModuleType(
+                "torchvision.transforms.functional_tensor"
+            )
             functional_tensor.rgb_to_grayscale = TF.rgb_to_grayscale
-            sys.modules["torchvision.transforms.functional_tensor"] = functional_tensor
+            sys.modules["torchvision.transforms.functional_tensor"] = (
+                functional_tensor
+            )
 
         self.outscale = outscale
         self.upsampler = None
@@ -80,15 +86,16 @@ class ImageRealESRGAN:
             model_path = snapshot_download(
                 repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
             )
-            model_path = os.path.join(model_path, suffix, "RealESRGAN_x4plus.pth")
+            model_path = os.path.join(
+                model_path, suffix, "RealESRGAN_x4plus.pth"
+            )
 
         self.model_path = model_path
-        
+
     def _lazy_init(self):
         if self.upsampler is None:
             from basicsr.archs.rrdbnet_arch import RRDBNet
             from realesrgan import RealESRGANer
-            from huggingface_hub import snapshot_download
 
             model = RRDBNet(
                 num_in_ch=3,
@@ -110,7 +117,7 @@ class ImageRealESRGAN:
     @spaces.GPU
     def __call__(self, image: Union[Image.Image, np.ndarray]) -> Image.Image:
         self._lazy_init()
-        
+
         if isinstance(image, Image.Image):
             image = np.array(image)
 

asset3d_gen/scripts/render_gs.py

@@ -2,9 +2,10 @@ import argparse
 import logging
 import math
 import os
-import spaces
+
 import cv2
 import numpy as np
+import spaces
 import torch
 from tqdm import tqdm
 from asset3d_gen.data.utils import (

asset3d_gen/utils/process_media.py

@@ -6,11 +6,12 @@ import subprocess
 from glob import glob
 from io import BytesIO
 from typing import Union
-import spaces
+
 import cv2
 import imageio
 import numpy as np
 import PIL.Image as Image
+import spaces
 from moviepy.editor import VideoFileClip, clips_array
 
 logging.basicConfig(level=logging.INFO)

asset3d_gen/validators/aesthetic_predictor.py

@@ -1,8 +1,9 @@
 import os
-import spaces
+
 import clip
 import numpy as np
 import pytorch_lightning as pl
+import spaces
 import torch
 import torch.nn as nn
 from huggingface_hub import snapshot_download

common.py

@@ -11,7 +11,6 @@ import gradio as gr
 import numpy as np
 import spaces
 import torch
-import torch
 import torch.nn.functional as F
 import trimesh
 from easydict import EasyDict as edict
@@ -58,8 +57,11 @@ from thirdparty.TRELLIS.trellis.representations import (
     Gaussian,
     MeshExtractResult,
 )
-from thirdparty.TRELLIS.trellis.representations.gaussian.general_utils import inverse_sigmoid, strip_symmetric, build_scaling_rotation
-
+from thirdparty.TRELLIS.trellis.representations.gaussian.general_utils import (
+    build_scaling_rotation,
+    inverse_sigmoid,
+    strip_symmetric,
+)
 from thirdparty.TRELLIS.trellis.utils import postprocessing_utils
 from thirdparty.TRELLIS.trellis.utils.render_utils import (
     render_frames,
@@ -87,7 +89,9 @@ def patched_setup_functions(self):
     def inverse_softplus(x):
         return x + torch.log(-torch.expm1(-x))
 
-    def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):
+    def build_covariance_from_scaling_rotation(
+        scaling, scaling_modifier, rotation
+    ):
         L = build_scaling_rotation(scaling_modifier * scaling, rotation)
         actual_covariance = L @ L.transpose(1, 2)
         symm = strip_symmetric(actual_covariance)
@@ -105,10 +109,15 @@ def patched_setup_functions(self):
     self.inverse_opacity_activation = inverse_sigmoid
     self.rotation_activation = F.normalize
 
-    self.scale_bias = self.inverse_scaling_activation(torch.tensor(self.scaling_bias)).to(self.device)
+    self.scale_bias = self.inverse_scaling_activation(
+        torch.tensor(self.scaling_bias)
+    ).to(self.device)
     self.rots_bias = torch.zeros((4)).to(self.device)
     self.rots_bias[0] = 1
-    self.opacity_bias = self.inverse_opacity_activation(torch.tensor(self.opacity_bias)).to(self.device)
+    self.opacity_bias = self.inverse_opacity_activation(
+        torch.tensor(self.opacity_bias)
+    ).to(self.device)
+
 
 Gaussian.setup_functions = patched_setup_functions
 
@@ -357,9 +366,7 @@ def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
     }
 
 
-# @spaces.GPU
 def unpack_state(state: dict, device: str = "cpu") -> tuple[Gaussian, dict]:
-    print("debug11")
     gs = Gaussian(
         aabb=state["gaussian"]["aabb"],
         sh_degree=state["gaussian"]["sh_degree"],
@@ -369,7 +376,6 @@ def unpack_state(state: dict, device: str = "cpu") -> tuple[Gaussian, dict]:
         scaling_activation=state["gaussian"]["scaling_activation"],
         device=device,
     )
-    print("debug12")
     gs._xyz = torch.tensor(state["gaussian"]["_xyz"], device=device)
     gs._features_dc = torch.tensor(
         state["gaussian"]["_features_dc"], device=device
@@ -382,7 +388,6 @@ def unpack_state(state: dict, device: str = "cpu") -> tuple[Gaussian, dict]:
         vertices=torch.tensor(state["mesh"]["vertices"], device=device),
         faces=torch.tensor(state["mesh"]["faces"], device=device),
     )
-    print("debug13")
 
     return gs, mesh
 
@@ -484,21 +489,62 @@ def image_to_3d(
     return state, video_path
 
 
+@spaces.GPU
+def extract_3d_representations(
+    state: dict, enable_delight: bool, req: gr.Request
+):
+    output_root = TMP_DIR
+    output_root = os.path.join(output_root, str(req.session_hash))
+    gs_model, mesh_model = unpack_state(state)
+
+    mesh = postprocessing_utils.to_glb(
+        gs_model,
+        mesh_model,
+        simplify=0.9,
+        texture_size=1024,
+        verbose=True,
+    )
+    filename = "sample"
+    gs_path = os.path.join(output_root, f"{filename}_gs.ply")
+    gs_model.save_ply(gs_path)
+
+    # Rotate mesh and GS by 90 degrees around Z-axis.
+    rot_matrix = [[0, 0, -1], [0, 1, 0], [1, 0, 0]]
+    # Addtional rotation for GS to align mesh.
+    gs_rot = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]]) @ np.array(
+        rot_matrix
+    )
+    pose = GaussianOperator.trans_to_quatpose(gs_rot)
+    aligned_gs_path = gs_path.replace(".ply", "_aligned.ply")
+    GaussianOperator.resave_ply(
+        in_ply=gs_path,
+        out_ply=aligned_gs_path,
+        instance_pose=pose,
+    )
+
+    mesh.vertices = mesh.vertices @ np.array(rot_matrix)
+    mesh_obj_path = os.path.join(output_root, f"{filename}.obj")
+    mesh.export(mesh_obj_path)
+    mesh_glb_path = os.path.join(output_root, f"{filename}.glb")
+    mesh.export(mesh_glb_path)
+
+    torch.cuda.empty_cache()
+
+    return mesh_glb_path, gs_path, mesh_obj_path, aligned_gs_path
+
+
 def extract_3d_representations_v2(
     state: dict,
     enable_delight: bool,
     req: gr.Request,
 ):
-    print("debug1")
     output_root = TMP_DIR
     user_dir = os.path.join(output_root, str(req.session_hash))
     gs_model, mesh_model = unpack_state(state, device="cpu")
-    print("debug2")
 
     filename = "sample"
     gs_path = os.path.join(user_dir, f"{filename}_gs.ply")
     gs_model.save_ply(gs_path)
-    print("debug3")
 
     # Rotate mesh and GS by 90 degrees around Z-axis.
     rot_matrix = [[0, 0, -1], [0, 1, 0], [1, 0, 0]]
@@ -509,17 +555,14 @@ def extract_3d_representations_v2(
     gs_rot = np.array(gs_add_rot) @ np.array(rot_matrix)
     pose = GaussianOperator.trans_to_quatpose(gs_rot)
     aligned_gs_path = gs_path.replace(".ply", "_aligned.ply")
-    print("debug4")
     GaussianOperator.resave_ply(
         in_ply=gs_path,
         out_ply=aligned_gs_path,
         instance_pose=pose,
         device="cpu",
     )
-    print("debug5")
     color_path = os.path.join(user_dir, "color.png")
     render_gs_api(aligned_gs_path, color_path)
-    print("debug6")
 
     mesh = trimesh.Trimesh(
         vertices=mesh_model.vertices.cpu().numpy(),
@@ -530,7 +573,7 @@ def extract_3d_representations_v2(
 
     mesh_obj_path = os.path.join(user_dir, f"{filename}.obj")
     mesh.export(mesh_obj_path)
-    print("debug7")
+
     mesh = backproject_api(
         delight_model=DELIGHT,
         imagesr_model=IMAGESR_MODEL,
@@ -540,7 +583,6 @@ def extract_3d_representations_v2(
         skip_fix_mesh=False,
         delight=enable_delight,
     )
-    print("debug8")
 
     mesh_glb_path = os.path.join(user_dir, f"{filename}.glb")
     mesh.export(mesh_glb_path)
@@ -560,7 +602,7 @@ def extract_urdf(
     output_root = TMP_DIR
     if req is not None:
         output_root = os.path.join(output_root, str(req.session_hash))
-    
+
     # Convert to URDF and recover attrs by GPT.
     filename = "sample"
     urdf_convertor = URDFGenerator(GPT_CLIENT, render_view_num=4)