update

app.py

@@ -267,7 +267,7 @@ with gr.Blocks(delete_cache=(43200, 43200), theme=custom_theme) as demo:
 
     demo.load(start_session)
     demo.unload(end_session)
-
+    no_mesh_post_process = gr.State(True)
     mesh_input.change(
         lambda: tuple(
             [
@@ -368,6 +368,7 @@ with gr.Blocks(delete_cache=(43200, 43200), theme=custom_theme) as demo:
             texture_size,
             project_delight,
             fix_mesh,
+            no_mesh_post_process,
         ],
         outputs=[mesh_output, mesh_outpath, download_btn],
     ).success(

app_style.py

@@ -1,10 +1,26 @@
+# Project EmbodiedGen
+#
+# Copyright (c) 2025 Horizon Robotics. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied. See the License for the specific language governing
+# permissions and limitations under the License.
+
 from gradio.themes import Soft
 from gradio.themes.utils.colors import gray, neutral, slate, stone, teal, zinc
 
 lighting_css = """
 <style>
 #lighter_mesh canvas {
-    filter: brightness(2.0) !important;
+    filter: brightness(2.3) !important;
 }
 </style>
 """

common.py

@@ -53,6 +53,7 @@ from embodied_gen.scripts.text2image import (
 from embodied_gen.utils.gpt_clients import GPT_CLIENT
 from embodied_gen.utils.process_media import (
     filter_image_small_connected_components,
+    keep_largest_connected_component,
     merge_images_video,
 )
 from embodied_gen.utils.tags import VERSION
@@ -151,6 +152,21 @@ if os.getenv("GRADIO_APP") == "imageto3d":
         os.path.dirname(os.path.abspath(__file__)), "sessions/imageto3d"
     )
     os.makedirs(TMP_DIR, exist_ok=True)
+elif os.getenv("GRADIO_APP") == "imageto3d_sam3d":
+    from embodied_gen.models.sam3d import Sam3dInference
+
+    RBG_REMOVER = RembgRemover()
+    RBG14_REMOVER = BMGG14Remover()
+    SAM_PREDICTOR = SAMPredictor(model_type="vit_h", device="cpu")
+    PIPELINE = Sam3dInference()
+    SEG_CHECKER = ImageSegChecker(GPT_CLIENT)
+    GEO_CHECKER = MeshGeoChecker(GPT_CLIENT)
+    AESTHETIC_CHECKER = ImageAestheticChecker()
+    CHECKERS = [GEO_CHECKER, SEG_CHECKER, AESTHETIC_CHECKER]
+    TMP_DIR = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "sessions/imageto3d"
+    )
+    os.makedirs(TMP_DIR, exist_ok=True)
 elif os.getenv("GRADIO_APP") == "textto3d":
     RBG_REMOVER = RembgRemover()
     RBG14_REMOVER = BMGG14Remover()
@@ -169,6 +185,23 @@ elif os.getenv("GRADIO_APP") == "textto3d":
         os.path.dirname(os.path.abspath(__file__)), "sessions/textto3d"
     )
     os.makedirs(TMP_DIR, exist_ok=True)
+elif os.getenv("GRADIO_APP") == "textto3d_sam3d":
+    from embodied_gen.models.sam3d import Sam3dInference
+
+    RBG_REMOVER = RembgRemover()
+    RBG14_REMOVER = BMGG14Remover()
+    PIPELINE = Sam3dInference()
+    text_model_dir = "weights/Kolors"
+    PIPELINE_IMG_IP = build_text2img_ip_pipeline(text_model_dir, ref_scale=0.3)
+    PIPELINE_IMG = build_text2img_pipeline(text_model_dir)
+    SEG_CHECKER = ImageSegChecker(GPT_CLIENT)
+    GEO_CHECKER = MeshGeoChecker(GPT_CLIENT)
+    AESTHETIC_CHECKER = ImageAestheticChecker()
+    CHECKERS = [GEO_CHECKER, SEG_CHECKER, AESTHETIC_CHECKER]
+    TMP_DIR = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "sessions/textto3d"
+    )
+    os.makedirs(TMP_DIR, exist_ok=True)
 elif os.getenv("GRADIO_APP") == "texture_edit":
     DELIGHT = DelightingModel()
     IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
@@ -201,18 +234,23 @@ def end_session(req: gr.Request) -> None:
 
 @spaces.GPU
 def preprocess_image_fn(
-    image: str | np.ndarray | Image.Image, rmbg_tag: str = "rembg"
+    image: str | np.ndarray | Image.Image,
+    rmbg_tag: str = "rembg",
+    preprocess: bool = True,
 ) -> tuple[Image.Image, Image.Image]:
     if isinstance(image, str):
         image = Image.open(image)
     elif isinstance(image, np.ndarray):
         image = Image.fromarray(image)
 
-    image_cache = image.copy().resize((512, 512))
+    image_cache = image.copy()  # resize_pil(image.copy(), 1024)
 
     bg_remover = RBG_REMOVER if rmbg_tag == "rembg" else RBG14_REMOVER
     image = bg_remover(image)
-    image = trellis_preprocess(image)
+    image = keep_largest_connected_component(image)
+
+    if preprocess:
+        image = trellis_preprocess(image)
 
     return image, image_cache
 
