update

common.py

@@ -11,6 +11,8 @@ 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
 from PIL import Image
@@ -56,6 +58,8 @@ 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.utils import postprocessing_utils
 from thirdparty.TRELLIS.trellis.utils.render_utils import (
     render_frames,
@@ -79,6 +83,36 @@ DELIGHT = DelightingModel()
 IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
 
 
+def inverse_softplus(x):
+    return x + torch.log(-torch.expm1(-x))
+
+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)
+    return symm
+
+def patched_setup_functions(self):
+    if self.scaling_activation_type == "exp":
+        self.scaling_activation = torch.exp
+        self.inverse_scaling_activation = torch.log
+    elif self.scaling_activation_type == "softplus":
+        self.scaling_activation = F.softplus
+        self.inverse_scaling_activation = inverse_softplus
+
+    self.covariance_activation = build_covariance_from_scaling_rotation
+    self.opacity_activation = torch.sigmoid
+    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.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)
+
+Gaussian.setup_functions = patched_setup_functions
+
+
 def download_kolors_weights() -> None:
     logger.info(f"Download kolors weights from huggingface...")
     subprocess.run(
@@ -323,7 +357,7 @@ def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
     }
 
 
-@spaces.GPU
+# @spaces.GPU
 def unpack_state(state: dict, device: str = "cpu") -> tuple[Gaussian, dict]:
     print("debug11")
     gs = Gaussian(
@@ -333,6 +367,7 @@ def unpack_state(state: dict, device: str = "cpu") -> tuple[Gaussian, dict]:
         scaling_bias=state["gaussian"]["scaling_bias"],
         opacity_bias=state["gaussian"]["opacity_bias"],
         scaling_activation=state["gaussian"]["scaling_activation"],
+        device=device,
     )
     print("debug12")
     gs._xyz = torch.tensor(state["gaussian"]["_xyz"], device=device)
@@ -449,50 +484,6 @@ 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,