@@ -224,7 +262,7 @@ def preprocess_sam_image_fn(
         image = Image.fromarray(image)
 
     sam_image = SAM_PREDICTOR.preprocess_image(image)
-    image_cache = Image.fromarray(sam_image).resize((512, 512))
+    image_cache = sam_image.copy()
     SAM_PREDICTOR.predictor.set_image(sam_image)
 
     return sam_image, image_cache
@@ -349,11 +387,11 @@ def select_point(
 def image_to_3d(
     image: Image.Image,
     seed: int,
-    ss_guidance_strength: float,
     ss_sampling_steps: int,
-    slat_guidance_strength: float,
     slat_sampling_steps: int,
     raw_image_cache: Image.Image,
+    ss_guidance_strength: float,
+    slat_guidance_strength: float,
     sam_image: Image.Image = None,
     is_sam_image: bool = False,
     req: gr.Request = None,
@@ -392,8 +430,56 @@ def image_to_3d(
 
     gs_model = outputs["gaussian"][0]
     mesh_model = outputs["mesh"][0]
-    color_images = render_video(gs_model)["color"]
-    normal_images = render_video(mesh_model)["normal"]
+    color_images = render_video(gs_model, r=1.85)["color"]
+    normal_images = render_video(mesh_model, r=1.85)["normal"]
+
+    video_path = os.path.join(output_root, "gs_mesh.mp4")
+    merge_images_video(color_images, normal_images, video_path)
+    state = pack_state(gs_model, mesh_model)
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    return state, video_path
+
+
+@spaces.GPU
+def image_to_3d_sam3d(
+    image: Image.Image,
+    seed: int,
+    ss_sampling_steps: int,
+    slat_sampling_steps: int,
+    raw_image_cache: Image.Image,
+    ss_guidance_strength: float = None,
+    slat_guidance_strength: float = None,
+    sam_image: Image.Image = None,
+    is_sam_image: bool = False,
+    req: gr.Request = None,
+) -> tuple[dict, str]:
+    if is_sam_image:
+        seg_image = filter_image_small_connected_components(sam_image)
+        seg_image = Image.fromarray(seg_image, mode="RGBA")
+    else:
+        seg_image = image
+
+    if isinstance(seg_image, np.ndarray):
+        seg_image = Image.fromarray(seg_image)
+
+    output_root = os.path.join(TMP_DIR, str(req.session_hash))
+    os.makedirs(output_root, exist_ok=True)
+    seg_image.save(f"{output_root}/seg_image.png")
+    raw_image_cache.save(f"{output_root}/raw_image.png")
+    outputs = PIPELINE.run(
+        seg_image,
+        seed=seed,
+        stage1_inference_steps=ss_sampling_steps,
+        stage2_inference_steps=slat_sampling_steps,
+    )
+
+    gs_model = outputs["gaussian"][0]
+    mesh_model = outputs["mesh"][0]
+    color_images = render_video(gs_model, r=1.85)["color"]
+    normal_images = render_video(mesh_model, r=1.85)["normal"]
 
     video_path = os.path.join(output_root, "gs_mesh.mp4")
     merge_images_video(color_images, normal_images, video_path)
@@ -688,6 +774,7 @@ def text2image_fn(
     image_wh: int | tuple[int, int] = [1024, 1024],
     rmbg_tag: str = "rembg",
     seed: int = None,
+    enable_pre_resize: bool = True,
     n_sample: int = 3,
     req: gr.Request = None,
 ):
@@ -715,7 +802,9 @@ def text2image_fn(
 
     for idx in range(len(images)):
         image = images[idx]
-        images[idx], _ = preprocess_image_fn(image, rmbg_tag)
+        images[idx], _ = preprocess_image_fn(
+            image, rmbg_tag, enable_pre_resize
+        )
 
     save_paths = []
     for idx, image in enumerate(images):
@@ -841,6 +930,7 @@ def backproject_texture_v2(
     texture_size: int,
     enable_delight: bool = True,
     fix_mesh: bool = False,
+    no_mesh_post_process: bool = False,
     uuid: str = "sample",
     req: gr.Request = None,
 ) -> str:
@@ -857,6 +947,7 @@ def backproject_texture_v2(
         skip_fix_mesh=not fix_mesh,
         delight=enable_delight,
         texture_wh=[texture_size, texture_size],
+        no_mesh_post_process=no_mesh_post_process,
     )
 
     output_obj_mesh = os.path.join(output_dir, f"{uuid}.obj")

embodied_gen/data/backproject_v2.py

@@ -274,6 +274,7 @@ class TextureBacker:
         mask_thresh (float, optional): Threshold for visibility masks.
         smooth_texture (bool, optional): Apply post-processing to texture.
         inpaint_smooth (bool, optional): Apply inpainting smoothing.
+        mesh_post_process (bool, optional): False for preventing modification of vertices.
 
     Example:
         ```py
@@ -308,6 +309,7 @@ class TextureBacker:
         mask_thresh: float = 0.5,
         smooth_texture: bool = True,
         inpaint_smooth: bool = False,
+        mesh_post_process: bool = True,
     ) -> None:
         self.camera_params = camera_params
         self.renderer = None
@@ -318,6 +320,7 @@ class TextureBacker:
         self.mask_thresh = mask_thresh
         self.smooth_texture = smooth_texture
         self.inpaint_smooth = inpaint_smooth
+        self.mesh_post_process = mesh_post_process
 
         self.bake_angle_thresh = bake_angle_thresh
         self.bake_unreliable_kernel_size = int(
@@ -668,7 +671,12 @@ class TextureBacker:
             mesh, self.scale, self.center
         )
         textured_mesh = save_mesh_with_mtl(
-            vertices, faces, uv_map, texture_np, output_path
+            vertices,
+            faces,
+            uv_map,
+            texture_np,
+            output_path,
+            mesh_process=self.mesh_post_process,
         )
 
         return textured_mesh
@@ -766,6 +774,7 @@ def parse_args():
         help="Disable saving delight image",
     )
     parser.add_argument("--n_max_faces", type=int, default=30000)
+    parser.add_argument("--no_mesh_post_process", action="store_true")
     args, unknown = parser.parse_known_args()
 
     return args
@@ -856,6 +865,7 @@ def entrypoint(
         render_wh=args.resolution_hw,
         texture_wh=args.texture_wh,
         smooth_texture=not args.no_smooth_texture,
+        mesh_post_process=not args.no_mesh_post_process,
     )
 
     textured_mesh = texture_backer(multiviews, mesh, args.output_path)

embodied_gen/data/backproject_v3.py

@@ -14,10 +14,10 @@
 # implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
-
 import argparse
 import logging
 import math
+import os
 from typing import Literal, Union
 
 import cv2
@@ -353,8 +353,8 @@ def parse_args():
     parser.add_argument(
         "--distance",
         type=float,
-        default=5,
-        help="Camera distance (default: 5)",
+        default=4.5,
+        help="Camera distance (default: 4.5)",
     )
     parser.add_argument(
         "--resolution_hw",
@@ -400,8 +400,8 @@ def parse_args():
     parser.add_argument(
         "--mesh_sipmlify_ratio",
         type=float,
-        default=0.9,
-        help="Mesh simplification ratio (default: 0.9)",
+        default=0.85,
+        help="Mesh simplification ratio (default: 0.85)",
     )
     parser.add_argument(
         "--delight", action="store_true", help="Use delighting model."
@@ -425,6 +425,7 @@ def parse_args():
     return args
 
 
+@spaces.GPU
 def entrypoint(
     delight_model: DelightingModel = None,
     imagesr_model: ImageRealESRGAN = None,
@@ -499,7 +500,7 @@ def entrypoint(
     faces = mesh.faces.astype(np.int32)
     vertices = vertices.astype(np.float32)
 
-    if not args.skip_fix_mesh and len(faces) > 10 * args.n_max_faces:
+    if not args.skip_fix_mesh:
         mesh_fixer = MeshFixer(vertices, faces, args.device)
         vertices, faces = mesh_fixer(
             filter_ratio=args.mesh_sipmlify_ratio,
@@ -511,7 +512,7 @@ def entrypoint(
         if len(faces) > args.n_max_faces:
             mesh_fixer = MeshFixer(vertices, faces, args.device)
             vertices, faces = mesh_fixer(
-                filter_ratio=max(0.05, args.mesh_sipmlify_ratio - 0.2),
+                filter_ratio=max(0.1, args.mesh_sipmlify_ratio - 0.1),
                 max_hole_size=0.04,
                 resolution=1024,
                 num_views=1000,

embodied_gen/data/utils.py

@@ -15,10 +15,13 @@
 # permissions and limitations under the License.
 
 
+import logging
 import math
 import os
-import random
+import time
 import zipfile
+from contextlib import contextmanager
+from dataclasses import dataclass, field
 from shutil import rmtree
 from typing import List, Tuple, Union
 
@@ -28,20 +31,9 @@ import numpy as np
 import nvdiffrast.torch as dr
 import torch
 import torch.nn.functional as F
-from PIL import Image, ImageEnhance
-
-try:
-    from kolors.models.modeling_chatglm import ChatGLMModel
-    from kolors.models.tokenization_chatglm import ChatGLMTokenizer
-except ImportError:
-    ChatGLMTokenizer = None
-    ChatGLMModel = None
-import logging
-from dataclasses import dataclass, field
-
 import trimesh
 from kaolin.render.camera import Camera
-from torch import nn
+from PIL import Image, ImageEnhance
 
 logger = logging.getLogger(__name__)
 
@@ -50,10 +42,8 @@ __all__ = [
     "DiffrastRender",
     "save_images",
     "render_pbr",
-    "prelabel_text_feature",
     "calc_vertex_normals",
     "normalize_vertices_array",
-    "load_mesh_to_unit_cube",
     "as_list",
     "CameraSetting",
     "import_kaolin_mesh",
@@ -67,6 +57,7 @@ __all__ = [
     "trellis_preprocess",
     "delete_dir",
     "kaolin_to_opencv_view",
+    "model_device_ctx",
 ]
 
 
@@ -520,114 +511,6 @@ def render_pbr(
     return image, albedo, diffuse, normal
 
 
-def _move_to_target_device(data, device: str):
-    if isinstance(data, dict):
-        for key, value in data.items():
-            data[key] = _move_to_target_device(value, device)
-    elif isinstance(data, torch.Tensor):
-        return data.to(device)
-
-    return data
-
-
-def _encode_prompt(
-    prompt_batch,
-    text_encoders,
-    tokenizers,
-    proportion_empty_prompts=0,
-    is_train=True,
-):
-    prompt_embeds_list = []
-
-    captions = []
-    for caption in prompt_batch:
-        if random.random() < proportion_empty_prompts:
-            captions.append("")
-        elif isinstance(caption, str):
-            captions.append(caption)
-        elif isinstance(caption, (list, np.ndarray)):
-            captions.append(random.choice(caption) if is_train else caption[0])
-
-    with torch.no_grad():
-        for tokenizer, text_encoder in zip(tokenizers, text_encoders):
-            text_inputs = tokenizer(
-                captions,
-                padding="max_length",
-                max_length=256,
-                truncation=True,
-                return_tensors="pt",
-            ).to(text_encoder.device)
-
-            output = text_encoder(
-                input_ids=text_inputs.input_ids,
-                attention_mask=text_inputs.attention_mask,
-                position_ids=text_inputs.position_ids,
-                output_hidden_states=True,
-            )
-
-            # We are only interested in the pooled output of the text encoder.
-            prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone()
-            pooled_prompt_embeds = output.hidden_states[-1][-1, :, :].clone()
-            bs_embed, seq_len, _ = prompt_embeds.shape
-            prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
-            prompt_embeds_list.append(prompt_embeds)
-
-    prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
-    pooled_prompt_embeds = pooled_prompt_embeds.view(bs_embed, -1)
-
-    return prompt_embeds, pooled_prompt_embeds
-
-
-def load_llm_models(pretrained_model_name_or_path: str, device: str):
-    tokenizer = ChatGLMTokenizer.from_pretrained(
-        pretrained_model_name_or_path,
-        subfolder="text_encoder",
-    )
-    text_encoder = ChatGLMModel.from_pretrained(
-        pretrained_model_name_or_path,
-        subfolder="text_encoder",
-    ).to(device)
-
-    text_encoders = [
-        text_encoder,
-    ]
-    tokenizers = [
-        tokenizer,
-    ]
-
-    logger.info(f"Load model from {pretrained_model_name_or_path} done.")
-
-    return tokenizers, text_encoders
-
-
-def prelabel_text_feature(
-    prompt_batch: List[str],
-    output_dir: str,
-    tokenizers: nn.Module,
-    text_encoders: nn.Module,
-) -> List[str]:
-    os.makedirs(output_dir, exist_ok=True)
-
-    # prompt_batch ["text..."]
-    prompt_embeds, pooled_prompt_embeds = _encode_prompt(
-        prompt_batch, text_encoders, tokenizers
-    )
-
-    prompt_embeds = _move_to_target_device(prompt_embeds, device="cpu")
-    pooled_prompt_embeds = _move_to_target_device(
-        pooled_prompt_embeds, device="cpu"
-    )
-
-    data_dict = dict(
-        prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds
-    )
-
-    save_path = os.path.join(output_dir, "text_feat.pth")
-    torch.save(data_dict, save_path)
-
-    return save_path
-
-
 def _calc_face_normals(
     vertices: torch.Tensor,  # V,3 first vertex may be unreferenced
     faces: torch.Tensor,  # F,3 long, first face may be all zero
@@ -683,25 +566,6 @@ def normalize_vertices_array(
     return vertices, scale, center
 
 
-def load_mesh_to_unit_cube(
-    mesh_file: str,
-    mesh_scale: float = 1.0,
-) -> tuple[trimesh.Trimesh, float, list[float]]:
-    if not os.path.exists(mesh_file):
-        raise FileNotFoundError(f"mesh_file path {mesh_file} not exists.")
-
-    mesh = trimesh.load(mesh_file)
-    if isinstance(mesh, trimesh.Scene):
-        mesh = trimesh.utils.concatenate(mesh)
-
-    vertices, scale, center = normalize_vertices_array(
-        mesh.vertices, mesh_scale
-    )
-    mesh.vertices = vertices
-
-    return mesh, scale, center
-
-
 def as_list(obj):
     if isinstance(obj, (list, tuple)):
         return obj
@@ -862,6 +726,7 @@ def save_mesh_with_mtl(
     texture: Union[Image.Image, np.ndarray],
     output_path: str,
     material_base=(250, 250, 250, 255),
+    mesh_process: bool = True,
 ) -> trimesh.Trimesh:
     if isinstance(texture, np.ndarray):
         texture = Image.fromarray(texture)
@@ -870,6 +735,7 @@ def save_mesh_with_mtl(
         vertices,
         faces,
         visual=trimesh.visual.TextureVisuals(uv=uvs, image=texture),
+        process=mesh_process,  # True for preventing modification of vertices
     )
     mesh.visual.material = trimesh.visual.material.SimpleMaterial(
         image=texture,
@@ -998,8 +864,9 @@ def gamma_shs(shs: torch.Tensor, gamma: float) -> torch.Tensor:
 
 
 def resize_pil(image: Image.Image, max_size: int = 1024) -> Image.Image:
-    max_size = max(image.size)
-    scale = min(1, 1024 / max_size)
+    current_max_dim = max(image.size)
+    scale = min(1, max_size / current_max_dim)
+
     if scale < 1:
         new_size = (int(image.width * scale), int(image.height * scale))
         image = image.resize(new_size, Image.Resampling.LANCZOS)
@@ -1068,3 +935,34 @@ def delete_dir(folder_path: str, keep_subs: list[str] = None) -> None:
             rmtree(item_path)
         else:
             os.remove(item_path)
+
+
+@contextmanager
+def model_device_ctx(
+    *models,
+    src_device: str = "cpu",
+    dst_device: str = "cuda",
+    verbose: bool = False,
+):
+    start = time.perf_counter()
+    for m in models:
+        if m is None:
+            continue
+        m.to(dst_device)
+    to_cuda_time = time.perf_counter() - start
+
+    try:
+        yield
+    finally:
+        start = time.perf_counter()
+        for m in models:
+            if m is None:
+                continue
+            m.to(src_device)
+        to_cpu_time = time.perf_counter() - start
+
+        if verbose:
+            model_names = [m.__class__.__name__ for m in models]
+            logger.debug(
+                f"[model_device_ctx] {model_names} to cuda: {to_cuda_time:.1f}s, to cpu: {to_cpu_time:.1f}s"
+            )

embodied_gen/models/sam3d.py

@@ -0,0 +1,147 @@
+# Project EmbodiedGen
+#
+# Copyright (c) 2025 Horizon Robotics. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied. See the License for the specific language governing
+# permissions and limitations under the License.
+
+from embodied_gen.utils.monkey_patches import monkey_patch_sam3d
+
+monkey_patch_sam3d()
+import os
+import sys
+from typing import Optional, Union
+
+import numpy as np
+from hydra.utils import instantiate
+from modelscope import snapshot_download
+from omegaconf import OmegaConf
+from PIL import Image
+
+current_file_path = os.path.abspath(__file__)
+current_dir = os.path.dirname(current_file_path)
+sys.path.append(os.path.join(current_dir, "../.."))
+from thirdparty.sam3d.sam3d_objects.pipeline.inference_pipeline_pointmap import (
+    InferencePipelinePointMap,
+)
+
+__all__ = ["Sam3dInference"]
+
+
+def load_image(path: str) -> np.ndarray:
+    image = Image.open(path)
+    image = np.array(image)
+    image = image.astype(np.uint8)
+    return image
+
+
+def load_mask(path: str) -> np.ndarray:
+    mask = load_image(path)
+    mask = mask > 0
+    if mask.ndim == 3:
+        mask = mask[..., -1]
+    return mask
+
+
+class Sam3dInference:
+    def __init__(
+        self, local_dir: str = "weights/sam-3d-objects", compile: bool = False
+    ) -> None:
+        if not os.path.exists(local_dir):
+            snapshot_download("facebook/sam-3d-objects", local_dir=local_dir)
+        config_file = os.path.join(local_dir, "checkpoints/pipeline.yaml")
+        config = OmegaConf.load(config_file)
+        config.rendering_engine = "nvdiffrast"
+        config.compile_model = compile
+        config.workspace_dir = os.path.dirname(config_file)
+        # Generate 4 gs in each pixel.
+        config["slat_decoder_gs_config_path"] = config.pop(
+            "slat_decoder_gs_4_config_path", "slat_decoder_gs_4.yaml"
+        )
+        config["slat_decoder_gs_ckpt_path"] = config.pop(
+            "slat_decoder_gs_4_ckpt_path", "slat_decoder_gs_4.ckpt"
+        )
+        self.pipeline: InferencePipelinePointMap = instantiate(config)
+
+    def merge_mask_to_rgba(
+        self, image: np.ndarray, mask: np.ndarray
+    ) -> np.ndarray:
+        mask = mask.astype(np.uint8) * 255
+        mask = mask[..., None]
+        rgba_image = np.concatenate([image[..., :3], mask], axis=-1)
+
+        return rgba_image
+
+    def run(
+        self,
+        image: np.ndarray | Image.Image,
+        mask: np.ndarray = None,
+        seed: int = None,
+        pointmap: np.ndarray = None,
+        use_stage1_distillation: bool = False,
+        use_stage2_distillation: bool = False,
+        stage1_inference_steps: int = 25,
+        stage2_inference_steps: int = 25,
+    ) -> dict:
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        if mask is not None:
+            image = self.merge_mask_to_rgba(image, mask)
+        return self.pipeline.run(
+            image,
+            None,
+            seed,
+            stage1_only=False,
+            with_mesh_postprocess=False,
+            with_texture_baking=False,
+            with_layout_postprocess=False,
+            use_vertex_color=True,
+            use_stage1_distillation=use_stage1_distillation,
+            use_stage2_distillation=use_stage2_distillation,
+            stage1_inference_steps=stage1_inference_steps,
+            stage2_inference_steps=stage2_inference_steps,
+            pointmap=pointmap,
+        )
+
+
+if __name__ == "__main__":
+    pipeline = Sam3dInference()
+
+    # load image
+    image = load_image(
+        "/home/users/xinjie.wang/xinjie/sam-3d-objects/notebook/images/shutterstock_stylish_kidsroom_1640806567/image.png"
+    )
+    mask = load_mask(
+        "/home/users/xinjie.wang/xinjie/sam-3d-objects/notebook/images/shutterstock_stylish_kidsroom_1640806567/13.png"
+    )
+
+    import torch
+
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()
+        torch.cuda.empty_cache()
+
+    from time import time
+
+    start = time()
+
+    output = pipeline.run(image, mask, seed=42)
+    print(f"Running cost: {round(time()-start, 1)}")
+
+    if torch.cuda.is_available():
+        max_memory = torch.cuda.max_memory_allocated() / (1024**3)
+        print(f"(Max VRAM): {max_memory:.2f} GB")
+
+    print(f"End: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
+
+    output["gs"].save_ply(f"outputs/splat.ply")
+    print("Your reconstruction has been saved to outputs/splat.ply")

embodied_gen/models/segment_model.py

@@ -43,6 +43,7 @@ __all__ = [
     "SAMRemover",
     "SAMPredictor",
     "RembgRemover",
+    "BMGG14Remover",
     "get_segmented_image_by_agent",
 ]
 
@@ -376,7 +377,7 @@ class BMGG14Remover(object):
 
     def __call__(
         self, image: Union[str, Image.Image, np.ndarray], save_path: str = None
-    ):
+    ) -> Image.Image:
         """Removes background from an image.
 
         Args:
@@ -496,13 +497,18 @@ if __name__ == "__main__":
     # input_image = "outputs/text2image/tmp/bucket.jpeg"
     # output_image = "outputs/text2image/tmp/bucket_seg.png"
 
-    remover = SAMRemover(model_type="vit_h")
-    remover = RembgRemover()
-    clean_image = remover(input_image)
-    clean_image.save(output_image)
-    get_segmented_image_by_agent(
-        Image.open(input_image), remover, remover, None, "./test_seg.png"
-    )
+    # remover = SAMRemover(model_type="vit_h")
+    # remover = RembgRemover()
+    # clean_image = remover(input_image)
+    # clean_image.save(output_image)
+    # get_segmented_image_by_agent(
+    #     Image.open(input_image), remover, remover, None, "./test_seg.png"
+    # )
 
     remover = BMGG14Remover()
-    remover("embodied_gen/models/test_seg.jpg", "./seg.png")
+    clean_image = remover("./camera.jpeg", "./seg.png")
+    from embodied_gen.utils.process_media import (
+        keep_largest_connected_component,
+    )
+
+    keep_largest_connected_component(clean_image).save("./seg_post.png")

embodied_gen/scripts/gen_texture.py

@@ -94,6 +94,7 @@ def entrypoint() -> None:
             delight=cfg.delight,
             no_save_delight_img=True,
             texture_wh=[cfg.texture_size, cfg.texture_size],
+            no_mesh_post_process=True,
         )
         drender_api(
             mesh_path=f"{output_root}/texture_mesh/{uuid}.obj",

embodied_gen/scripts/imageto3d.py

@@ -26,12 +26,14 @@ import numpy as np
 import torch
 import trimesh
 from PIL import Image
-from embodied_gen.data.backproject_v2 import entrypoint as backproject_api
+from embodied_gen.data.backproject_v3 import entrypoint as backproject_api
 from embodied_gen.data.utils import delete_dir, trellis_preprocess
-from embodied_gen.models.delight_model import DelightingModel
+
+# from embodied_gen.models.delight_model import DelightingModel
 from embodied_gen.models.gs_model import GaussianOperator
 from embodied_gen.models.segment_model import RembgRemover
-from embodied_gen.models.sr_model import ImageRealESRGAN
+
+# from embodied_gen.models.sr_model import ImageRealESRGAN
 from embodied_gen.scripts.render_gs import entrypoint as render_gs_api
 from embodied_gen.utils.gpt_clients import GPT_CLIENT
 from embodied_gen.utils.log import logger
@@ -59,8 +61,8 @@ os.environ["SPCONV_ALGO"] = "native"
 random.seed(0)
 
 logger.info("Loading Image3D Models...")
-DELIGHT = DelightingModel()
-IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
+# DELIGHT = DelightingModel()
+# IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
 RBG_REMOVER = RembgRemover()
 PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
     "microsoft/TRELLIS-image-large"
@@ -108,9 +110,7 @@ def parse_args():
         default=2,
     )
     parser.add_argument("--disable_decompose_convex", action="store_true")
-    parser.add_argument(
-        "--texture_wh", type=int, nargs=2, default=[2048, 2048]
-    )
+    parser.add_argument("--texture_size", type=int, default=2048)
     args, unknown = parser.parse_known_args()
 
     return args
@@ -248,16 +248,14 @@ def entrypoint(**kwargs):
             mesh.export(mesh_obj_path)
 
             mesh = backproject_api(
-                delight_model=DELIGHT,
-                imagesr_model=IMAGESR_MODEL,
-                color_path=color_path,
+                # delight_model=DELIGHT,
+                # imagesr_model=IMAGESR_MODEL,
+                gs_path=aligned_gs_path,
                 mesh_path=mesh_obj_path,
                 output_path=mesh_obj_path,
                 skip_fix_mesh=False,
-                delight=True,
-                texture_wh=args.texture_wh,
-                elevation=[20, -10, 60, -50],
-                num_images=12,
+                texture_size=args.texture_size,
+                delight=False,
             )
 
             mesh_glb_path = os.path.join(output_root, f"{filename}.glb")

embodied_gen/utils/monkey_patches.py

@@ -25,6 +25,12 @@ from omegaconf import OmegaConf
 from PIL import Image
 from torchvision import transforms
 
+__all__ = [
+    "monkey_patch_pano2room",
+    "monkey_patch_maniskill",
+    "monkey_patch_sam3d",
+]
+
 
 def monkey_patch_pano2room():
     current_file_path = os.path.abspath(__file__)
@@ -216,3 +222,374 @@ def monkey_patch_maniskill():
     ManiSkillScene.get_human_render_camera_images = (
         get_human_render_camera_images
     )
+
+
+def monkey_patch_sam3d():
+    from typing import Optional, Union
+
+    from embodied_gen.data.utils import model_device_ctx
+    from embodied_gen.utils.log import logger
+
+    os.environ["LIDRA_SKIP_INIT"] = "true"
+
+    current_file_path = os.path.abspath(__file__)
+    current_dir = os.path.dirname(current_file_path)
+    sam3d_root = os.path.abspath(
+        os.path.join(current_dir, "../../thirdparty/sam3d")
+    )
+    if sam3d_root not in sys.path:
+        sys.path.insert(0, sam3d_root)
+
+    print(f"[MonkeyPatch] Added to sys.path: {sam3d_root}")
+
+    def patch_pointmap_infer_pipeline():
+        from copy import deepcopy
+
+        try:
+            from sam3d_objects.pipeline.inference_pipeline_pointmap import (
+                InferencePipelinePointMap,
+            )
+        except ImportError:
+            logger.error(
+                "[MonkeyPatch]: Could not import sam3d_objects directly. Check paths."
+            )
+            return
+
+        def patch_run(
+            self,
+            image: Union[None, Image.Image, np.ndarray],
+            mask: Union[None, Image.Image, np.ndarray] = None,
+            seed: Optional[int] = None,
+            stage1_only=False,
+            with_mesh_postprocess=True,
+            with_texture_baking=True,
+            with_layout_postprocess=True,
+            use_vertex_color=False,
+            stage1_inference_steps=None,
+            stage2_inference_steps=None,
+            use_stage1_distillation=False,
+            use_stage2_distillation=False,
+            pointmap=None,
+            decode_formats=None,
+            estimate_plane=False,
+        ) -> dict:
+            image = self.merge_image_and_mask(image, mask)
+            with self.device:
+                pointmap_dict = self.compute_pointmap(image, pointmap)
+                pointmap = pointmap_dict["pointmap"]
+                pts = type(self)._down_sample_img(pointmap)
+                pts_colors = type(self)._down_sample_img(
+                    pointmap_dict["pts_color"]
+                )
+
+                if estimate_plane:
+                    return self.estimate_plane(pointmap_dict, image)
+
+                ss_input_dict = self.preprocess_image(
+                    image, self.ss_preprocessor, pointmap=pointmap
+                )
+
+                slat_input_dict = self.preprocess_image(
+                    image, self.slat_preprocessor
+                )
+                if seed is not None:
+                    torch.manual_seed(seed)
+
+                with model_device_ctx(
+                    self.models["ss_generator"],
+                    self.models["ss_decoder"],
+                    self.condition_embedders["ss_condition_embedder"],
+                ):
+                    ss_return_dict = self.sample_sparse_structure(
+                        ss_input_dict,
+                        inference_steps=stage1_inference_steps,
+                        use_distillation=use_stage1_distillation,
+                    )
+
+                # We could probably use the decoder from the models themselves
+                pointmap_scale = ss_input_dict.get("pointmap_scale", None)
+                pointmap_shift = ss_input_dict.get("pointmap_shift", None)
+                ss_return_dict.update(
+                    self.pose_decoder(
+                        ss_return_dict,
+                        scene_scale=pointmap_scale,
+                        scene_shift=pointmap_shift,
+                    )
+                )
+
+                logger.info(
+                    f"Rescaling scale by {ss_return_dict['downsample_factor']} after downsampling"
+                )
+                ss_return_dict["scale"] = (
+                    ss_return_dict["scale"]
+                    * ss_return_dict["downsample_factor"]
+                )
+
+                if stage1_only:
+                    logger.info("Finished!")
+                    ss_return_dict["voxel"] = (
+                        ss_return_dict["coords"][:, 1:] / 64 - 0.5
+                    )
+                    return {
+                        **ss_return_dict,
+                        "pointmap": pts.cpu().permute((1, 2, 0)),  # HxWx3
+                        "pointmap_colors": pts_colors.cpu().permute(
+                            (1, 2, 0)
+                        ),  # HxWx3
+                    }
+                    # return ss_return_dict
+
+                coords = ss_return_dict["coords"]
+                with model_device_ctx(
+                    self.models["slat_generator"],
+                    self.condition_embedders["slat_condition_embedder"],
+                ):
+                    slat = self.sample_slat(
+                        slat_input_dict,
+                        coords,
+                        inference_steps=stage2_inference_steps,
+                        use_distillation=use_stage2_distillation,
+                    )
+
+                with model_device_ctx(
+                    self.models["slat_decoder_mesh"],
+                    self.models["slat_decoder_gs"],
+                    self.models["slat_decoder_gs_4"],
+                ):
+                    outputs = self.decode_slat(
+                        slat,
+                        (
+                            self.decode_formats
+                            if decode_formats is None
+                            else decode_formats
+                        ),
+                    )
+
+                outputs = self.postprocess_slat_output(
+                    outputs,
+                    with_mesh_postprocess,
+                    with_texture_baking,
+                    use_vertex_color,
+                )
+                glb = outputs.get("glb", None)
+
+                try:
+                    if (
+                        with_layout_postprocess
+                        and self.layout_post_optimization_method is not None
+                    ):
+                        assert (
+                            glb is not None
+                        ), "require mesh to run postprocessing"
+                        logger.info(
+                            "Running layout post optimization method..."
+                        )
+                        postprocessed_pose = self.run_post_optimization(
+                            deepcopy(glb),
+                            pointmap_dict["intrinsics"],
+                            ss_return_dict,
+                            ss_input_dict,
+                        )
+                        ss_return_dict.update(postprocessed_pose)
+                except Exception as e:
+                    logger.error(
+                        f"Error during layout post optimization: {e}",
+                        exc_info=True,
+                    )
+
+                result = {
+                    **ss_return_dict,
+                    **outputs,
+                    "pointmap": pts.cpu().permute((1, 2, 0)),
+                    "pointmap_colors": pts_colors.cpu().permute((1, 2, 0)),
+                }
+                return result
+
+        InferencePipelinePointMap.run = patch_run
+
+    def patch_infer_init():
+        import torch
+
+        try:
+            from sam3d_objects.pipeline import preprocess_utils
+            from sam3d_objects.pipeline.inference_pipeline_pointmap import (
+                InferencePipeline,
+            )
+            from sam3d_objects.pipeline.inference_utils import (
+                SLAT_MEAN,
+                SLAT_STD,
+            )
+        except ImportError:
+            print(
+                "[MonkeyPatch] Error: Could not import sam3d_objects directly for infer pipeline."
+            )
+            return
+
+        def patch_init(
+            self,
+            ss_generator_config_path,
+            ss_generator_ckpt_path,
+            slat_generator_config_path,
+            slat_generator_ckpt_path,
+            ss_decoder_config_path,
+            ss_decoder_ckpt_path,
+            slat_decoder_gs_config_path,
+            slat_decoder_gs_ckpt_path,
+            slat_decoder_mesh_config_path,
+            slat_decoder_mesh_ckpt_path,
+            slat_decoder_gs_4_config_path=None,
+            slat_decoder_gs_4_ckpt_path=None,
+            ss_encoder_config_path=None,
+            ss_encoder_ckpt_path=None,
+            decode_formats=["gaussian", "mesh"],
+            dtype="bfloat16",
+            pad_size=1.0,
+            version="v0",
+            device="cuda",
+            ss_preprocessor=preprocess_utils.get_default_preprocessor(),
+            slat_preprocessor=preprocess_utils.get_default_preprocessor(),
+            ss_condition_input_mapping=["image"],
+            slat_condition_input_mapping=["image"],
+            pose_decoder_name="default",
+            workspace_dir="",
+            downsample_ss_dist=0,  # the distance we use to downsample
+            ss_inference_steps=25,
+            ss_rescale_t=3,
+            ss_cfg_strength=7,
+            ss_cfg_interval=[0, 500],
+            ss_cfg_strength_pm=0.0,
+            slat_inference_steps=25,
+            slat_rescale_t=3,
+            slat_cfg_strength=5,
+            slat_cfg_interval=[0, 500],
+            rendering_engine: str = "nvdiffrast",  # nvdiffrast OR pytorch3d,
+            shape_model_dtype=None,
+            compile_model=False,
+            slat_mean=SLAT_MEAN,
+            slat_std=SLAT_STD,
+        ):
+            self.rendering_engine = rendering_engine
+            self.device = torch.device(device)
+            self.compile_model = compile_model
+            logger.info(f"self.device: {self.device}")
+            logger.info(
+                f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', None)}"
+            )
+            logger.info(f"Actually using GPU: {torch.cuda.current_device()}")
+            with self.device:
+                self.decode_formats = decode_formats
+                self.pad_size = pad_size
+                self.version = version
+                self.ss_condition_input_mapping = ss_condition_input_mapping
+                self.slat_condition_input_mapping = (
+                    slat_condition_input_mapping
+                )
+                self.workspace_dir = workspace_dir
+                self.downsample_ss_dist = downsample_ss_dist
+                self.ss_inference_steps = ss_inference_steps
+                self.ss_rescale_t = ss_rescale_t
+                self.ss_cfg_strength = ss_cfg_strength
+                self.ss_cfg_interval = ss_cfg_interval
+                self.ss_cfg_strength_pm = ss_cfg_strength_pm
+                self.slat_inference_steps = slat_inference_steps
+                self.slat_rescale_t = slat_rescale_t
+                self.slat_cfg_strength = slat_cfg_strength
+                self.slat_cfg_interval = slat_cfg_interval
+
+                self.dtype = self._get_dtype(dtype)
+                if shape_model_dtype is None:
+                    self.shape_model_dtype = self.dtype
+                else:
+                    self.shape_model_dtype = self._get_dtype(shape_model_dtype)
+
+                # Setup preprocessors
+                self.pose_decoder = self.init_pose_decoder(
+                    ss_generator_config_path, pose_decoder_name
+                )
+                self.ss_preprocessor = self.init_ss_preprocessor(
+                    ss_preprocessor, ss_generator_config_path
+                )
+                self.slat_preprocessor = slat_preprocessor
+
+                logger.info("Loading model weights...")
+                raw_device = self.device
+                self.device = torch.device("cpu")
+                ss_generator = self.init_ss_generator(
+                    ss_generator_config_path, ss_generator_ckpt_path
+                )
+                slat_generator = self.init_slat_generator(
+                    slat_generator_config_path, slat_generator_ckpt_path
+                )
+                ss_decoder = self.init_ss_decoder(
+                    ss_decoder_config_path, ss_decoder_ckpt_path
+                )
+                ss_encoder = self.init_ss_encoder(
+                    ss_encoder_config_path, ss_encoder_ckpt_path
+                )
+                slat_decoder_gs = self.init_slat_decoder_gs(
+                    slat_decoder_gs_config_path, slat_decoder_gs_ckpt_path
+                )
+                slat_decoder_gs_4 = self.init_slat_decoder_gs(
+                    slat_decoder_gs_4_config_path, slat_decoder_gs_4_ckpt_path
+                )
+                slat_decoder_mesh = self.init_slat_decoder_mesh(
+                    slat_decoder_mesh_config_path, slat_decoder_mesh_ckpt_path
+                )
+
+                # Load conditioner embedder so that we only load it once
+                ss_condition_embedder = self.init_ss_condition_embedder(
+                    ss_generator_config_path, ss_generator_ckpt_path
+                )
+                slat_condition_embedder = self.init_slat_condition_embedder(
+                    slat_generator_config_path, slat_generator_ckpt_path
+                )
+                self.device = raw_device
+
+                self.condition_embedders = {
+                    "ss_condition_embedder": ss_condition_embedder,
+                    "slat_condition_embedder": slat_condition_embedder,
+                }
+
+                # override generator and condition embedder setting
+                self.override_ss_generator_cfg_config(
+                    ss_generator,
+                    cfg_strength=ss_cfg_strength,
+                    inference_steps=ss_inference_steps,
+                    rescale_t=ss_rescale_t,
+                    cfg_interval=ss_cfg_interval,
+                    cfg_strength_pm=ss_cfg_strength_pm,
+                )
+                self.override_slat_generator_cfg_config(
+                    slat_generator,
+                    cfg_strength=slat_cfg_strength,
+                    inference_steps=slat_inference_steps,
+                    rescale_t=slat_rescale_t,
+                    cfg_interval=slat_cfg_interval,
+                )
+
+                self.models = torch.nn.ModuleDict(
+                    {
+                        "ss_generator": ss_generator,
+                        "slat_generator": slat_generator,
+                        "ss_encoder": ss_encoder,
+                        "ss_decoder": ss_decoder,
+                        "slat_decoder_gs": slat_decoder_gs,
+                        "slat_decoder_gs_4": slat_decoder_gs_4,
+                        "slat_decoder_mesh": slat_decoder_mesh,
+                    }
+                )
+                logger.info("Loading model weights completed!")
+
+                if self.compile_model:
+                    logger.info("Compiling model...")
+                    self._compile()
+                    logger.info("Model compilation completed!")
+                self.slat_mean = torch.tensor(slat_mean)
+                self.slat_std = torch.tensor(slat_std)
+
+        InferencePipeline.__init__ = patch_init
+
+    patch_pointmap_infer_pipeline()
+    patch_infer_init()
+
+    return

embodied_gen/utils/process_media.py

@@ -230,6 +230,29 @@ def filter_image_small_connected_components(
     return image
 
 
+def keep_largest_connected_component(pil_img: Image.Image) -> Image.Image:
+    if pil_img.mode != "RGBA":
+        pil_img = pil_img.convert("RGBA")
+
+    img_arr = np.array(pil_img)
+    alpha_channel = img_arr[:, :, 3]
+
+    _, binary_mask = cv2.threshold(alpha_channel, 0, 255, cv2.THRESH_BINARY)
+    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
+        binary_mask, connectivity=8
+    )
+    if num_labels < 2:
+        return pil_img
+
+    largest_label = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
+    new_alpha = np.where(labels == largest_label, alpha_channel, 0).astype(
+        np.uint8
+    )
+    img_arr[:, :, 3] = new_alpha
+
+    return Image.fromarray(img_arr)
+
+
 def combine_images_to_grid(
     images: list[str | Image.Image],
     cat_row_col: tuple[int, int] = None,
@@ -439,7 +462,7 @@ class SceneTreeVisualizer:
         plt.axis("off")
 
         legend_handles = [
-            Patch(facecolor=color, edgecolor='black', label=role)
+            Patch(facecolor=color, edgecolor="black", label=role)
             for role, color in self.role_colors.items()
         ]
         plt.legend(
@@ -465,7 +488,7 @@ def load_scene_dict(file_path: str) -> dict:
         dict: Mapping from scene ID to description.
     """
     scene_dict = {}
-    with open(file_path, "r", encoding='utf-8') as f:
+    with open(file_path, "r", encoding="utf-8") as f:
         for line in f:
             line = line.strip()
             if not line or ":" not in line:
@@ -487,7 +510,7 @@ def is_image_file(filename: str) -> bool:
     """
     mime_type, _ = mimetypes.guess_type(filename)
 
-    return mime_type is not None and mime_type.startswith('image')
+    return mime_type is not None and mime_type.startswith("image")
 
 
 def parse_text_prompts(prompts: list[str]) -> list[str]:

embodied_gen/utils/trender.py

@@ -16,6 +16,7 @@
 
 import os
 import sys
+from collections import defaultdict
 
 import numpy as np
 import spaces
@@ -25,10 +26,8 @@ from tqdm import tqdm
 current_file_path = os.path.abspath(__file__)
 current_dir = os.path.dirname(current_file_path)
 sys.path.append(os.path.join(current_dir, "../.."))
-from thirdparty.TRELLIS.trellis.renderers.mesh_renderer import MeshRenderer
-from thirdparty.TRELLIS.trellis.representations import MeshExtractResult
+from thirdparty.TRELLIS.trellis.renderers import GaussianRenderer, MeshRenderer
 from thirdparty.TRELLIS.trellis.utils.render_utils import (
-    render_frames,
     yaw_pitch_r_fov_to_extrinsics_intrinsics,
 )
 
@@ -38,7 +37,7 @@ __all__ = [
 
 
 @spaces.GPU
-def render_mesh(sample, extrinsics, intrinsics, options={}, **kwargs):
+def render_mesh_frames(sample, extrinsics, intrinsics, options={}, **kwargs):
     renderer = MeshRenderer()
     renderer.rendering_options.resolution = options.get("resolution", 512)
     renderer.rendering_options.near = options.get("near", 1)
@@ -60,6 +59,57 @@ def render_mesh(sample, extrinsics, intrinsics, options={}, **kwargs):
     return rets
 
 
+@spaces.GPU
+def render_gs_frames(
+    sample,
+    extrinsics,
+    intrinsics,
+    options=None,
+    colors_overwrite=None,
+    verbose=True,
+    **kwargs,
+):
+    def to_img(tensor):
+        return np.clip(
+            tensor.detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255
+        ).astype(np.uint8)
+
+    def to_numpy(tensor):
+        return tensor.detach().cpu().numpy()
+
+    renderer = GaussianRenderer()
+    renderer.pipe.kernel_size = kwargs.get("kernel_size", 0.1)
+    renderer.pipe.use_mip_gaussian = True
+
+    defaults = {
+        "resolution": 512,
+        "near": 0.8,
+        "far": 1.6,
+        "bg_color": (0, 0, 0),
+        "ssaa": 1,
+    }
+    final_options = {**defaults, **(options or {})}
+
+    for k, v in final_options.items():
+        if hasattr(renderer.rendering_options, k):
+            setattr(renderer.rendering_options, k, v)
+
+    outputs = defaultdict(list)
+    iterator = zip(extrinsics, intrinsics)
+    if verbose:
+        iterator = tqdm(iterator, total=len(extrinsics), desc="Rendering")
+
+    for extr, intr in iterator:
+        res = renderer.render(
+            sample, extr, intr, colors_overwrite=colors_overwrite
+        )
+        outputs["color"].append(to_img(res["color"]))
+        depth = res.get("percent_depth") or res.get("depth")
+        outputs["depth"].append(to_numpy(depth) if depth is not None else None)
+
+    return dict(outputs)
+
+
 @spaces.GPU
 def render_video(
     sample,
@@ -77,7 +127,9 @@ def render_video(
         yaws, pitch, r, fov
     )
     render_fn = (
-        render_mesh if isinstance(sample, MeshExtractResult) else render_frames
+        render_mesh_frames
+        if sample.__class__.__name__ == "MeshExtractResult"
+        else render_gs_frames
     )
     result = render_fn(
         sample,

thirdparty/TRELLIS/trellis/utils/postprocessing_utils.py

@@ -440,7 +440,7 @@ def to_glb(
     vertices, faces, uvs = parametrize_mesh(vertices, faces)
 
     # bake texture
-    observations, extrinsics, intrinsics = render_multiview(app_rep, resolution=1024, nviews=200)
+    observations, extrinsics, intrinsics = render_multiview(app_rep, resolution=1024, nviews=100)
     masks = [np.any(observation > 0, axis=-1) for observation in observations]
     extrinsics = [extrinsics[i].cpu().numpy() for i in range(len(extrinsics))]
     intrinsics = [intrinsics[i].cpu().numpy() for i in range(len(intrinsics))]