Initial clean commit

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,15 @@
+---
+title: EmbodiedGen Text To 3D
+emoji: 📝
+colorFrom: blue
+colorTo: green
+sdk: gradio
+sdk_version: 5.33.1
+app_file: app.py
+pinned: false
+license: apache-2.0
+short_description: Create 3D models from text descriptions
+paper: https://huggingface.co/papers/2506.10600
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,479 @@
+# 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.
+
+
+import os
+
+os.environ["GRADIO_APP"] = "textto3d"
+
+
+import gradio as gr
+from app_style import custom_theme, image_css, lighting_css
+from common import (
+    MAX_SEED,
+    VERSION,
+    active_btn_by_text_content,
+    end_session,
+    extract_3d_representations_v2,
+    extract_urdf,
+    get_cached_image,
+    get_seed,
+    get_selected_image,
+    image_to_3d,
+    start_session,
+    text2image_fn,
+)
+
+with gr.Blocks(delete_cache=(43200, 43200), theme=custom_theme) as demo:
+    gr.HTML(image_css, visible=False)
+    # gr.HTML(lighting_css, visible=False)
+    gr.Markdown(
+        """
+        ## ***EmbodiedGen***: Text-to-3D Asset
+        **🔖 Version**: {VERSION}
+        <p style="display: flex; gap: 10px; flex-wrap: nowrap;">
+            <a href="https://horizonrobotics.github.io/EmbodiedGen">
+                <img alt="📖 Documentation" src="https://img.shields.io/badge/📖-Documentation-blue">
+            </a>
+            <a href="https://arxiv.org/abs/2506.10600">
+                <img alt="📄 arXiv" src="https://img.shields.io/badge/📄-arXiv-b31b1b">
+            </a>
+            <a href="https://github.com/HorizonRobotics/EmbodiedGen">
+                <img alt="💻 GitHub" src="https://img.shields.io/badge/GitHub-000000?logo=github">
+            </a>
+            <a href="https://www.youtube.com/watch?v=rG4odybuJRk">
+                <img alt="🎥 Video" src="https://img.shields.io/badge/🎥-Video-red">
+            </a>
+        </p>
+
+        📝 Create 3D assets from text descriptions for a wide range of geometry and styles.
+        """.format(
+            VERSION=VERSION
+        ),
+        elem_classes=["header"],
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            raw_image_cache = gr.Image(
+                format="png",
+                image_mode="RGB",
+                type="pil",
+                visible=False,
+            )
+            text_prompt = gr.Textbox(
+                label="Text Prompt (Chinese or English)",
+                placeholder="Input text prompt here",
+            )
+            ip_image = gr.Image(
+                label="Reference Image(optional)",
+                format="png",
+                image_mode="RGB",
+                type="filepath",
+                height=250,
+                elem_classes=["image_fit"],
+            )
+            gr.Markdown(
+                "Note: The `reference image` is optional, if use, "
+                "please provide image in nearly square resolution."
+            )
+
+            with gr.Accordion(label="Image Generation Settings", open=False):
+                with gr.Row():
+                    seed = gr.Slider(
+                        0, MAX_SEED, label="Seed", value=0, step=1
+                    )
+                    randomize_seed = gr.Checkbox(
+                        label="Randomize Seed", value=False
+                    )
+                rmbg_tag = gr.Radio(
+                    choices=["rembg", "rmbg14"],
+                    value="rembg",
+                    label="Background Removal Model",
+                )
+                ip_adapt_scale = gr.Slider(
+                    0, 1, label="IP-adapter Scale", value=0.3, step=0.05
+                )
+                img_guidance_scale = gr.Slider(
+                    1, 30, label="Text Guidance Scale", value=12, step=0.2
+                )
+                img_inference_steps = gr.Slider(
+                    10, 100, label="Sampling Steps", value=50, step=5
+                )
+                img_resolution = gr.Slider(
+                    512,
+                    1536,
+                    label="Image Resolution",
+                    value=1024,
+                    step=128,
+                )
+
+            generate_img_btn = gr.Button(
+                "🎨 1. Generate Images(~1min)",
+                variant="primary",
+                interactive=False,
+            )
+            dropdown = gr.Radio(
+                choices=["sample1", "sample2", "sample3"],
+                value="sample1",
+                label="Choose your favorite sample style.",
+            )
+            select_img = gr.Image(
+                visible=False,
+                format="png",
+                image_mode="RGBA",
+                type="pil",
+                height=300,
+            )
+
+            # text to 3d
+            with gr.Accordion(label="Generation Settings", open=False):
+                seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
+                texture_size = gr.Slider(
+                    1024, 4096, label="UV texture size", value=2048, step=256
+                )
+                with gr.Row():
+                    randomize_seed = gr.Checkbox(
+                        label="Randomize Seed", value=False
+                    )
+                    project_delight = gr.Checkbox(
+                        label="Back-project Delight", value=True
+                    )
+                gr.Markdown("Geo Structure Generation")
+                with gr.Row():
+                    ss_guidance_strength = gr.Slider(
+                        0.0,
+                        10.0,
+                        label="Guidance Strength",
+                        value=7.5,
+                        step=0.1,
+                    )
+                    ss_sampling_steps = gr.Slider(
+                        1, 50, label="Sampling Steps", value=12, step=1
+                    )
+                gr.Markdown("Visual Appearance Generation")
+                with gr.Row():
+                    slat_guidance_strength = gr.Slider(
+                        0.0,
+                        10.0,
+                        label="Guidance Strength",
+                        value=3.0,
+                        step=0.1,
+                    )
+                    slat_sampling_steps = gr.Slider(
+                        1, 50, label="Sampling Steps", value=12, step=1
+                    )
+
+            generate_btn = gr.Button(
+                "🚀 2. Generate 3D(~0.5 mins)",
+                variant="primary",
+                interactive=False,
+            )
+            model_output_obj = gr.Textbox(label="raw mesh .obj", visible=False)
+            with gr.Row():
+                extract_rep3d_btn = gr.Button(
+                    "🔍 3. Extract 3D Representation(~1 mins)",
+                    variant="primary",
+                    interactive=False,
+                )
+            with gr.Accordion(
+                label="Enter Asset Attributes(optional)", open=False
+            ):
+                asset_cat_text = gr.Textbox(
+                    label="Enter Asset Category (e.g., chair)"
+                )
+                height_range_text = gr.Textbox(
+                    label="Enter Height Range in meter (e.g., 0.5-0.6)"
+                )
+                mass_range_text = gr.Textbox(
+                    label="Enter Mass Range in kg (e.g., 1.1-1.2)"
+                )
+                asset_version_text = gr.Textbox(
+                    label=f"Enter version (e.g., {VERSION})"
+                )
+            with gr.Row():
+                extract_urdf_btn = gr.Button(
+                    "🧩 4. Extract URDF with physics(~1 mins)",
+                    variant="primary",
+                    interactive=False,
+                )
+            with gr.Row():
+                download_urdf = gr.DownloadButton(
+                    label="⬇️ 5. Download URDF",
+                    variant="primary",
+                    interactive=False,
+                )
+
+        with gr.Column(scale=3):
+            with gr.Row():
+                image_sample1 = gr.Image(
+                    label="sample1",
+                    format="png",
+                    image_mode="RGBA",
+                    type="filepath",
+                    height=300,
+                    interactive=False,
+                    elem_classes=["image_fit"],
+                )
+                image_sample2 = gr.Image(
+                    label="sample2",
+                    format="png",
+                    image_mode="RGBA",
+                    type="filepath",
+                    height=300,
+                    interactive=False,
+                    elem_classes=["image_fit"],
+                )
+                image_sample3 = gr.Image(
+                    label="sample3",
+                    format="png",
+                    image_mode="RGBA",
+                    type="filepath",
+                    height=300,
+                    interactive=False,
+                    elem_classes=["image_fit"],
+                )
+                usample1 = gr.Image(
+                    format="png",
+                    image_mode="RGBA",
+                    type="filepath",
+                    visible=False,
+                )
+                usample2 = gr.Image(
+                    format="png",
+                    image_mode="RGBA",
+                    type="filepath",
+                    visible=False,
+                )
+                usample3 = gr.Image(
+                    format="png",
+                    image_mode="RGBA",
+                    type="filepath",
+                    visible=False,
+                )
+            gr.Markdown(
+                "Generated image may be poor quality due to auto seg."
+                "Retry by adjusting text prompt, seed or switch seg model in `Image Gen Settings`."
+            )
+            with gr.Row():
+                video_output = gr.Video(
+                    label="Generated 3D Asset",
+                    autoplay=True,
+                    loop=True,
+                    height=300,
+                    interactive=False,
+                )
+                model_output_gs = gr.Model3D(
+                    label="Gaussian Representation",
+                    height=300,
+                    interactive=False,
+                )
+                aligned_gs = gr.Textbox(visible=False)
+
+                model_output_mesh = gr.Model3D(
+                    label="Mesh Representation",
+                    clear_color=[0.8, 0.8, 0.8, 1],
+                    height=300,
+                    interactive=False,
+                    elem_id="lighter_mesh",
+                )
+
+            gr.Markdown("Estimated Asset 3D Attributes(No input required)")
+            with gr.Row():
+                est_type_text = gr.Textbox(
+                    label="Asset category", interactive=False
+                )
+                est_height_text = gr.Textbox(
+                    label="Real height(.m)", interactive=False
+                )
+                est_mass_text = gr.Textbox(
+                    label="Mass(.kg)", interactive=False
+                )
+                est_mu_text = gr.Textbox(
+                    label="Friction coefficient", interactive=False
+                )
+
+            prompt_examples = [
+                "satin gold tea cup with saucer",
+                "small bronze figurine of a lion",
+                "brown leather bag",
+                "Miniature cup with floral design",
+                "带木质底座, 具有经纬线的地球仪",
+                "橙色电动手钻, 有磨损细节",
+                "手工制作的皮革笔记本",
+            ]
+            examples = gr.Examples(
+                label="Gallery",
+                examples=prompt_examples,
+                inputs=[text_prompt],
+                examples_per_page=10,
+            )
+
+    output_buf = gr.State()
+
+    demo.load(start_session)
+    demo.unload(end_session)
+
+    text_prompt.change(
+        active_btn_by_text_content,
+        inputs=[text_prompt],
+        outputs=[generate_img_btn],
+    )
+
+    generate_img_btn.click(
+        lambda: tuple(
+            [
+                gr.Button(interactive=False),
+                gr.Button(interactive=False),
+                gr.Button(interactive=False),
+                gr.Button(interactive=False),
+                None,
+                "",
+                None,
+                None,
+                "",
+                "",
+                "",
+                "",
+                "",
+                "",
+                "",
+                "",
+                None,
+                None,
+                None,
+            ]
+        ),
+        outputs=[
+            extract_rep3d_btn,
+            extract_urdf_btn,
+            download_urdf,
+            generate_btn,
+            model_output_gs,
+            aligned_gs,
+            model_output_mesh,
+            video_output,
+            asset_cat_text,
+            height_range_text,
+            mass_range_text,
+            asset_version_text,
+            est_type_text,
+            est_height_text,
+            est_mass_text,
+            est_mu_text,
+            image_sample1,
+            image_sample2,
+            image_sample3,
+        ],
+    ).success(
+        text2image_fn,
+        inputs=[
+            text_prompt,
+            img_guidance_scale,
+            img_inference_steps,
+            ip_image,
+            ip_adapt_scale,
+            img_resolution,
+            rmbg_tag,
+            seed,
+        ],
+        outputs=[
+            image_sample1,
+            image_sample2,
+            image_sample3,
+            usample1,
+            usample2,
+            usample3,
+        ],
+    ).success(
+        lambda: gr.Button(interactive=True),
+        outputs=[generate_btn],
+    )
+
+    generate_btn.click(
+        get_seed,
+        inputs=[randomize_seed, seed],
+        outputs=[seed],
+    ).success(
+        get_selected_image,
+        inputs=[dropdown, usample1, usample2, usample3],
+        outputs=select_img,
+    ).success(
+        get_cached_image,
+        inputs=[select_img],
+        outputs=[raw_image_cache],
+    ).success(
+        image_to_3d,
+        inputs=[
+            select_img,
+            seed,
+            ss_guidance_strength,
+            ss_sampling_steps,
+            slat_guidance_strength,
+            slat_sampling_steps,
+            raw_image_cache,
+        ],
+        outputs=[output_buf, video_output],
+    ).success(
+        lambda: gr.Button(interactive=True),
+        outputs=[extract_rep3d_btn],
+    )
+
+    extract_rep3d_btn.click(
+        extract_3d_representations_v2,
+        inputs=[
+            output_buf,
+            project_delight,
+            texture_size,
+        ],
+        outputs=[
+            model_output_mesh,
+            model_output_gs,
+            model_output_obj,
+            aligned_gs,
+        ],
+    ).success(
+        lambda: gr.Button(interactive=True),
+        outputs=[extract_urdf_btn],
+    )
+
+    extract_urdf_btn.click(
+        extract_urdf,
+        inputs=[
+            aligned_gs,
+            model_output_obj,
+            asset_cat_text,
+            height_range_text,
+            mass_range_text,
+            asset_version_text,
+        ],
+        outputs=[
+            download_urdf,
+            est_type_text,
+            est_height_text,
+            est_mass_text,
+            est_mu_text,
+        ],
+        queue=True,
+        show_progress="full",
+    ).success(
+        lambda: gr.Button(interactive=True),
+        outputs=[download_urdf],
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()

app_style.py

@@ -0,0 +1,27 @@
+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(1.9) !important;
+}
+</style>
+"""
+
+image_css = """
+<style>
+.image_fit .image-frame {
+object-fit: contain !important;
+height: 100% !important;
+}
+</style>
+"""
+
+custom_theme = Soft(
+    primary_hue=stone,
+    secondary_hue=gray,
+    radius_size="md",
+    text_size="sm",
+    spacing_size="sm",
+)

common.py

@@ -0,0 +1,849 @@
+# 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.
+
+import gc
+import logging
+import os
+import shutil
+import subprocess
+import sys
+from glob import glob
+
+import cv2
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+import torch.nn.functional as F
+import trimesh
+from easydict import EasyDict as edict
+from PIL import Image
+from embodied_gen.data.backproject_v2 import entrypoint as backproject_api
+from embodied_gen.data.differentiable_render import entrypoint as render_api
+from embodied_gen.data.utils import trellis_preprocess, zip_files
+from embodied_gen.models.delight_model import DelightingModel
+from embodied_gen.models.gs_model import GaussianOperator
+from embodied_gen.models.segment_model import (
+    BMGG14Remover,
+    RembgRemover,
+    SAMPredictor,
+)
+from embodied_gen.models.sr_model import ImageRealESRGAN, ImageStableSR
+from embodied_gen.scripts.render_gs import entrypoint as render_gs_api
+from embodied_gen.scripts.render_mv import build_texture_gen_pipe, infer_pipe
+from embodied_gen.scripts.text2image import (
+    build_text2img_ip_pipeline,
+    build_text2img_pipeline,
+    text2img_gen,
+)
+from embodied_gen.utils.gpt_clients import GPT_CLIENT
+from embodied_gen.utils.process_media import (
+    filter_image_small_connected_components,
+    merge_images_video,
+)
+from embodied_gen.utils.tags import VERSION
+from embodied_gen.utils.trender import render_video
+from embodied_gen.validators.quality_checkers import (
+    BaseChecker,
+    ImageAestheticChecker,
+    ImageSegChecker,
+    MeshGeoChecker,
+)
+from embodied_gen.validators.urdf_convertor import URDFGenerator
+
+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.pipelines import TrellisImageTo3DPipeline
+from thirdparty.TRELLIS.trellis.representations import (
+    Gaussian,
+    MeshExtractResult,
+)
+from thirdparty.TRELLIS.trellis.representations.gaussian.general_utils import (
+    build_scaling_rotation,
+    inverse_sigmoid,
+    strip_symmetric,
+)
+from thirdparty.TRELLIS.trellis.utils import postprocessing_utils
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+os.environ["TORCH_EXTENSIONS_DIR"] = os.path.expanduser(
+    "~/.cache/torch_extensions"
+)
+os.environ["GRADIO_ANALYTICS_ENABLED"] = "false"
+os.environ["SPCONV_ALGO"] = "native"
+
+MAX_SEED = 100000
+
+
+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
+    ):
+        L = build_scaling_rotation(scaling_modifier * scaling, rotation)
+        actual_covariance = L @ L.transpose(1, 2)
+        symm = strip_symmetric(actual_covariance)
+        return symm
+
+    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
+
+
+DELIGHT = DelightingModel()
+IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
+# IMAGESR_MODEL = ImageStableSR()
+if os.getenv("GRADIO_APP") == "imageto3d":
+    RBG_REMOVER = RembgRemover()
+    RBG14_REMOVER = BMGG14Remover()
+    SAM_PREDICTOR = SAMPredictor(model_type="vit_h", device="cpu")
+    PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
+        "microsoft/TRELLIS-image-large"
+    )
+    # PIPELINE.cuda()
+    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()
+    PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
+        "microsoft/TRELLIS-image-large"
+    )
+    # PIPELINE.cuda()
+    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":
+    PIPELINE_IP = build_texture_gen_pipe(
+        base_ckpt_dir="./weights",
+        ip_adapt_scale=0.7,
+        device="cuda",
+    )
+    PIPELINE = build_texture_gen_pipe(
+        base_ckpt_dir="./weights",
+        ip_adapt_scale=0,
+        device="cuda",
+    )
+    TMP_DIR = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "sessions/texture_edit"
+    )
+    os.makedirs(TMP_DIR, exist_ok=True)
+
+
+def start_session(req: gr.Request) -> None:
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    os.makedirs(user_dir, exist_ok=True)
+
+
+def end_session(req: gr.Request) -> None:
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    if os.path.exists(user_dir):
+        shutil.rmtree(user_dir)
+
+
+@spaces.GPU
+def preprocess_image_fn(
+    image: str | np.ndarray | Image.Image, rmbg_tag: str = "rembg"
+) -> 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))
+
+    bg_remover = RBG_REMOVER if rmbg_tag == "rembg" else RBG14_REMOVER
+    image = bg_remover(image)
+    image = trellis_preprocess(image)
+
+    return image, image_cache
+
+
+def preprocess_sam_image_fn(
+    image: Image.Image,
+) -> tuple[Image.Image, Image.Image]:
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+
+    sam_image = SAM_PREDICTOR.preprocess_image(image)
+    image_cache = Image.fromarray(sam_image).resize((512, 512))
+    SAM_PREDICTOR.predictor.set_image(sam_image)
+
+    return sam_image, image_cache
+
+
+def active_btn_by_content(content: gr.Image) -> gr.Button:
+    interactive = True if content is not None else False
+
+    return gr.Button(interactive=interactive)
+
+
+def active_btn_by_text_content(content: gr.Textbox) -> gr.Button:
+    if content is not None and len(content) > 0:
+        interactive = True
+    else:
+        interactive = False
+
+    return gr.Button(interactive=interactive)
+
+
+def get_selected_image(
+    choice: str, sample1: str, sample2: str, sample3: str
+) -> str:
+    if choice == "sample1":
+        return sample1
+    elif choice == "sample2":
+        return sample2
+    elif choice == "sample3":
+        return sample3
+    else:
+        raise ValueError(f"Invalid choice: {choice}")
+
+
+def get_cached_image(image_path: str) -> Image.Image:
+    if isinstance(image_path, Image.Image):
+        return image_path
+    return Image.open(image_path).resize((512, 512))
+
+
+@spaces.GPU
+def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
+    return {
+        "gaussian": {
+            **gs.init_params,
+            "_xyz": gs._xyz.cpu().numpy(),
+            "_features_dc": gs._features_dc.cpu().numpy(),
+            "_scaling": gs._scaling.cpu().numpy(),
+            "_rotation": gs._rotation.cpu().numpy(),
+            "_opacity": gs._opacity.cpu().numpy(),
+        },
+        "mesh": {
+            "vertices": mesh.vertices.cpu().numpy(),
+            "faces": mesh.faces.cpu().numpy(),
+        },
+    }
+
+
+def unpack_state(state: dict, device: str = "cpu") -> tuple[Gaussian, dict]:
+    gs = Gaussian(
+        aabb=state["gaussian"]["aabb"],
+        sh_degree=state["gaussian"]["sh_degree"],
+        mininum_kernel_size=state["gaussian"]["mininum_kernel_size"],
+        scaling_bias=state["gaussian"]["scaling_bias"],
+        opacity_bias=state["gaussian"]["opacity_bias"],
+        scaling_activation=state["gaussian"]["scaling_activation"],
+        device=device,
+    )
+    gs._xyz = torch.tensor(state["gaussian"]["_xyz"], device=device)
+    gs._features_dc = torch.tensor(
+        state["gaussian"]["_features_dc"], device=device
+    )
+    gs._scaling = torch.tensor(state["gaussian"]["_scaling"], device=device)
+    gs._rotation = torch.tensor(state["gaussian"]["_rotation"], device=device)
+    gs._opacity = torch.tensor(state["gaussian"]["_opacity"], device=device)
+
+    mesh = edict(
+        vertices=torch.tensor(state["mesh"]["vertices"], device=device),
+        faces=torch.tensor(state["mesh"]["faces"], device=device),
+    )
+
+    return gs, mesh
+
+
+def get_seed(randomize_seed: bool, seed: int, max_seed: int = MAX_SEED) -> int:
+    return np.random.randint(0, max_seed) if randomize_seed else seed
+
+
+def select_point(
+    image: np.ndarray,
+    sel_pix: list,
+    point_type: str,
+    evt: gr.SelectData,
+):
+    if point_type == "foreground_point":
+        sel_pix.append((evt.index, 1))  # append the foreground_point
+    elif point_type == "background_point":
+        sel_pix.append((evt.index, 0))  # append the background_point
+    else:
+        sel_pix.append((evt.index, 1))  # default foreground_point
+
+    masks = SAM_PREDICTOR.generate_masks(image, sel_pix)
+    seg_image = SAM_PREDICTOR.get_segmented_image(image, masks)
+
+    for point, label in sel_pix:
+        color = (255, 0, 0) if label == 0 else (0, 255, 0)
+        marker_type = 1 if label == 0 else 5
+        cv2.drawMarker(
+            image,
+            point,
+            color,
+            markerType=marker_type,
+            markerSize=15,
+            thickness=10,
+        )
+
+    torch.cuda.empty_cache()
+
+    return (image, masks), seg_image
+
+
+@spaces.GPU
+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,
+    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")
+        seg_image = trellis_preprocess(seg_image)
+    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")
+    PIPELINE.cuda()
+    outputs = PIPELINE.run(
+        seg_image,
+        seed=seed,
+        formats=["gaussian", "mesh"],
+        preprocess_image=False,
+        sparse_structure_sampler_params={
+            "steps": ss_sampling_steps,
+            "cfg_strength": ss_guidance_strength,
+        },
+        slat_sampler_params={
+            "steps": slat_sampling_steps,
+            "cfg_strength": slat_guidance_strength,
+        },
+    )
+    # Set to cpu for memory saving.
+    PIPELINE.cpu()
+
+    gs_model = outputs["gaussian"][0]
+    mesh_model = outputs["mesh"][0]
+    color_images = render_video(gs_model)["color"]
+    normal_images = render_video(mesh_model)["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 extract_3d_representations(
+    state: dict, enable_delight: bool, texture_size: int, 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, device="cuda")
+
+    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,
+    texture_size: int,
+    req: gr.Request,
+):
+    output_root = TMP_DIR
+    user_dir = os.path.join(output_root, str(req.session_hash))
+    gs_model, mesh_model = unpack_state(state, device="cpu")
+
+    filename = "sample"
+    gs_path = os.path.join(user_dir, 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]]
+    gs_add_rot = [[1, 0, 0], [0, -1, 0], [0, 0, -1]]
+    mesh_add_rot = [[1, 0, 0], [0, 0, -1], [0, 1, 0]]
+
+    # Addtional rotation for GS to align mesh.
+    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")
+    GaussianOperator.resave_ply(
+        in_ply=gs_path,
+        out_ply=aligned_gs_path,
+        instance_pose=pose,
+        device="cpu",
+    )
+    color_path = os.path.join(user_dir, "color.png")
+    render_gs_api(
+        input_gs=aligned_gs_path,
+        output_path=color_path,
+        elevation=[20, -10, 60, -50],
+        num_images=12,
+    )
+
+    mesh = trimesh.Trimesh(
+        vertices=mesh_model.vertices.cpu().numpy(),
+        faces=mesh_model.faces.cpu().numpy(),
+    )
+    mesh.vertices = mesh.vertices @ np.array(mesh_add_rot)
+    mesh.vertices = mesh.vertices @ np.array(rot_matrix)
+
+    mesh_obj_path = os.path.join(user_dir, f"{filename}.obj")
+    mesh.export(mesh_obj_path)
+
+    mesh = backproject_api(
+        delight_model=DELIGHT,
+        imagesr_model=IMAGESR_MODEL,
+        color_path=color_path,
+        mesh_path=mesh_obj_path,
+        output_path=mesh_obj_path,
+        skip_fix_mesh=False,
+        delight=enable_delight,
+        texture_wh=[texture_size, texture_size],
+        elevation=[20, -10, 60, -50],
+        num_images=12,
+    )
+
+    mesh_glb_path = os.path.join(user_dir, f"{filename}.glb")
+    mesh.export(mesh_glb_path)
+
+    return mesh_glb_path, gs_path, mesh_obj_path, aligned_gs_path
+
+
+def extract_urdf(
+    gs_path: str,
+    mesh_obj_path: str,
+    asset_cat_text: str,
+    height_range_text: str,
+    mass_range_text: str,
+    asset_version_text: str,
+    req: gr.Request = None,
+):
+    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, decompose_convex=True
+    )
+    asset_attrs = {
+        "version": VERSION,
+        "gs_model": f"{urdf_convertor.output_mesh_dir}/{filename}_gs.ply",
+    }
+    if asset_version_text:
+        asset_attrs["version"] = asset_version_text
+    if asset_cat_text:
+        asset_attrs["category"] = asset_cat_text.lower()
+    if height_range_text:
+        try:
+            min_height, max_height = map(float, height_range_text.split("-"))
+            asset_attrs["min_height"] = min_height
+            asset_attrs["max_height"] = max_height
+        except ValueError:
+            return "Invalid height input format. Use the format: min-max."
+    if mass_range_text:
+        try:
+            min_mass, max_mass = map(float, mass_range_text.split("-"))
+            asset_attrs["min_mass"] = min_mass
+            asset_attrs["max_mass"] = max_mass
+        except ValueError:
+            return "Invalid mass input format. Use the format: min-max."
+
+    urdf_path = urdf_convertor(
+        mesh_path=mesh_obj_path,
+        output_root=f"{output_root}/URDF_{filename}",
+        **asset_attrs,
+    )
+
+    # Rescale GS and save to URDF/mesh folder.
+    real_height = urdf_convertor.get_attr_from_urdf(
+        urdf_path, attr_name="real_height"
+    )
+    out_gs = f"{output_root}/URDF_{filename}/{urdf_convertor.output_mesh_dir}/{filename}_gs.ply"  # noqa
+    GaussianOperator.resave_ply(
+        in_ply=gs_path,
+        out_ply=out_gs,
+        real_height=real_height,
+        device="cpu",
+    )
+
+    # Quality check and update .urdf file.
+    mesh_out = f"{output_root}/URDF_{filename}/{urdf_convertor.output_mesh_dir}/{filename}.obj"  # noqa
+    trimesh.load(mesh_out).export(mesh_out.replace(".obj", ".glb"))
+    # image_paths = render_asset3d(
+    #     mesh_path=mesh_out,
+    #     output_root=f"{output_root}/URDF_{filename}",
+    #     output_subdir="qa_renders",
+    #     num_images=8,
+    #     elevation=(30, -30),
+    #     distance=5.5,
+    # )
+
+    image_dir = f"{output_root}/URDF_{filename}/{urdf_convertor.output_render_dir}/image_color"  # noqa
+    image_paths = glob(f"{image_dir}/*.png")
+    images_list = []
+    for checker in CHECKERS:
+        images = image_paths
+        if isinstance(checker, ImageSegChecker):
+            images = [
+                f"{TMP_DIR}/{req.session_hash}/raw_image.png",
+                f"{TMP_DIR}/{req.session_hash}/seg_image.png",
+            ]
+        images_list.append(images)
+
+    results = BaseChecker.validate(CHECKERS, images_list)
+    urdf_convertor.add_quality_tag(urdf_path, results)
+
+    # Zip urdf files
+    urdf_zip = zip_files(
+        input_paths=[
+            f"{output_root}/URDF_{filename}/{urdf_convertor.output_mesh_dir}",
+            f"{output_root}/URDF_{filename}/{filename}.urdf",
+        ],
+        output_zip=f"{output_root}/urdf_{filename}.zip",
+    )
+
+    estimated_type = urdf_convertor.estimated_attrs["category"]
+    estimated_height = urdf_convertor.estimated_attrs["height"]
+    estimated_mass = urdf_convertor.estimated_attrs["mass"]
+    estimated_mu = urdf_convertor.estimated_attrs["mu"]
+
+    return (
+        urdf_zip,
+        estimated_type,
+        estimated_height,
+        estimated_mass,
+        estimated_mu,
+    )
+
+
+@spaces.GPU
+def text2image_fn(
+    prompt: str,
+    guidance_scale: float,
+    infer_step: int = 50,
+    ip_image: Image.Image | str = None,
+    ip_adapt_scale: float = 0.3,
+    image_wh: int | tuple[int, int] = [1024, 1024],
+    rmbg_tag: str = "rembg",
+    seed: int = None,
+    n_sample: int = 3,
+    req: gr.Request = None,
+):
+    if isinstance(image_wh, int):
+        image_wh = (image_wh, image_wh)
+    output_root = TMP_DIR
+    if req is not None:
+        output_root = os.path.join(output_root, str(req.session_hash))
+        os.makedirs(output_root, exist_ok=True)
+
+    pipeline = PIPELINE_IMG if ip_image is None else PIPELINE_IMG_IP
+    if ip_image is not None:
+        pipeline.set_ip_adapter_scale([ip_adapt_scale])
+
+    images = text2img_gen(
+        prompt=prompt,
+        n_sample=n_sample,
+        guidance_scale=guidance_scale,
+        pipeline=pipeline,
+        ip_image=ip_image,
+        image_wh=image_wh,
+        infer_step=infer_step,
+        seed=seed,
+    )
+
+    for idx in range(len(images)):
+        image = images[idx]
+        images[idx], _ = preprocess_image_fn(image, rmbg_tag)
+
+    save_paths = []
+    for idx, image in enumerate(images):
+        save_path = f"{output_root}/sample_{idx}.png"
+        image.save(save_path)
+        save_paths.append(save_path)
+
+    logger.info(f"Images saved to {output_root}")
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    return save_paths + save_paths
+
+
+@spaces.GPU
+def generate_condition(mesh_path: str, req: gr.Request, uuid: str = "sample"):
+    output_root = os.path.join(TMP_DIR, str(req.session_hash))
+
+    _ = render_api(
+        mesh_path=mesh_path,
+        output_root=f"{output_root}/condition",
+        uuid=str(uuid),
+    )
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    return None, None, None
+
+
+@spaces.GPU
+def generate_texture_mvimages(
+    prompt: str,
+    controlnet_cond_scale: float = 0.55,
+    guidance_scale: float = 9,
+    strength: float = 0.9,
+    num_inference_steps: int = 50,
+    seed: int = 0,
+    ip_adapt_scale: float = 0,
+    ip_img_path: str = None,
+    uid: str = "sample",
+    sub_idxs: tuple[tuple[int]] = ((0, 1, 2), (3, 4, 5)),
+    req: gr.Request = None,
+) -> list[str]:
+    output_root = os.path.join(TMP_DIR, str(req.session_hash))
+    use_ip_adapter = True if ip_img_path and ip_adapt_scale > 0 else False
+    PIPELINE_IP.set_ip_adapter_scale([ip_adapt_scale])
+    img_save_paths = infer_pipe(
+        index_file=f"{output_root}/condition/index.json",
+        controlnet_cond_scale=controlnet_cond_scale,
+        guidance_scale=guidance_scale,
+        strength=strength,
+        num_inference_steps=num_inference_steps,
+        ip_adapt_scale=ip_adapt_scale,
+        ip_img_path=ip_img_path,
+        uid=uid,
+        prompt=prompt,
+        save_dir=f"{output_root}/multi_view",
+        sub_idxs=sub_idxs,
+        pipeline=PIPELINE_IP if use_ip_adapter else PIPELINE,
+        seed=seed,
+    )
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    return img_save_paths + img_save_paths
+
+
+def backproject_texture(
+    mesh_path: str,
+    input_image: str,
+    texture_size: int,
+    uuid: str = "sample",
+    req: gr.Request = None,
+) -> str:
+    output_root = os.path.join(TMP_DIR, str(req.session_hash))
+    output_dir = os.path.join(output_root, "texture_mesh")
+    os.makedirs(output_dir, exist_ok=True)
+    command = [
+        "backproject-cli",
+        "--mesh_path",
+        mesh_path,
+        "--input_image",
+        input_image,
+        "--output_root",
+        output_dir,
+        "--uuid",
+        f"{uuid}",
+        "--texture_size",
+        str(texture_size),
+        "--skip_fix_mesh",
+    ]
+
+    _ = subprocess.run(
+        command, capture_output=True, text=True, encoding="utf-8"
+    )
+    output_obj_mesh = os.path.join(output_dir, f"{uuid}.obj")
+    output_glb_mesh = os.path.join(output_dir, f"{uuid}.glb")
+    _ = trimesh.load(output_obj_mesh).export(output_glb_mesh)
+
+    zip_file = zip_files(
+        input_paths=[
+            output_glb_mesh,
+            output_obj_mesh,
+            os.path.join(output_dir, "material.mtl"),
+            os.path.join(output_dir, "material_0.png"),
+        ],
+        output_zip=os.path.join(output_dir, f"{uuid}.zip"),
+    )
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    return output_glb_mesh, output_obj_mesh, zip_file
+
+
+@spaces.GPU
+def backproject_texture_v2(
+    mesh_path: str,
+    input_image: str,
+    texture_size: int,
+    enable_delight: bool = True,
+    fix_mesh: bool = False,
+    uuid: str = "sample",
+    req: gr.Request = None,
+) -> str:
+    output_root = os.path.join(TMP_DIR, str(req.session_hash))
+    output_dir = os.path.join(output_root, "texture_mesh")
+    os.makedirs(output_dir, exist_ok=True)
+
+    textured_mesh = backproject_api(
+        delight_model=DELIGHT,
+        imagesr_model=IMAGESR_MODEL,
+        color_path=input_image,
+        mesh_path=mesh_path,
+        output_path=f"{output_dir}/{uuid}.obj",
+        skip_fix_mesh=not fix_mesh,
+        delight=enable_delight,
+        texture_wh=[texture_size, texture_size],
+    )
+
+    output_obj_mesh = os.path.join(output_dir, f"{uuid}.obj")
+    output_glb_mesh = os.path.join(output_dir, f"{uuid}.glb")
+    _ = textured_mesh.export(output_glb_mesh)
+
+    zip_file = zip_files(
+        input_paths=[
+            output_glb_mesh,
+            output_obj_mesh,
+            os.path.join(output_dir, "material.mtl"),
+            os.path.join(output_dir, "material_0.png"),
+        ],
+        output_zip=os.path.join(output_dir, f"{uuid}.zip"),
+    )
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    return output_glb_mesh, output_obj_mesh, zip_file
+
+
+@spaces.GPU
+def render_result_video(
+    mesh_path: str, video_size: int, req: gr.Request, uuid: str = ""
+) -> str:
+    output_root = os.path.join(TMP_DIR, str(req.session_hash))
+    output_dir = os.path.join(output_root, "texture_mesh")
+
+    _ = render_api(
+        mesh_path=mesh_path,
+        output_root=output_dir,
+        num_images=90,
+        elevation=[20],
+        with_mtl=True,
+        pbr_light_factor=1,
+        uuid=str(uuid),
+        gen_color_mp4=True,
+        gen_glonormal_mp4=True,
+        distance=5.5,
+        resolution_hw=(video_size, video_size),
+    )
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    return f"{output_dir}/color.mp4"

embodied_gen/data/asset_converter.py

@@ -0,0 +1,663 @@
+from __future__ import annotations
+
+import logging
+import os
+import xml.etree.ElementTree as ET
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from glob import glob
+from shutil import copy
+
+import trimesh
+from scipy.spatial.transform import Rotation
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "AssetConverterFactory",
+    "AssetType",
+    "MeshtoMJCFConverter",
+    "MeshtoUSDConverter",
+    "URDFtoUSDConverter",
+]
+
+
+@dataclass
+class AssetType(str):
+    """Asset type enumeration."""
+
+    MJCF = "mjcf"
+    USD = "usd"
+    URDF = "urdf"
+    MESH = "mesh"
+
+
+class AssetConverterBase(ABC):
+    """Converter abstract base class."""
+
+    @abstractmethod
+    def convert(self, urdf_path: str, output_path: str, **kwargs) -> str:
+        pass
+
+    def transform_mesh(
+        self, input_mesh: str, output_mesh: str, mesh_origin: ET.Element
+    ) -> None:
+        """Apply transform to the mesh based on the origin element in URDF."""
+        mesh = trimesh.load(input_mesh)
+        rpy = list(map(float, mesh_origin.get("rpy").split(" ")))
+        rotation = Rotation.from_euler("xyz", rpy, degrees=False)
+        offset = list(map(float, mesh_origin.get("xyz").split(" ")))
+        mesh.vertices = (mesh.vertices @ rotation.as_matrix().T) + offset
+
+        os.makedirs(os.path.dirname(output_mesh), exist_ok=True)
+        _ = mesh.export(output_mesh)
+
+        return
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return False
+
+
+class MeshtoMJCFConverter(AssetConverterBase):
+    """Convert URDF files into MJCF format."""
+
+    def __init__(
+        self,
+        **kwargs,
+    ) -> None:
+        self.kwargs = kwargs
+
+    def _copy_asset_file(self, src: str, dst: str) -> None:
+        if os.path.exists(dst):
+            return
+        os.makedirs(os.path.dirname(dst), exist_ok=True)
+        copy(src, dst)
+
+    def add_geometry(
+        self,
+        mujoco_element: ET.Element,
+        link: ET.Element,
+        body: ET.Element,
+        tag: str,
+        input_dir: str,
+        output_dir: str,
+        mesh_name: str,
+        material: ET.Element | None = None,
+        is_collision: bool = False,
+    ) -> None:
+        """Add geometry to the MJCF body from the URDF link."""
+        element = link.find(tag)
+        geometry = element.find("geometry")
+        mesh = geometry.find("mesh")
+        filename = mesh.get("filename")
+        scale = mesh.get("scale", "1.0 1.0 1.0")
+
+        mesh_asset = ET.SubElement(
+            mujoco_element, "mesh", name=mesh_name, file=filename, scale=scale
+        )
+        geom = ET.SubElement(body, "geom", type="mesh", mesh=mesh_name)
+
+        self._copy_asset_file(
+            f"{input_dir}/{filename}",
+            f"{output_dir}/{filename}",
+        )
+
+        # Preprocess the mesh by applying rotation.
+        input_mesh = f"{input_dir}/{filename}"
+        output_mesh = f"{output_dir}/{filename}"
+        mesh_origin = element.find("origin")
+        if mesh_origin is not None:
+            self.transform_mesh(input_mesh, output_mesh, mesh_origin)
+
+        if material is not None:
+            geom.set("material", material.get("name"))
+
+        if is_collision:
+            geom.set("contype", "1")
+            geom.set("conaffinity", "1")
+            geom.set("rgba", "1 1 1 0")
+
+    def add_materials(
+        self,
+        mujoco_element: ET.Element,
+        link: ET.Element,
+        tag: str,
+        input_dir: str,
+        output_dir: str,
+        name: str,
+        reflectance: float = 0.2,
+    ) -> ET.Element:
+        """Add materials to the MJCF asset from the URDF link."""
+        element = link.find(tag)
+        geometry = element.find("geometry")
+        mesh = geometry.find("mesh")
+        filename = mesh.get("filename")
+        dirname = os.path.dirname(filename)
+
+        material = ET.SubElement(
+            mujoco_element,
+            "material",
+            name=f"material_{name}",
+            texture=f"texture_{name}",
+            reflectance=str(reflectance),
+        )
+
+        for path in glob(f"{input_dir}/{dirname}/*.png"):
+            file_name = os.path.basename(path)
+            self._copy_asset_file(
+                path,
+                f"{output_dir}/{dirname}/{file_name}",
+            )
+            ET.SubElement(
+                mujoco_element,
+                "texture",
+                name=f"texture_{name}_{os.path.splitext(file_name)[0]}",
+                type="2d",
+                file=f"{dirname}/{file_name}",
+            )
+
+        return material
+
+    def convert(self, urdf_path: str, mjcf_path: str):
+        """Convert a URDF file to MJCF format."""
+        tree = ET.parse(urdf_path)
+        root = tree.getroot()
+
+        mujoco_struct = ET.Element("mujoco")
+        mujoco_struct.set("model", root.get("name"))
+        mujoco_asset = ET.SubElement(mujoco_struct, "asset")
+        mujoco_worldbody = ET.SubElement(mujoco_struct, "worldbody")
+
+        input_dir = os.path.dirname(urdf_path)
+        output_dir = os.path.dirname(mjcf_path)
+        os.makedirs(output_dir, exist_ok=True)
+        for idx, link in enumerate(root.findall("link")):
+            link_name = link.get("name", "unnamed_link")
+            body = ET.SubElement(mujoco_worldbody, "body", name=link_name)
+
+            material = self.add_materials(
+                mujoco_asset,
+                link,
+                "visual",
+                input_dir,
+                output_dir,
+                name=str(idx),
+            )
+            self.add_geometry(
+                mujoco_asset,
+                link,
+                body,
+                "visual",
+                input_dir,
+                output_dir,
+                f"visual_mesh_{idx}",
+                material,
+            )
+            self.add_geometry(
+                mujoco_asset,
+                link,
+                body,
+                "collision",
+                input_dir,
+                output_dir,
+                f"collision_mesh_{idx}",
+                is_collision=True,
+            )
+
+        tree = ET.ElementTree(mujoco_struct)
+        ET.indent(tree, space="  ", level=0)
+
+        tree.write(mjcf_path, encoding="utf-8", xml_declaration=True)
+        logger.info(f"Successfully converted {urdf_path} → {mjcf_path}")
+
+
+class URDFtoMJCFConverter(MeshtoMJCFConverter):
+    """Convert URDF files with joints to MJCF format, handling transformations from joints."""
+
+    def add_materials(
+        self,
+        mujoco_element: ET.Element,
+        link: ET.Element,
+        tag: str,
+        input_dir: str,
+        output_dir: str,
+        name: str,
+        reflectance: float = 0.2,
+    ) -> ET.Element:
+        """Add materials to the MJCF asset from the URDF link."""
+        element = link.find(tag)
+        geometry = element.find("geometry")
+        mesh = geometry.find("mesh")
+        filename = mesh.get("filename")
+        dirname = os.path.dirname(filename)
+
+        diffuse_texture = None
+        for path in glob(f"{input_dir}/{dirname}/*.png"):
+            file_name = os.path.basename(path)
+            self._copy_asset_file(
+                path,
+                f"{output_dir}/{dirname}/{file_name}",
+            )
+            texture_name = f"texture_{name}_{os.path.splitext(file_name)[0]}"
+            ET.SubElement(
+                mujoco_element,
+                "texture",
+                name=texture_name,
+                type="2d",
+                file=f"{dirname}/{file_name}",
+            )
+            if "diffuse" in file_name.lower():
+                diffuse_texture = texture_name
+
+        if diffuse_texture is None:
+            return None
+
+        material = ET.SubElement(
+            mujoco_element,
+            "material",
+            name=f"material_{name}",
+            texture=diffuse_texture,
+            reflectance=str(reflectance),
+        )
+
+        return material
+
+    def convert(self, urdf_path: str, mjcf_path: str, **kwargs) -> str:
+        """Convert a URDF file with joints to MJCF format."""
+        tree = ET.parse(urdf_path)
+        root = tree.getroot()
+
+        mujoco_struct = ET.Element("mujoco")
+        mujoco_struct.set("model", root.get("name"))
+        mujoco_asset = ET.SubElement(mujoco_struct, "asset")
+        mujoco_worldbody = ET.SubElement(mujoco_struct, "worldbody")
+
+        input_dir = os.path.dirname(urdf_path)
+        output_dir = os.path.dirname(mjcf_path)
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Create a dictionary to store body elements for each link
+        body_dict = {}
+
+        # Process all links first
+        for idx, link in enumerate(root.findall("link")):
+            link_name = link.get("name", f"unnamed_link_{idx}")
+            body = ET.SubElement(mujoco_worldbody, "body", name=link_name)
+            body_dict[link_name] = body
+
+            # Add materials and geometry
+            visual_element = link.find("visual")
+            if visual_element is not None:
+                material = self.add_materials(
+                    mujoco_asset,
+                    link,
+                    "visual",
+                    input_dir,
+                    output_dir,
+                    name=str(idx),
+                )
+                self.add_geometry(
+                    mujoco_asset,
+                    link,
+                    body,
+                    "visual",
+                    input_dir,
+                    output_dir,
+                    f"visual_mesh_{idx}",
+                    material,
+                )
+
+            collision_element = link.find("collision")
+            if collision_element is not None:
+                self.add_geometry(
+                    mujoco_asset,
+                    link,
+                    body,
+                    "collision",
+                    input_dir,
+                    output_dir,
+                    f"collision_mesh_{idx}",
+                    is_collision=True,
+                )
+
+        # Process joints to set transformations and hierarchy
+        for joint in root.findall("joint"):
+            joint_type = joint.get("type")
+            if joint_type != "fixed":
+                logger.warning(
+                    f"Skipping non-fixed joint: {joint.get('name')}"
+                )
+                continue
+
+            parent_link = joint.find("parent").get("link")
+            child_link = joint.find("child").get("link")
+            origin = joint.find("origin")
+
+            if parent_link not in body_dict or child_link not in body_dict:
+                logger.warning(
+                    f"Parent or child link not found for joint: {joint.get('name')}"
+                )
+                continue
+
+            # Move child body under parent body in MJCF hierarchy
+            child_body = body_dict[child_link]
+            mujoco_worldbody.remove(child_body)
+            parent_body = body_dict[parent_link]
+            parent_body.append(child_body)
+
+            # Apply joint origin transformation to child body
+            if origin is not None:
+                xyz = origin.get("xyz", "0 0 0")
+                rpy = origin.get("rpy", "0 0 0")
+                child_body.set("pos", xyz)
+                # Convert rpy to MJCF euler format (degrees)
+                rpy_floats = list(map(float, rpy.split()))
+                rotation = Rotation.from_euler(
+                    "xyz", rpy_floats, degrees=False
+                )
+                euler_deg = rotation.as_euler("xyz", degrees=True)
+                child_body.set(
+                    "euler", f"{euler_deg[0]} {euler_deg[1]} {euler_deg[2]}"
+                )
+
+        tree = ET.ElementTree(mujoco_struct)
+        ET.indent(tree, space="  ", level=0)
+        tree.write(mjcf_path, encoding="utf-8", xml_declaration=True)
+        logger.info(f"Successfully converted {urdf_path} → {mjcf_path}")
+
+        return mjcf_path
+
+
+class MeshtoUSDConverter(AssetConverterBase):
+    """Convert Mesh file from URDF into USD format."""
+
+    DEFAULT_BIND_APIS = [
+        "MaterialBindingAPI",
+        "PhysicsMeshCollisionAPI",
+        "PhysicsCollisionAPI",
+        "PhysxCollisionAPI",
+        "PhysicsMassAPI",
+        "PhysicsRigidBodyAPI",
+        "PhysxRigidBodyAPI",
+    ]
+
+    def __init__(
+        self,
+        force_usd_conversion: bool = True,
+        make_instanceable: bool = False,
+        simulation_app=None,
+        **kwargs,
+    ):
+        self.usd_parms = dict(
+            force_usd_conversion=force_usd_conversion,
+            make_instanceable=make_instanceable,
+            **kwargs,
+        )
+        if simulation_app is not None:
+            self.simulation_app = simulation_app
+
+    def __enter__(self):
+        from isaaclab.app import AppLauncher
+
+        if not hasattr(self, "simulation_app"):
+            launch_args = dict(
+                headless=True,
+                no_splash=True,
+                fast_shutdown=True,
+                disable_gpu=True,
+            )
+            self.app_launcher = AppLauncher(launch_args)
+            self.simulation_app = self.app_launcher.app
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # Close the simulation app if it was created here
+        if hasattr(self, "app_launcher"):
+            self.simulation_app.close()
+
+        if exc_val is not None:
+            logger.error(f"Exception occurred: {exc_val}.")
+
+        return False
+
+    def convert(self, urdf_path: str, output_file: str):
+        """Convert a URDF file to USD and post-process collision meshes."""
+        from isaaclab.sim.converters import MeshConverter, MeshConverterCfg
+        from pxr import PhysxSchema, Sdf, Usd, UsdShade
+
+        tree = ET.parse(urdf_path)
+        root = tree.getroot()
+        mesh_file = root.find("link/visual/geometry/mesh").get("filename")
+        input_mesh = os.path.join(os.path.dirname(urdf_path), mesh_file)
+        output_dir = os.path.abspath(os.path.dirname(output_file))
+        output_mesh = f"{output_dir}/mesh/{os.path.basename(mesh_file)}"
+        mesh_origin = root.find("link/visual/origin")
+        if mesh_origin is not None:
+            self.transform_mesh(input_mesh, output_mesh, mesh_origin)
+
+        cfg = MeshConverterCfg(
+            asset_path=output_mesh,
+            usd_dir=output_dir,
+            usd_file_name=os.path.basename(output_file),
+            **self.usd_parms,
+        )
+        urdf_converter = MeshConverter(cfg)
+        usd_path = urdf_converter.usd_path
+
+        stage = Usd.Stage.Open(usd_path)
+        layer = stage.GetRootLayer()
+        with Usd.EditContext(stage, layer):
+            for prim in stage.Traverse():
+                # Change texture path to relative path.
+                if prim.GetName() == "material_0":
+                    shader = UsdShade.Shader(prim).GetInput("diffuse_texture")
+                    if shader.Get() is not None:
+                        relative_path = shader.Get().path.replace(
+                            f"{output_dir}/", ""
+                        )
+                        shader.Set(Sdf.AssetPath(relative_path))
+
+                # Add convex decomposition collision and set ShrinkWrap.
+                elif prim.GetName() == "mesh":
+                    approx_attr = prim.GetAttribute("physics:approximation")
+                    if not approx_attr:
+                        approx_attr = prim.CreateAttribute(
+                            "physics:approximation", Sdf.ValueTypeNames.Token
+                        )
+                    approx_attr.Set("convexDecomposition")
+
+                    physx_conv_api = (
+                        PhysxSchema.PhysxConvexDecompositionCollisionAPI.Apply(
+                            prim
+                        )
+                    )
+                    physx_conv_api.GetShrinkWrapAttr().Set(True)
+
+                    api_schemas = prim.GetMetadata("apiSchemas")
+                    if api_schemas is None:
+                        api_schemas = Sdf.TokenListOp()
+
+                    api_list = list(api_schemas.GetAddedOrExplicitItems())
+                    for api in self.DEFAULT_BIND_APIS:
+                        if api not in api_list:
+                            api_list.append(api)
+
+                    api_schemas.appendedItems = api_list
+                    prim.SetMetadata("apiSchemas", api_schemas)
+
+        layer.Save()
+        logger.info(f"Successfully converted {urdf_path} → {usd_path}")
+
+
+class URDFtoUSDConverter(MeshtoUSDConverter):
+    """Convert URDF files into USD format.
+
+    Args:
+        fix_base (bool): Whether to fix the base link.
+        merge_fixed_joints (bool): Whether to merge fixed joints.
+        make_instanceable (bool): Whether to make prims instanceable.
+        force_usd_conversion (bool): Force conversion to USD.
+        collision_from_visuals (bool): Generate collisions from visuals if not provided.
+    """
+
+    def __init__(
+        self,
+        fix_base: bool = False,
+        merge_fixed_joints: bool = False,
+        make_instanceable: bool = True,
+        force_usd_conversion: bool = True,
+        collision_from_visuals: bool = True,
+        joint_drive=None,
+        rotate_wxyz: tuple[float] | None = None,
+        simulation_app=None,
+        **kwargs,
+    ):
+        self.usd_parms = dict(
+            fix_base=fix_base,
+            merge_fixed_joints=merge_fixed_joints,
+            make_instanceable=make_instanceable,
+            force_usd_conversion=force_usd_conversion,
+            collision_from_visuals=collision_from_visuals,
+            joint_drive=joint_drive,
+            **kwargs,
+        )
+        self.rotate_wxyz = rotate_wxyz
+        if simulation_app is not None:
+            self.simulation_app = simulation_app
+
+    def convert(self, urdf_path: str, output_file: str):
+        """Convert a URDF file to USD and post-process collision meshes."""
+        from isaaclab.sim.converters import UrdfConverter, UrdfConverterCfg
+        from pxr import Gf, PhysxSchema, Sdf, Usd, UsdGeom
+
+        cfg = UrdfConverterCfg(
+            asset_path=urdf_path,
+            usd_dir=os.path.abspath(os.path.dirname(output_file)),
+            usd_file_name=os.path.basename(output_file),
+            **self.usd_parms,
+        )
+
+        urdf_converter = UrdfConverter(cfg)
+        usd_path = urdf_converter.usd_path
+
+        stage = Usd.Stage.Open(usd_path)
+        layer = stage.GetRootLayer()
+        with Usd.EditContext(stage, layer):
+            for prim in stage.Traverse():
+                if prim.GetName() == "collisions":
+                    approx_attr = prim.GetAttribute("physics:approximation")
+                    if not approx_attr:
+                        approx_attr = prim.CreateAttribute(
+                            "physics:approximation", Sdf.ValueTypeNames.Token
+                        )
+                    approx_attr.Set("convexDecomposition")
+
+                    physx_conv_api = (
+                        PhysxSchema.PhysxConvexDecompositionCollisionAPI.Apply(
+                            prim
+                        )
+                    )
+                    physx_conv_api.GetShrinkWrapAttr().Set(True)
+
+                    api_schemas = prim.GetMetadata("apiSchemas")
+                    if api_schemas is None:
+                        api_schemas = Sdf.TokenListOp()
+
+                    api_list = list(api_schemas.GetAddedOrExplicitItems())
+                    for api in self.DEFAULT_BIND_APIS:
+                        if api not in api_list:
+                            api_list.append(api)
+
+                    api_schemas.appendedItems = api_list
+                    prim.SetMetadata("apiSchemas", api_schemas)
+
+        if self.rotate_wxyz is not None:
+            inner_prim = next(
+                p
+                for p in stage.GetDefaultPrim().GetChildren()
+                if p.IsA(UsdGeom.Xform)
+            )
+            xformable = UsdGeom.Xformable(inner_prim)
+            xformable.ClearXformOpOrder()
+            orient_op = xformable.AddOrientOp(UsdGeom.XformOp.PrecisionDouble)
+            orient_op.Set(Gf.Quatd(*self.rotate_wxyz))
+
+        layer.Save()
+        logger.info(f"Successfully converted {urdf_path} → {usd_path}")
+
+
+class AssetConverterFactory:
+    """Factory class for creating asset converters based on target and source types."""
+
+    @staticmethod
+    def create(
+        target_type: AssetType, source_type: AssetType = "urdf", **kwargs
+    ) -> AssetConverterBase:
+        """Create an asset converter instance based on target and source types."""
+        if target_type == AssetType.MJCF and source_type == AssetType.URDF:
+            converter = MeshtoMJCFConverter(**kwargs)
+        elif target_type == AssetType.USD and source_type == AssetType.URDF:
+            converter = URDFtoUSDConverter(**kwargs)
+        elif target_type == AssetType.USD and source_type == AssetType.MESH:
+            converter = MeshtoUSDConverter(**kwargs)
+        else:
+            raise ValueError(
+                f"Unsupported converter type: {source_type} -> {target_type}."
+            )
+
+        return converter
+
+
+if __name__ == "__main__":
+    # # target_asset_type = AssetType.MJCF
+    # target_asset_type = AssetType.USD
+
+    # urdf_paths = [
+    #     "outputs/embodiedgen_assets/demo_assets/remote_control/result/remote_control.urdf",
+    # ]
+
+    # if target_asset_type == AssetType.MJCF:
+    #     output_files = [
+    #         "outputs/embodiedgen_assets/demo_assets/remote_control/mjcf/remote_control.mjcf",
+    #     ]
+    #     asset_converter = AssetConverterFactory.create(
+    #         target_type=AssetType.MJCF,
+    #         source_type=AssetType.URDF,
+    #     )
+
+    # elif target_asset_type == AssetType.USD:
+    #     output_files = [
+    #         "outputs/embodiedgen_assets/demo_assets/remote_control/usd/remote_control.usd",
+    #     ]
+    #     asset_converter = AssetConverterFactory.create(
+    #         target_type=AssetType.USD,
+    #         source_type=AssetType.MESH,
+    #     )
+
+    # with asset_converter:
+    #     for urdf_path, output_file in zip(urdf_paths, output_files):
+    #         asset_converter.convert(urdf_path, output_file)
+
+    # urdf_path = "outputs/embodiedgen_assets/demo_assets/remote_control/result/remote_control.urdf"
+    # output_file = "outputs/embodiedgen_assets/demo_assets/remote_control/usd/remote_control.usd"
+
+    # asset_converter = AssetConverterFactory.create(
+    #     target_type=AssetType.USD,
+    #     source_type=AssetType.URDF,
+    #     rotate_wxyz=(0.7071, 0.7071, 0, 0),  # rotate 90 deg around the X-axis
+    # )
+
+    # with asset_converter:
+    #     asset_converter.convert(urdf_path, output_file)
+
+    urdf_path = "/home/users/xinjie.wang/xinjie/infinigen/outputs/exports/kitchen_simple_solve_nos_i_urdf/export_scene/scene.urdf"
+    output_file = "/home/users/xinjie.wang/xinjie/infinigen/outputs/exports/kitchen_simple_solve_nos_i_urdf/mjcf/scene.urdf"
+    asset_converter = URDFtoMJCFConverter()
+    with asset_converter:
+        asset_converter.convert(urdf_path, output_file)

embodied_gen/data/backproject.py

@@ -0,0 +1,518 @@
+# 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.
+
+
+import argparse
+import logging
+import math
+import os
+from typing import List, Literal, Union
+
+import cv2
+import numpy as np
+import nvdiffrast.torch as dr
+import torch
+import trimesh
+import utils3d
+import xatlas
+from tqdm import tqdm
+from embodied_gen.data.mesh_operator import MeshFixer
+from embodied_gen.data.utils import (
+    CameraSetting,
+    get_images_from_grid,
+    init_kal_camera,
+    normalize_vertices_array,
+    post_process_texture,
+    save_mesh_with_mtl,
+)
+from embodied_gen.models.delight_model import DelightingModel
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+class TextureBaker(object):
+    """Baking textures onto a mesh from multiple observations.
+
+    This class take 3D mesh data, camera settings and texture baking parameters
+    to generate texture map by projecting images to the mesh from diff views.
+    It supports both a fast texture baking approach and a more optimized method
+    with total variation regularization.
+
+    Attributes:
+        vertices (torch.Tensor): The vertices of the mesh.
+        faces (torch.Tensor): The faces of the mesh, defined by vertex indices.
+        uvs (torch.Tensor): The UV coordinates of the mesh.
+        camera_params (CameraSetting): Camera setting (intrinsics, extrinsics).
+        device (str): The device to run computations on ("cpu" or "cuda").
+        w2cs (torch.Tensor): World-to-camera transformation matrices.
+        projections (torch.Tensor): Camera projection matrices.
+
+    Example:
+        >>> vertices, faces, uvs = TextureBaker.parametrize_mesh(vertices, faces)  # noqa
+        >>> texture_backer = TextureBaker(vertices, faces, uvs, camera_params)
+        >>> images = get_images_from_grid(args.color_path, image_size)
+        >>> texture = texture_backer.bake_texture(
+        ...     images, texture_size=args.texture_size, mode=args.baker_mode
+        ... )
+        >>> texture = post_process_texture(texture)
+    """
+
+    def __init__(
+        self,
+        vertices: np.ndarray,
+        faces: np.ndarray,
+        uvs: np.ndarray,
+        camera_params: CameraSetting,
+        device: str = "cuda",
+    ) -> None:
+        self.vertices = (
+            torch.tensor(vertices, device=device)
+            if isinstance(vertices, np.ndarray)
+            else vertices.to(device)
+        )
+        self.faces = (
+            torch.tensor(faces.astype(np.int32), device=device)
+            if isinstance(faces, np.ndarray)
+            else faces.to(device)
+        )
+        self.uvs = (
+            torch.tensor(uvs, device=device)
+            if isinstance(uvs, np.ndarray)
+            else uvs.to(device)
+        )
+        self.camera_params = camera_params
+        self.device = device
+
+        camera = init_kal_camera(camera_params)
+        matrix_mv = camera.view_matrix()  # (n_cam 4 4) world2cam
+        matrix_mv = kaolin_to_opencv_view(matrix_mv)
+        matrix_p = (
+            camera.intrinsics.projection_matrix()
+        )  # (n_cam 4 4) cam2pixel
+        self.w2cs = matrix_mv.to(self.device)
+        self.projections = matrix_p.to(self.device)
+
+    @staticmethod
+    def parametrize_mesh(
+        vertices: np.array, faces: np.array
+    ) -> Union[np.array, np.array, np.array]:
+        vmapping, indices, uvs = xatlas.parametrize(vertices, faces)
+
+        vertices = vertices[vmapping]
+        faces = indices
+
+        return vertices, faces, uvs
+
+    def _bake_fast(self, observations, w2cs, projections, texture_size, masks):
+        texture = torch.zeros(
+            (texture_size * texture_size, 3), dtype=torch.float32
+        ).cuda()
+        texture_weights = torch.zeros(
+            (texture_size * texture_size), dtype=torch.float32
+        ).cuda()
+        rastctx = utils3d.torch.RastContext(backend="cuda")
+        for observation, w2c, projection in tqdm(
+            zip(observations, w2cs, projections),
+            total=len(observations),
+            desc="Texture baking (fast)",
+        ):
+            with torch.no_grad():
+                rast = utils3d.torch.rasterize_triangle_faces(
+                    rastctx,
+                    self.vertices[None],
+                    self.faces,
+                    observation.shape[1],
+                    observation.shape[0],
+                    uv=self.uvs[None],
+                    view=w2c,
+                    projection=projection,
+                )
+                uv_map = rast["uv"][0].detach().flip(0)
+                mask = rast["mask"][0].detach().bool() & masks[0]
+
+            # nearest neighbor interpolation
+            uv_map = (uv_map * texture_size).floor().long()
+            obs = observation[mask]
+            uv_map = uv_map[mask]
+            idx = (
+                uv_map[:, 0] + (texture_size - uv_map[:, 1] - 1) * texture_size
+            )
+            texture = texture.scatter_add(
+                0, idx.view(-1, 1).expand(-1, 3), obs
+            )
+            texture_weights = texture_weights.scatter_add(
+                0,
+                idx,
+                torch.ones(
+                    (obs.shape[0]), dtype=torch.float32, device=texture.device
+                ),
+            )
+
+        mask = texture_weights > 0
+        texture[mask] /= texture_weights[mask][:, None]
+        texture = np.clip(
+            texture.reshape(texture_size, texture_size, 3).cpu().numpy() * 255,
+            0,
+            255,
+        ).astype(np.uint8)
+
+        # inpaint
+        mask = (
+            (texture_weights == 0)
+            .cpu()
+            .numpy()
+            .astype(np.uint8)
+            .reshape(texture_size, texture_size)
+        )
+        texture = cv2.inpaint(texture, mask, 3, cv2.INPAINT_TELEA)
+
+        return texture
+
+    def _bake_opt(
+        self,
+        observations,
+        w2cs,
+        projections,
+        texture_size,
+        lambda_tv,
+        masks,
+        total_steps,
+    ):
+        rastctx = utils3d.torch.RastContext(backend="cuda")
+        observations = [observations.flip(0) for observations in observations]
+        masks = [m.flip(0) for m in masks]
+        _uv = []
+        _uv_dr = []
+        for observation, w2c, projection in tqdm(
+            zip(observations, w2cs, projections),
+            total=len(w2cs),
+        ):
+            with torch.no_grad():
+                rast = utils3d.torch.rasterize_triangle_faces(
+                    rastctx,
+                    self.vertices[None],
+                    self.faces,
+                    observation.shape[1],
+                    observation.shape[0],
+                    uv=self.uvs[None],
+                    view=w2c,
+                    projection=projection,
+                )
+                _uv.append(rast["uv"].detach())
+                _uv_dr.append(rast["uv_dr"].detach())
+
+        texture = torch.nn.Parameter(
+            torch.zeros(
+                (1, texture_size, texture_size, 3), dtype=torch.float32
+            ).cuda()
+        )
+        optimizer = torch.optim.Adam([texture], betas=(0.5, 0.9), lr=1e-2)
+
+        def cosine_anealing(step, total_steps, start_lr, end_lr):
+            return end_lr + 0.5 * (start_lr - end_lr) * (
+                1 + np.cos(np.pi * step / total_steps)
+            )
+
+        def tv_loss(texture):
+            return torch.nn.functional.l1_loss(
+                texture[:, :-1, :, :], texture[:, 1:, :, :]
+            ) + torch.nn.functional.l1_loss(
+                texture[:, :, :-1, :], texture[:, :, 1:, :]
+            )
+
+        with tqdm(total=total_steps, desc="Texture baking") as pbar:
+            for step in range(total_steps):
+                optimizer.zero_grad()
+                selected = np.random.randint(0, len(w2cs))
+                uv, uv_dr, observation, mask = (
+                    _uv[selected],
+                    _uv_dr[selected],
+                    observations[selected],
+                    masks[selected],
+                )
+                render = dr.texture(texture, uv, uv_dr)[0]
+                loss = torch.nn.functional.l1_loss(
+                    render[mask], observation[mask]
+                )
+                if lambda_tv > 0:
+                    loss += lambda_tv * tv_loss(texture)
+                loss.backward()
+                optimizer.step()
+
+                optimizer.param_groups[0]["lr"] = cosine_anealing(
+                    step, total_steps, 1e-2, 1e-5
+                )
+                pbar.set_postfix({"loss": loss.item()})
+                pbar.update()
+        texture = np.clip(
+            texture[0].flip(0).detach().cpu().numpy() * 255, 0, 255
+        ).astype(np.uint8)
+        mask = 1 - utils3d.torch.rasterize_triangle_faces(
+            rastctx,
+            (self.uvs * 2 - 1)[None],
+            self.faces,
+            texture_size,
+            texture_size,
+        )["mask"][0].detach().cpu().numpy().astype(np.uint8)
+        texture = cv2.inpaint(texture, mask, 3, cv2.INPAINT_TELEA)
+
+        return texture
+
+    def bake_texture(
+        self,
+        images: List[np.array],
+        texture_size: int = 1024,
+        mode: Literal["fast", "opt"] = "opt",
+        lambda_tv: float = 1e-2,
+        opt_step: int = 2000,
+    ):
+        masks = [np.any(img > 0, axis=-1) for img in images]
+        masks = [torch.tensor(m > 0).bool().to(self.device) for m in masks]
+        images = [
+            torch.tensor(obs / 255.0).float().to(self.device) for obs in images
+        ]
+
+        if mode == "fast":
+            return self._bake_fast(
+                images, self.w2cs, self.projections, texture_size, masks
+            )
+        elif mode == "opt":
+            return self._bake_opt(
+                images,
+                self.w2cs,
+                self.projections,
+                texture_size,
+                lambda_tv,
+                masks,
+                opt_step,
+            )
+        else:
+            raise ValueError(f"Unknown mode: {mode}")
+
+
+def kaolin_to_opencv_view(raw_matrix):
+    R_orig = raw_matrix[:, :3, :3]
+    t_orig = raw_matrix[:, :3, 3]
+
+    R_target = torch.zeros_like(R_orig)
+    R_target[:, :, 0] = R_orig[:, :, 2]
+    R_target[:, :, 1] = R_orig[:, :, 0]
+    R_target[:, :, 2] = R_orig[:, :, 1]
+
+    t_target = t_orig
+
+    target_matrix = (
+        torch.eye(4, device=raw_matrix.device)
+        .unsqueeze(0)
+        .repeat(raw_matrix.size(0), 1, 1)
+    )
+    target_matrix[:, :3, :3] = R_target
+    target_matrix[:, :3, 3] = t_target
+
+    return target_matrix
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Render settings")
+
+    parser.add_argument(
+        "--mesh_path",
+        type=str,
+        nargs="+",
+        required=True,
+        help="Paths to the mesh files for rendering.",
+    )
+    parser.add_argument(
+        "--color_path",
+        type=str,
+        nargs="+",
+        required=True,
+        help="Paths to the mesh files for rendering.",
+    )
+    parser.add_argument(
+        "--output_root",
+        type=str,
+        default="./outputs",
+        help="Root directory for output",
+    )
+    parser.add_argument(
+        "--uuid",
+        type=str,
+        nargs="+",
+        default=None,
+        help="uuid for rendering saving.",
+    )
+    parser.add_argument(
+        "--num_images", type=int, default=6, help="Number of images to render."
+    )
+    parser.add_argument(
+        "--elevation",
+        type=float,
+        nargs="+",
+        default=[20.0, -10.0],
+        help="Elevation angles for the camera (default: [20.0, -10.0])",
+    )
+    parser.add_argument(
+        "--distance",
+        type=float,
+        default=5,
+        help="Camera distance (default: 5)",
+    )
+    parser.add_argument(
+        "--resolution_hw",
+        type=int,
+        nargs=2,
+        default=(512, 512),
+        help="Resolution of the output images (default: (512, 512))",
+    )
+    parser.add_argument(
+        "--fov",
+        type=float,
+        default=30,
+        help="Field of view in degrees (default: 30)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        choices=["cpu", "cuda"],
+        default="cuda",
+        help="Device to run on (default: `cuda`)",
+    )
+    parser.add_argument(
+        "--texture_size",
+        type=int,
+        default=1024,
+        help="Texture size for texture baking (default: 1024)",
+    )
+    parser.add_argument(
+        "--baker_mode",
+        type=str,
+        default="opt",
+        help="Texture baking mode, `fast` or `opt` (default: opt)",
+    )
+    parser.add_argument(
+        "--opt_step",
+        type=int,
+        default=2500,
+        help="Optimization steps for texture baking (default: 2500)",
+    )
+    parser.add_argument(
+        "--mesh_sipmlify_ratio",
+        type=float,
+        default=0.9,
+        help="Mesh simplification ratio (default: 0.9)",
+    )
+    parser.add_argument(
+        "--no_coor_trans",
+        action="store_true",
+        help="Do not transform the asset coordinate system.",
+    )
+    parser.add_argument(
+        "--delight", action="store_true", help="Use delighting model."
+    )
+    parser.add_argument(
+        "--skip_fix_mesh", action="store_true", help="Fix mesh geometry."
+    )
+
+    args = parser.parse_args()
+
+    if args.uuid is None:
+        args.uuid = []
+        for path in args.mesh_path:
+            uuid = os.path.basename(path).split(".")[0]
+            args.uuid.append(uuid)
+
+    return args
+
+
+def entrypoint() -> None:
+    args = parse_args()
+    camera_params = CameraSetting(
+        num_images=args.num_images,
+        elevation=args.elevation,
+        distance=args.distance,
+        resolution_hw=args.resolution_hw,
+        fov=math.radians(args.fov),
+        device=args.device,
+    )
+
+    for mesh_path, uuid, img_path in zip(
+        args.mesh_path, args.uuid, args.color_path
+    ):
+        mesh = trimesh.load(mesh_path)
+        if isinstance(mesh, trimesh.Scene):
+            mesh = mesh.dump(concatenate=True)
+        vertices, scale, center = normalize_vertices_array(mesh.vertices)
+
+        if not args.no_coor_trans:
+            x_rot = np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]])
+            z_rot = np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]])
+            vertices = vertices @ x_rot
+            vertices = vertices @ z_rot
+
+        faces = mesh.faces.astype(np.int32)
+        vertices = vertices.astype(np.float32)
+
+        if not args.skip_fix_mesh:
+            mesh_fixer = MeshFixer(vertices, faces, args.device)
+            vertices, faces = mesh_fixer(
+                filter_ratio=args.mesh_sipmlify_ratio,
+                max_hole_size=0.04,
+                resolution=1024,
+                num_views=1000,
+                norm_mesh_ratio=0.5,
+            )
+
+        vertices, faces, uvs = TextureBaker.parametrize_mesh(vertices, faces)
+        texture_backer = TextureBaker(
+            vertices,
+            faces,
+            uvs,
+            camera_params,
+        )
+        images = get_images_from_grid(
+            img_path, img_size=camera_params.resolution_hw[0]
+        )
+        if args.delight:
+            delight_model = DelightingModel()
+            images = [delight_model(img) for img in images]
+
+        images = [np.array(img) for img in images]
+        texture = texture_backer.bake_texture(
+            images=[img[..., :3] for img in images],
+            texture_size=args.texture_size,
+            mode=args.baker_mode,
+            opt_step=args.opt_step,
+        )
+        texture = post_process_texture(texture)
+
+        if not args.no_coor_trans:
+            vertices = vertices @ np.linalg.inv(z_rot)
+            vertices = vertices @ np.linalg.inv(x_rot)
+        vertices = vertices / scale
+        vertices = vertices + center
+
+        output_path = os.path.join(args.output_root, f"{uuid}.obj")
+        mesh = save_mesh_with_mtl(vertices, faces, uvs, texture, output_path)
+
+    return
+
+
+if __name__ == "__main__":
+    entrypoint()

embodied_gen/data/backproject_v2.py

@@ -0,0 +1,721 @@
+# 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.
+
+
+import argparse
+import logging
+import math
+import os
+
+import cv2
+import numpy as np
+import nvdiffrast.torch as dr
+import spaces
+import torch
+import torch.nn.functional as F
+import trimesh
+import xatlas
+from PIL import Image
+from embodied_gen.data.mesh_operator import MeshFixer
+from embodied_gen.data.utils import (
+    CameraSetting,
+    DiffrastRender,
+    as_list,
+    get_images_from_grid,
+    init_kal_camera,
+    normalize_vertices_array,
+    post_process_texture,
+    save_mesh_with_mtl,
+)
+from embodied_gen.models.delight_model import DelightingModel
+from embodied_gen.models.sr_model import ImageRealESRGAN
+from embodied_gen.utils.process_media import vcat_pil_images
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "TextureBacker",
+]
+
+
+def _transform_vertices(
+    mtx: torch.Tensor, pos: torch.Tensor, keepdim: bool = False
+) -> torch.Tensor:
+    """Transform 3D vertices using a projection matrix."""
+    t_mtx = torch.as_tensor(mtx, device=pos.device, dtype=pos.dtype)
+    if pos.size(-1) == 3:
+        pos = torch.cat([pos, torch.ones_like(pos[..., :1])], dim=-1)
+
+    result = pos @ t_mtx.T
+
+    return result if keepdim else result.unsqueeze(0)
+
+
+def _bilinear_interpolation_scattering(
+    image_h: int, image_w: int, coords: torch.Tensor, values: torch.Tensor
+) -> torch.Tensor:
+    """Bilinear interpolation scattering for grid-based value accumulation."""
+    device = values.device
+    dtype = values.dtype
+    C = values.shape[-1]
+
+    indices = coords * torch.tensor(
+        [image_h - 1, image_w - 1], dtype=dtype, device=device
+    )
+    i, j = indices.unbind(-1)
+
+    i0, j0 = (
+        indices.floor()
+        .long()
+        .clamp(0, image_h - 2)
+        .clamp(0, image_w - 2)
+        .unbind(-1)
+    )
+    i1, j1 = i0 + 1, j0 + 1
+
+    w_i = i - i0.float()
+    w_j = j - j0.float()
+    weights = torch.stack(
+        [(1 - w_i) * (1 - w_j), (1 - w_i) * w_j, w_i * (1 - w_j), w_i * w_j],
+        dim=1,
+    )
+
+    indices_comb = torch.stack(
+        [
+            torch.stack([i0, j0], dim=1),
+            torch.stack([i0, j1], dim=1),
+            torch.stack([i1, j0], dim=1),
+            torch.stack([i1, j1], dim=1),
+        ],
+        dim=1,
+    )
+
+    grid = torch.zeros(image_h, image_w, C, device=device, dtype=dtype)
+    cnt = torch.zeros(image_h, image_w, 1, device=device, dtype=dtype)
+
+    for k in range(4):
+        idx = indices_comb[:, k]
+        w = weights[:, k].unsqueeze(-1)
+
+        stride = torch.tensor([image_w, 1], device=device, dtype=torch.long)
+        flat_idx = (idx * stride).sum(-1)
+
+        grid.view(-1, C).scatter_add_(
+            0, flat_idx.unsqueeze(-1).expand(-1, C), values * w
+        )
+        cnt.view(-1, 1).scatter_add_(0, flat_idx.unsqueeze(-1), w)
+
+    mask = cnt.squeeze(-1) > 0
+    grid[mask] = grid[mask] / cnt[mask].repeat(1, C)
+
+    return grid
+
+
+def _texture_inpaint_smooth(
+    texture: np.ndarray,
+    mask: np.ndarray,
+    vertices: np.ndarray,
+    faces: np.ndarray,
+    uv_map: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Perform texture inpainting using vertex-based color propagation."""
+    image_h, image_w, C = texture.shape
+    N = vertices.shape[0]
+
+    # Initialize vertex data structures
+    vtx_mask = np.zeros(N, dtype=np.float32)
+    vtx_colors = np.zeros((N, C), dtype=np.float32)
+    unprocessed = []
+    adjacency = [[] for _ in range(N)]
+
+    # Build adjacency graph and initial color assignment
+    for face_idx in range(faces.shape[0]):
+        for k in range(3):
+            uv_idx_k = faces[face_idx, k]
+            v_idx = faces[face_idx, k]
+
+            # Convert UV to pixel coordinates with boundary clamping
+            u = np.clip(
+                int(round(uv_map[uv_idx_k, 0] * (image_w - 1))), 0, image_w - 1
+            )
+            v = np.clip(
+                int(round((1.0 - uv_map[uv_idx_k, 1]) * (image_h - 1))),
+                0,
+                image_h - 1,
+            )
+
+            if mask[v, u]:
+                vtx_mask[v_idx] = 1.0
+                vtx_colors[v_idx] = texture[v, u]
+            elif v_idx not in unprocessed:
+                unprocessed.append(v_idx)
+
+            # Build undirected adjacency graph
+            neighbor = faces[face_idx, (k + 1) % 3]
+            if neighbor not in adjacency[v_idx]:
+                adjacency[v_idx].append(neighbor)
+            if v_idx not in adjacency[neighbor]:
+                adjacency[neighbor].append(v_idx)
+
+    # Color propagation with dynamic stopping
+    remaining_iters, prev_count = 2, 0
+    while remaining_iters > 0:
+        current_unprocessed = []
+
+        for v_idx in unprocessed:
+            valid_neighbors = [n for n in adjacency[v_idx] if vtx_mask[n] > 0]
+            if not valid_neighbors:
+                current_unprocessed.append(v_idx)
+                continue
+
+            # Calculate inverse square distance weights
+            neighbors_pos = vertices[valid_neighbors]
+            dist_sq = np.sum((vertices[v_idx] - neighbors_pos) ** 2, axis=1)
+            weights = 1 / np.maximum(dist_sq, 1e-8)
+
+            vtx_colors[v_idx] = np.average(
+                vtx_colors[valid_neighbors], weights=weights, axis=0
+            )
+            vtx_mask[v_idx] = 1.0
+
+        # Update iteration control
+        if len(current_unprocessed) == prev_count:
+            remaining_iters -= 1
+        else:
+            remaining_iters = min(remaining_iters + 1, 2)
+        prev_count = len(current_unprocessed)
+        unprocessed = current_unprocessed
+
+    # Generate output texture
+    inpainted_texture, updated_mask = texture.copy(), mask.copy()
+    for face_idx in range(faces.shape[0]):
+        for k in range(3):
+            v_idx = faces[face_idx, k]
+            if not vtx_mask[v_idx]:
+                continue
+
+            # UV coordinate conversion
+            uv_idx_k = faces[face_idx, k]
+            u = np.clip(
+                int(round(uv_map[uv_idx_k, 0] * (image_w - 1))), 0, image_w - 1
+            )
+            v = np.clip(
+                int(round((1.0 - uv_map[uv_idx_k, 1]) * (image_h - 1))),
+                0,
+                image_h - 1,
+            )
+
+            inpainted_texture[v, u] = vtx_colors[v_idx]
+            updated_mask[v, u] = 255
+
+    return inpainted_texture, updated_mask
+
+
+class TextureBacker:
+    """Texture baking pipeline for multi-view projection and fusion.
+
+    This class performs UV-based texture generation for a 3D mesh using
+    multi-view color images, depth, and normal information. The pipeline
+    includes mesh normalization and UV unwrapping, visibility-aware
+    back-projection, confidence-weighted texture fusion, and inpainting
+    of missing texture regions.
+
+    Args:
+        camera_params (CameraSetting): Camera intrinsics and extrinsics used
+            for rendering each view.
+        view_weights (list[float]): A list of weights for each view, used
+            to blend confidence maps during texture fusion.
+        render_wh (tuple[int, int], optional): Resolution (width, height) for
+            intermediate rendering passes. Defaults to (2048, 2048).
+        texture_wh (tuple[int, int], optional): Output texture resolution
+            (width, height). Defaults to (2048, 2048).
+        bake_angle_thresh (int, optional): Maximum angle (in degrees) between
+            view direction and surface normal for projection to be considered valid.
+            Defaults to 75.
+        mask_thresh (float, optional): Threshold applied to visibility masks
+            during rendering. Defaults to 0.5.
+        smooth_texture (bool, optional): If True, apply post-processing (e.g.,
+            blurring) to the final texture. Defaults to True.
+        inpaint_smooth (bool, optional): If True, apply inpainting to smooth.
+    """
+
+    def __init__(
+        self,
+        camera_params: CameraSetting,
+        view_weights: list[float],
+        render_wh: tuple[int, int] = (2048, 2048),
+        texture_wh: tuple[int, int] = (2048, 2048),
+        bake_angle_thresh: int = 75,
+        mask_thresh: float = 0.5,
+        smooth_texture: bool = True,
+        inpaint_smooth: bool = False,
+    ) -> None:
+        self.camera_params = camera_params
+        self.renderer = None
+        self.view_weights = view_weights
+        self.device = camera_params.device
+        self.render_wh = render_wh
+        self.texture_wh = texture_wh
+        self.mask_thresh = mask_thresh
+        self.smooth_texture = smooth_texture
+        self.inpaint_smooth = inpaint_smooth
+
+        self.bake_angle_thresh = bake_angle_thresh
+        self.bake_unreliable_kernel_size = int(
+            (2 / 512) * max(self.render_wh[0], self.render_wh[1])
+        )
+
+    def _lazy_init_render(self, camera_params, mask_thresh):
+        if self.renderer is None:
+            camera = init_kal_camera(camera_params)
+            mv = camera.view_matrix()  # (n 4 4) world2cam
+            p = camera.intrinsics.projection_matrix()
+            # NOTE: add a negative sign at P[0, 2] as the y axis is flipped in `nvdiffrast` output.  # noqa
+            p[:, 1, 1] = -p[:, 1, 1]
+            self.renderer = DiffrastRender(
+                p_matrix=p,
+                mv_matrix=mv,
+                resolution_hw=camera_params.resolution_hw,
+                context=dr.RasterizeCudaContext(),
+                mask_thresh=mask_thresh,
+                grad_db=False,
+                device=self.device,
+                antialias_mask=True,
+            )
+
+    def load_mesh(self, mesh: trimesh.Trimesh) -> trimesh.Trimesh:
+        mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
+        self.scale, self.center = scale, center
+
+        vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
+        uvs[:, 1] = 1 - uvs[:, 1]
+        mesh.vertices = mesh.vertices[vmapping]
+        mesh.faces = indices
+        mesh.visual.uv = uvs
+
+        return mesh
+
+    def get_mesh_np_attrs(
+        self,
+        mesh: trimesh.Trimesh,
+        scale: float = None,
+        center: np.ndarray = None,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        vertices = mesh.vertices.copy()
+        faces = mesh.faces.copy()
+        uv_map = mesh.visual.uv.copy()
+        uv_map[:, 1] = 1.0 - uv_map[:, 1]
+
+        if scale is not None:
+            vertices = vertices / scale
+        if center is not None:
+            vertices = vertices + center
+
+        return vertices, faces, uv_map
+
+    def _render_depth_edges(self, depth_image: torch.Tensor) -> torch.Tensor:
+        depth_image_np = depth_image.cpu().numpy()
+        depth_image_np = (depth_image_np * 255).astype(np.uint8)
+        depth_edges = cv2.Canny(depth_image_np, 30, 80)
+        sketch_image = (
+            torch.from_numpy(depth_edges).to(depth_image.device).float() / 255
+        )
+        sketch_image = sketch_image.unsqueeze(-1)
+
+        return sketch_image
+
+    def compute_enhanced_viewnormal(
+        self, mv_mtx: torch.Tensor, vertices: torch.Tensor, faces: torch.Tensor
+    ) -> torch.Tensor:
+        rast, _ = self.renderer.compute_dr_raster(vertices, faces)
+        rendered_view_normals = []
+        for idx in range(len(mv_mtx)):
+            pos_cam = _transform_vertices(mv_mtx[idx], vertices, keepdim=True)
+            pos_cam = pos_cam[:, :3] / pos_cam[:, 3:]
+            v0, v1, v2 = (pos_cam[faces[:, i]] for i in range(3))
+            face_norm = F.normalize(
+                torch.cross(v1 - v0, v2 - v0, dim=-1), dim=-1
+            )
+            vertex_norm = (
+                torch.from_numpy(
+                    trimesh.geometry.mean_vertex_normals(
+                        len(pos_cam), faces.cpu(), face_norm.cpu()
+                    )
+                )
+                .to(vertices.device)
+                .contiguous()
+            )
+            im_base_normals, _ = dr.interpolate(
+                vertex_norm[None, ...].float(),
+                rast[idx : idx + 1],
+                faces.to(torch.int32),
+            )
+            rendered_view_normals.append(im_base_normals)
+
+        rendered_view_normals = torch.cat(rendered_view_normals, dim=0)
+
+        return rendered_view_normals
+
+    def back_project(
+        self, image, vis_mask, depth, normal, uv
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        image = np.array(image)
+        image = torch.as_tensor(image, device=self.device, dtype=torch.float32)
+        if image.ndim == 2:
+            image = image.unsqueeze(-1)
+        image = image / 255
+
+        depth_inv = (1.0 - depth) * vis_mask
+        sketch_image = self._render_depth_edges(depth_inv)
+
+        cos = F.cosine_similarity(
+            torch.tensor([[0, 0, 1]], device=self.device),
+            normal.view(-1, 3),
+        ).view_as(normal[..., :1])
+        cos[cos < np.cos(np.radians(self.bake_angle_thresh))] = 0
+
+        k = self.bake_unreliable_kernel_size * 2 + 1
+        kernel = torch.ones((1, 1, k, k), device=self.device)
+
+        vis_mask = vis_mask.permute(2, 0, 1).unsqueeze(0).float()
+        vis_mask = F.conv2d(
+            1.0 - vis_mask,
+            kernel,
+            padding=k // 2,
+        )
+        vis_mask = 1.0 - (vis_mask > 0).float()
+        vis_mask = vis_mask.squeeze(0).permute(1, 2, 0)
+
+        sketch_image = sketch_image.permute(2, 0, 1).unsqueeze(0)
+        sketch_image = F.conv2d(sketch_image, kernel, padding=k // 2)
+        sketch_image = (sketch_image > 0).float()
+        sketch_image = sketch_image.squeeze(0).permute(1, 2, 0)
+        vis_mask = vis_mask * (sketch_image < 0.5)
+
+        cos[vis_mask == 0] = 0
+        valid_pixels = (vis_mask != 0).view(-1)
+
+        return (
+            self._scatter_texture(uv, image, valid_pixels),
+            self._scatter_texture(uv, cos, valid_pixels),
+        )
+
+    def _scatter_texture(self, uv, data, mask):
+        def __filter_data(data, mask):
+            return data.view(-1, data.shape[-1])[mask]
+
+        return _bilinear_interpolation_scattering(
+            self.texture_wh[1],
+            self.texture_wh[0],
+            __filter_data(uv, mask)[..., [1, 0]],
+            __filter_data(data, mask),
+        )
+
+    @torch.no_grad()
+    def fast_bake_texture(
+        self, textures: list[torch.Tensor], confidence_maps: list[torch.Tensor]
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        channel = textures[0].shape[-1]
+        texture_merge = torch.zeros(self.texture_wh + [channel]).to(
+            self.device
+        )
+        trust_map_merge = torch.zeros(self.texture_wh + [1]).to(self.device)
+        for texture, cos_map in zip(textures, confidence_maps):
+            view_sum = (cos_map > 0).sum()
+            painted_sum = ((cos_map > 0) * (trust_map_merge > 0)).sum()
+            if painted_sum / view_sum > 0.99:
+                continue
+            texture_merge += texture * cos_map
+            trust_map_merge += cos_map
+        texture_merge = texture_merge / torch.clamp(trust_map_merge, min=1e-8)
+
+        return texture_merge, trust_map_merge > 1e-8
+
+    def uv_inpaint(
+        self, mesh: trimesh.Trimesh, texture: np.ndarray, mask: np.ndarray
+    ) -> np.ndarray:
+        if self.inpaint_smooth:
+            vertices, faces, uv_map = self.get_mesh_np_attrs(mesh)
+            texture, mask = _texture_inpaint_smooth(
+                texture, mask, vertices, faces, uv_map
+            )
+
+        texture = texture.clip(0, 1)
+        texture = cv2.inpaint(
+            (texture * 255).astype(np.uint8),
+            255 - mask,
+            3,
+            cv2.INPAINT_NS,
+        )
+
+        return texture
+
+    @spaces.GPU
+    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)
+        view_normals = self.compute_enhanced_viewnormal(
+            self.renderer.mv_mtx, vertices, faces
+        )
+
+        textures, weighted_cos_maps = [], []
+        for color, mask, dep, normal, uv, weight in zip(
+            colors,
+            masks,
+            norm_deps,
+            view_normals,
+            render_uvs,
+            self.view_weights,
+        ):
+            texture, cos_map = self.back_project(color, mask, dep, normal, uv)
+            textures.append(texture)
+            weighted_cos_maps.append(weight * (cos_map**4))
+
+        texture, mask = self.fast_bake_texture(textures, weighted_cos_maps)
+
+        texture_np = texture.cpu().numpy()
+        mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)
+
+        return texture_np, mask_np
+
+    def __call__(
+        self,
+        colors: list[Image.Image],
+        mesh: trimesh.Trimesh,
+        output_path: str,
+    ) -> trimesh.Trimesh:
+        """Runs the texture baking and exports the textured mesh.
+
+        Args:
+            colors (list[Image.Image]): List of input view images.
+            mesh (trimesh.Trimesh): Input mesh to be textured.
+            output_path (str): Path to save the output textured mesh (.obj or .glb).
+
+        Returns:
+            trimesh.Trimesh: The textured mesh with UV and texture image.
+        """
+        mesh = self.load_mesh(mesh)
+        texture_np, mask_np = self.compute_texture(colors, mesh)
+
+        texture_np = self.uv_inpaint(mesh, texture_np, mask_np)
+        if self.smooth_texture:
+            texture_np = post_process_texture(texture_np)
+
+        vertices, faces, uv_map = self.get_mesh_np_attrs(
+            mesh, self.scale, self.center
+        )
+        textured_mesh = save_mesh_with_mtl(
+            vertices, faces, uv_map, texture_np, output_path
+        )
+
+        return textured_mesh
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Backproject texture")
+    parser.add_argument(
+        "--color_path",
+        nargs="+",
+        type=str,
+        help="Multiview color image in grid file paths",
+    )
+    parser.add_argument(
+        "--mesh_path",
+        type=str,
+        help="Mesh path, .obj, .glb or .ply",
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        help="Output mesh path with suffix",
+    )
+    parser.add_argument(
+        "--num_images", type=int, default=6, help="Number of images to render."
+    )
+    parser.add_argument(
+        "--elevation",
+        nargs="+",
+        type=float,
+        default=[20.0, -10.0],
+        help="Elevation angles for the camera (default: [20.0, -10.0])",
+    )
+    parser.add_argument(
+        "--distance",
+        type=float,
+        default=5,
+        help="Camera distance (default: 5)",
+    )
+    parser.add_argument(
+        "--resolution_hw",
+        type=int,
+        nargs=2,
+        default=(2048, 2048),
+        help="Resolution of the output images (default: (2048, 2048))",
+    )
+    parser.add_argument(
+        "--fov",
+        type=float,
+        default=30,
+        help="Field of view in degrees (default: 30)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        choices=["cpu", "cuda"],
+        default="cuda",
+        help="Device to run on (default: `cuda`)",
+    )
+    parser.add_argument(
+        "--skip_fix_mesh", action="store_true", help="Fix mesh geometry."
+    )
+    parser.add_argument(
+        "--texture_wh",
+        nargs=2,
+        type=int,
+        default=[2048, 2048],
+        help="Texture resolution width and height",
+    )
+    parser.add_argument(
+        "--mesh_sipmlify_ratio",
+        type=float,
+        default=0.9,
+        help="Mesh simplification ratio (default: 0.9)",
+    )
+    parser.add_argument(
+        "--delight", action="store_true", help="Use delighting model."
+    )
+    parser.add_argument(
+        "--no_smooth_texture",
+        action="store_true",
+        help="Do not smooth the texture.",
+    )
+    parser.add_argument(
+        "--save_glb_path", type=str, default=None, help="Save glb path."
+    )
+    parser.add_argument(
+        "--no_save_delight_img",
+        action="store_true",
+        help="Disable saving delight image",
+    )
+    parser.add_argument("--n_max_faces", type=int, default=30000)
+    args, unknown = parser.parse_known_args()
+
+    return args
+
+
+def entrypoint(
+    delight_model: DelightingModel = None,
+    imagesr_model: ImageRealESRGAN = None,
+    **kwargs,
+) -> trimesh.Trimesh:
+    args = parse_args()
+    for k, v in kwargs.items():
+        if hasattr(args, k) and v is not None:
+            setattr(args, k, v)
+
+    # Setup camera parameters.
+    camera_params = CameraSetting(
+        num_images=args.num_images,
+        elevation=args.elevation,
+        distance=args.distance,
+        resolution_hw=args.resolution_hw,
+        fov=math.radians(args.fov),
+        device=args.device,
+    )
+
+    args.color_path = as_list(args.color_path)
+    if args.delight and delight_model is None:
+        delight_model = DelightingModel()
+
+    color_grid = [Image.open(color_path) for color_path in args.color_path]
+    color_grid = vcat_pil_images(color_grid, image_mode="RGBA")
+    if args.delight:
+        color_grid = delight_model(color_grid)
+        if not args.no_save_delight_img:
+            save_dir = os.path.dirname(args.output_path)
+            os.makedirs(save_dir, exist_ok=True)
+            color_grid.save(f"{save_dir}/color_delight.png")
+
+    multiviews = get_images_from_grid(color_grid, img_size=512)
+    view_weights = [1, 0.1, 0.02, 0.1, 1, 0.02]
+    view_weights += [0.01] * (len(multiviews) - len(view_weights))
+
+    # Use RealESRGAN_x4plus for x4 (512->2048) image super resolution.
+    if imagesr_model is None:
+        imagesr_model = ImageRealESRGAN(outscale=4)
+    multiviews = [imagesr_model(img) for img in multiviews]
+    multiviews = [img.convert("RGB") for img in multiviews]
+    mesh = trimesh.load(args.mesh_path)
+    if isinstance(mesh, trimesh.Scene):
+        mesh = mesh.dump(concatenate=True)
+
+    if not args.skip_fix_mesh:
+        mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
+        mesh_fixer = MeshFixer(mesh.vertices, mesh.faces, args.device)
+        mesh.vertices, mesh.faces = mesh_fixer(
+            filter_ratio=args.mesh_sipmlify_ratio,
+            max_hole_size=0.04,
+            resolution=1024,
+            num_views=1000,
+            norm_mesh_ratio=0.5,
+        )
+        if len(mesh.faces) > args.n_max_faces:
+            mesh.vertices, mesh.faces = mesh_fixer(
+                filter_ratio=0.8,
+                max_hole_size=0.04,
+                resolution=1024,
+                num_views=1000,
+                norm_mesh_ratio=0.5,
+            )
+        # Restore scale.
+        mesh.vertices = mesh.vertices / scale
+        mesh.vertices = mesh.vertices + center
+
+    # Baking texture to mesh.
+    texture_backer = TextureBacker(
+        camera_params=camera_params,
+        view_weights=view_weights,
+        render_wh=args.resolution_hw,
+        texture_wh=args.texture_wh,
+        smooth_texture=not args.no_smooth_texture,
+    )
+
+    textured_mesh = texture_backer(multiviews, mesh, args.output_path)
+
+    if args.save_glb_path is not None:
+        os.makedirs(os.path.dirname(args.save_glb_path), exist_ok=True)
+        textured_mesh.export(args.save_glb_path)
+
+    return textured_mesh
+
+
+if __name__ == "__main__":
+    entrypoint()

embodied_gen/data/convex_decomposer.py

@@ -0,0 +1,190 @@
+# 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.
+
+import logging
+import multiprocessing as mp
+import os
+
+import coacd
+import numpy as np
+import trimesh
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "decompose_convex_coacd",
+    "decompose_convex_mesh",
+    "decompose_convex_mp",
+]
+
+
+def decompose_convex_coacd(
+    filename: str,
+    outfile: str,
+    params: dict,
+    verbose: bool = False,
+    auto_scale: bool = True,
+    scale_factor: float = 1.0,
+) -> None:
+    coacd.set_log_level("info" if verbose else "warn")
+
+    mesh = trimesh.load(filename, force="mesh")
+    mesh = coacd.Mesh(mesh.vertices, mesh.faces)
+
+    result = coacd.run_coacd(mesh, **params)
+
+    meshes = []
+    for v, f in result:
+        meshes.append(trimesh.Trimesh(v, f))
+
+    # Compute collision_scale because convex decomposition usually makes the mesh larger.
+    if auto_scale:
+        all_mesh = sum([trimesh.Trimesh(*m) for m in result])
+        convex_mesh_shape = np.ptp(all_mesh.vertices, axis=0)
+        visual_mesh_shape = np.ptp(mesh.vertices, axis=0)
+        scale_factor *= visual_mesh_shape / convex_mesh_shape
+
+    combined = trimesh.Scene()
+    for mesh_part in meshes:
+        mesh_part.vertices *= scale_factor
+        combined.add_geometry(mesh_part)
+
+    combined.export(outfile)
+
+
+def decompose_convex_mesh(
+    filename: str,
+    outfile: str,
+    threshold: float = 0.05,
+    max_convex_hull: int = -1,
+    preprocess_mode: str = "auto",
+    preprocess_resolution: int = 30,
+    resolution: int = 2000,
+    mcts_nodes: int = 20,
+    mcts_iterations: int = 150,
+    mcts_max_depth: int = 3,
+    pca: bool = False,
+    merge: bool = True,
+    seed: int = 0,
+    auto_scale: bool = True,
+    scale_factor: float = 1.005,
+    verbose: bool = False,
+) -> str:
+    """Decompose a mesh into convex parts using the CoACD algorithm."""
+    coacd.set_log_level("info" if verbose else "warn")
+
+    if os.path.exists(outfile):
+        logger.warning(f"Output file {outfile} already exists, removing it.")
+        os.remove(outfile)
+
+    params = dict(
+        threshold=threshold,
+        max_convex_hull=max_convex_hull,
+        preprocess_mode=preprocess_mode,
+        preprocess_resolution=preprocess_resolution,
+        resolution=resolution,
+        mcts_nodes=mcts_nodes,
+        mcts_iterations=mcts_iterations,
+        mcts_max_depth=mcts_max_depth,
+        pca=pca,
+        merge=merge,
+        seed=seed,
+    )
+
+    try:
+        decompose_convex_coacd(
+            filename, outfile, params, verbose, auto_scale, scale_factor
+        )
+        if os.path.exists(outfile):
+            return outfile
+    except Exception as e:
+        if verbose:
+            print(f"Decompose convex first attempt failed: {e}.")
+
+    if preprocess_mode != "on":
+        try:
+            params["preprocess_mode"] = "on"
+            decompose_convex_coacd(
+                filename, outfile, params, verbose, auto_scale, scale_factor
+            )
+            if os.path.exists(outfile):
+                return outfile
+        except Exception as e:
+            if verbose:
+                print(
+                    f"Decompose convex second attempt with preprocess_mode='on' failed: {e}"
+                )
+
+    raise RuntimeError(f"Convex decomposition failed on {filename}")
+
+
+def decompose_convex_mp(
+    filename: str,
+    outfile: str,
+    threshold: float = 0.05,
+    max_convex_hull: int = -1,
+    preprocess_mode: str = "auto",
+    preprocess_resolution: int = 30,
+    resolution: int = 2000,
+    mcts_nodes: int = 20,
+    mcts_iterations: int = 150,
+    mcts_max_depth: int = 3,
+    pca: bool = False,
+    merge: bool = True,
+    seed: int = 0,
+    verbose: bool = False,
+    auto_scale: bool = True,
+) -> str:
+    """Decompose a mesh into convex parts using the CoACD algorithm in a separate process.
+
+    See https://simulately.wiki/docs/toolkits/ConvexDecomp for details.
+    """
+    params = dict(
+        threshold=threshold,
+        max_convex_hull=max_convex_hull,
+        preprocess_mode=preprocess_mode,
+        preprocess_resolution=preprocess_resolution,
+        resolution=resolution,
+        mcts_nodes=mcts_nodes,
+        mcts_iterations=mcts_iterations,
+        mcts_max_depth=mcts_max_depth,
+        pca=pca,
+        merge=merge,
+        seed=seed,
+    )
+
+    ctx = mp.get_context("spawn")
+    p = ctx.Process(
+        target=decompose_convex_coacd,
+        args=(filename, outfile, params, verbose, auto_scale),
+    )
+    p.start()
+    p.join()
+    if p.exitcode == 0 and os.path.exists(outfile):
+        return outfile
+
+    if preprocess_mode != "on":
+        params["preprocess_mode"] = "on"
+        p = ctx.Process(
+            target=decompose_convex_coacd,
+            args=(filename, outfile, params, verbose, auto_scale),
+        )
+        p.start()
+        p.join()
+        if p.exitcode == 0 and os.path.exists(outfile):
+            return outfile
+
+    raise RuntimeError(f"Convex decomposition failed on {filename}")

embodied_gen/data/datasets.py

@@ -0,0 +1,320 @@
+# 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.
+
+
+import json
+import logging
+import os
+import random
+from typing import Any, Callable, Dict, List, Literal, Tuple
+
+import numpy as np
+import torch
+import torch.utils.checkpoint
+from PIL import Image
+from torch import nn
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "Asset3dGenDataset",
+    "PanoGSplatDataset",
+]
+
+
+class Asset3dGenDataset(Dataset):
+    def __init__(
+        self,
+        index_file: str,
+        target_hw: Tuple[int, int],
+        transform: Callable = None,
+        control_transform: Callable = None,
+        max_train_samples: int = None,
+        sub_idxs: List[List[int]] = None,
+        seed: int = 79,
+    ) -> None:
+        if not os.path.exists(index_file):
+            raise FileNotFoundError(f"{index_file} index_file not found.")
+
+        self.index_file = index_file
+        self.target_hw = target_hw
+        self.transform = transform
+        self.control_transform = control_transform
+        self.max_train_samples = max_train_samples
+        self.meta_info = self.prepare_data_index(index_file)
+        self.data_list = sorted(self.meta_info.keys())
+        self.sub_idxs = sub_idxs  # sub_idxs [[0,1,2], [3,4,5], [...], ...]
+        self.image_num = 6  # hardcode temp.
+        random.seed(seed)
+        logger.info(f"Trainset: {len(self)} asset3d instances.")
+
+    def __len__(self) -> int:
+        return len(self.meta_info)
+
+    def prepare_data_index(self, index_file: str) -> Dict[str, Any]:
+        with open(index_file, "r") as fin:
+            meta_info = json.load(fin)
+
+        meta_info_filtered = dict()
+        for idx, uid in enumerate(meta_info):
+            if "status" not in meta_info[uid]:
+                continue
+            if meta_info[uid]["status"] != "success":
+                continue
+            if self.max_train_samples and idx >= self.max_train_samples:
+                break
+
+            meta_info_filtered[uid] = meta_info[uid]
+
+        logger.info(
+            f"Load {len(meta_info)} assets, keep {len(meta_info_filtered)} valids."  # noqa
+        )
+
+        return meta_info_filtered
+
+    def fetch_sample_images(
+        self,
+        uid: str,
+        attrs: List[str],
+        sub_index: int = None,
+        transform: Callable = None,
+    ) -> torch.Tensor:
+        sample = self.meta_info[uid]
+        images = []
+        for attr in attrs:
+            item = sample[attr]
+            if sub_index is not None:
+                item = item[sub_index]
+            mode = "L" if attr == "image_mask" else "RGB"
+            image = Image.open(item).convert(mode)
+            if transform is not None:
+                image = transform(image)
+                if len(image.shape) == 2:
+                    image = image[..., None]
+            images.append(image)
+
+        images = torch.cat(images, dim=0)
+
+        return images
+
+    def fetch_sample_grid_images(
+        self,
+        uid: str,
+        attrs: List[str],
+        sub_idxs: List[List[int]],
+        transform: Callable = None,
+    ) -> torch.Tensor:
+        assert transform is not None
+
+        grid_image = []
+        for row_idxs in sub_idxs:
+            row_image = []
+            for row_idx in row_idxs:
+                image = self.fetch_sample_images(
+                    uid, attrs, row_idx, transform
+                )
+                row_image.append(image)
+            row_image = torch.cat(row_image, dim=2)  # (c h w)
+            grid_image.append(row_image)
+
+        grid_image = torch.cat(grid_image, dim=1)
+
+        return grid_image
+
+    def compute_text_embeddings(
+        self, embed_path: str, original_size: Tuple[int, int]
+    ) -> Dict[str, nn.Module]:
+        data_dict = torch.load(embed_path)
+        prompt_embeds = data_dict["prompt_embeds"][0]
+        add_text_embeds = data_dict["pooled_prompt_embeds"][0]
+
+        # Need changed if random crop, set as crop_top_left [y1, x1], center crop as [0, 0].  # noqa
+        crops_coords_top_left = (0, 0)
+        add_time_ids = list(
+            original_size + crops_coords_top_left + self.target_hw
+        )
+        add_time_ids = torch.tensor([add_time_ids])
+        # add_time_ids = add_time_ids.repeat((len(add_text_embeds), 1))
+
+        unet_added_cond_kwargs = {
+            "text_embeds": add_text_embeds,
+            "time_ids": add_time_ids,
+        }
+
+        return {"prompt_embeds": prompt_embeds, **unet_added_cond_kwargs}
+
+    def visualize_item(
+        self,
+        control: torch.Tensor,
+        color: torch.Tensor,
+        save_dir: str = None,
+    ) -> List[Image.Image]:
+        to_pil = transforms.ToPILImage()
+
+        color = (color + 1) / 2
+        color_pil = to_pil(color)
+        normal_pil = to_pil(control[0:3])
+        position_pil = to_pil(control[3:6])
+        mask_pil = to_pil(control[6:])
+
+        if save_dir is not None:
+            os.makedirs(save_dir, exist_ok=True)
+            color_pil.save(f"{save_dir}/rgb.jpg")
+            normal_pil.save(f"{save_dir}/normal.jpg")
+            position_pil.save(f"{save_dir}/position.jpg")
+            mask_pil.save(f"{save_dir}/mask.jpg")
+            logger.info(f"Visualization in {save_dir}")
+
+        return normal_pil, position_pil, mask_pil, color_pil
+
+    def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:
+        uid = self.data_list[index]
+
+        sub_idxs = self.sub_idxs
+        if sub_idxs is None:
+            sub_idxs = [[random.randint(0, self.image_num - 1)]]
+
+        input_image = self.fetch_sample_grid_images(
+            uid,
+            attrs=["image_view_normal", "image_position", "image_mask"],
+            sub_idxs=sub_idxs,
+            transform=self.control_transform,
+        )
+        assert input_image.shape[1:] == self.target_hw
+
+        output_image = self.fetch_sample_grid_images(
+            uid,
+            attrs=["image_color"],
+            sub_idxs=sub_idxs,
+            transform=self.transform,
+        )
+
+        sample = self.meta_info[uid]
+        text_feats = self.compute_text_embeddings(
+            sample["text_feat"], tuple(sample["image_hw"])
+        )
+
+        data = dict(
+            pixel_values=output_image,
+            conditioning_pixel_values=input_image,
+            prompt_embeds=text_feats["prompt_embeds"],
+            text_embeds=text_feats["text_embeds"],
+            time_ids=text_feats["time_ids"],
+        )
+
+        return data
+
+
+class PanoGSplatDataset(Dataset):
+    """A PyTorch Dataset for loading panorama-based 3D Gaussian Splatting data.
+
+    This dataset is designed to be compatible with train and eval pipelines
+    that use COLMAP-style camera conventions.
+
+    Args:
+        data_dir (str): Root directory where the dataset file is located.
+        split (str): Dataset split to use, either "train" or "eval".
+        data_name (str, optional): Name of the dataset file (default: "gs_data.pt").
+        max_sample_num (int, optional): Maximum number of samples to load. If None,
+            all available samples in the split will be used.
+    """
+
+    def __init__(
+        self,
+        data_dir: str,
+        split: str = Literal["train", "eval"],
+        data_name: str = "gs_data.pt",
+        max_sample_num: int = None,
+    ) -> None:
+        self.data_path = os.path.join(data_dir, data_name)
+        self.split = split
+        self.max_sample_num = max_sample_num
+        if not os.path.exists(self.data_path):
+            raise FileNotFoundError(
+                f"Dataset file {self.data_path} not found. Please provide the correct path."
+            )
+        self.data = torch.load(self.data_path, weights_only=False)
+        self.frames = self.data[split]
+        if max_sample_num is not None:
+            self.frames = self.frames[:max_sample_num]
+        self.points = self.data.get("points", None)
+        self.points_rgb = self.data.get("points_rgb", None)
+
+    def __len__(self) -> int:
+        return len(self.frames)
+
+    def cvt_blender_to_colmap_coord(self, c2w: np.ndarray) -> np.ndarray:
+        # change from OpenGL/Blender camera axes (Y up, Z back) to COLMAP (Y down, Z forward)
+        tranformed_c2w = np.copy(c2w)
+        tranformed_c2w[:3, 1:3] *= -1
+
+        return tranformed_c2w
+
+    def __getitem__(self, index: int) -> dict[str, any]:
+        data = self.frames[index]
+        c2w = self.cvt_blender_to_colmap_coord(data["camtoworld"])
+        item = dict(
+            camtoworld=c2w,
+            K=data["K"],
+            image_h=data["image_h"],
+            image_w=data["image_w"],
+        )
+        if "image" in data:
+            item["image"] = data["image"]
+        if "image_id" in data:
+            item["image_id"] = data["image_id"]
+
+        return item
+
+
+if __name__ == "__main__":
+    index_file = "datasets/objaverse/v1.0/statistics_1.0_gobjaverse_filter/view6s_v4/meta_ac2e0ddea8909db26d102c8465b5bcb2.json"  # noqa
+    target_hw = (512, 512)
+    transform_list = [
+        transforms.Resize(
+            target_hw, interpolation=transforms.InterpolationMode.BILINEAR
+        ),
+        transforms.CenterCrop(target_hw),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5]),
+    ]
+    image_transform = transforms.Compose(transform_list)
+    control_transform = transforms.Compose(transform_list[:-1])
+
+    sub_idxs = [[0, 1, 2], [3, 4, 5]]  # None
+    if sub_idxs is not None:
+        target_hw = (
+            target_hw[0] * len(sub_idxs),
+            target_hw[1] * len(sub_idxs[0]),
+        )
+
+    dataset = Asset3dGenDataset(
+        index_file,
+        target_hw,
+        image_transform,
+        control_transform,
+        sub_idxs=sub_idxs,
+    )
+    data = dataset[0]
+    dataset.visualize_item(
+        data["conditioning_pixel_values"], data["pixel_values"], save_dir="./"
+    )

embodied_gen/data/differentiable_render.py

@@ -0,0 +1,589 @@
+# 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.
+
+
+import argparse
+import json
+import logging
+import math
+import os
+from collections import defaultdict
+from typing import List, Union
+
+import cv2
+import imageio
+import numpy as np
+import nvdiffrast.torch as dr
+import PIL.Image as Image
+import torch
+from tqdm import tqdm
+from embodied_gen.data.utils import (
+    CameraSetting,
+    DiffrastRender,
+    as_list,
+    calc_vertex_normals,
+    import_kaolin_mesh,
+    init_kal_camera,
+    normalize_vertices_array,
+    render_pbr,
+    save_images,
+)
+from embodied_gen.utils.enum import RenderItems
+
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+os.environ["TORCH_EXTENSIONS_DIR"] = os.path.expanduser(
+    "~/.cache/torch_extensions"
+)
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "ImageRender",
+    "create_mp4_from_images",
+    "create_gif_from_images",
+]
+
+
+def create_mp4_from_images(
+    images: list[np.ndarray],
+    output_path: str,
+    fps: int = 10,
+    prompt: str = None,
+):
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    font_scale = 0.5
+    font_thickness = 1
+    color = (255, 255, 255)
+    position = (20, 25)
+
+    with imageio.get_writer(output_path, fps=fps) as writer:
+        for image in images:
+            image = image.clip(min=0, max=1)
+            image = (255.0 * image).astype(np.uint8)
+            image = image[..., :3]
+            if prompt is not None:
+                cv2.putText(
+                    image,
+                    prompt,
+                    position,
+                    font,
+                    font_scale,
+                    color,
+                    font_thickness,
+                )
+
+            writer.append_data(image)
+
+    logger.info(f"MP4 video saved to {output_path}")
+
+
+def create_gif_from_images(
+    images: list[np.ndarray], output_path: str, fps: int = 10
+) -> None:
+    pil_images = []
+    for image in images:
+        image = image.clip(min=0, max=1)
+        image = (255.0 * image).astype(np.uint8)
+        image = Image.fromarray(image, mode="RGBA")
+        pil_images.append(image.convert("RGB"))
+
+    duration = 1000 // fps
+    pil_images[0].save(
+        output_path,
+        save_all=True,
+        append_images=pil_images[1:],
+        duration=duration,
+        loop=0,
+    )
+
+    logger.info(f"GIF saved to {output_path}")
+
+
+class ImageRender(object):
+    """A differentiable mesh renderer supporting multi-view rendering.
+
+    This class wraps a differentiable rasterization using `nvdiffrast` to
+    render mesh geometry to various maps (normal, depth, alpha, albedo, etc.).
+
+    Args:
+        render_items (list[RenderItems]): A list of rendering targets to
+            generate (e.g., IMAGE, DEPTH, NORMAL, etc.).
+        camera_params (CameraSetting): The camera parameters for rendering,
+            including intrinsic and extrinsic matrices.
+        recompute_vtx_normal (bool, optional): If True, recomputes
+            vertex normals from the mesh geometry. Defaults to True.
+        with_mtl (bool, optional): Whether to load `.mtl` material files
+            for meshes. Defaults to False.
+        gen_color_gif (bool, optional): Generate a GIF of rendered
+            color images. Defaults to False.
+        gen_color_mp4 (bool, optional): Generate an MP4 video of rendered
+            color images. Defaults to False.
+        gen_viewnormal_mp4 (bool, optional): Generate an MP4 video of
+            view-space normals. Defaults to False.
+        gen_glonormal_mp4 (bool, optional): Generate an MP4 video of
+            global-space normals. Defaults to False.
+        no_index_file (bool, optional): If True, skip saving the `index.json`
+            summary file. Defaults to False.
+        light_factor (float, optional): A scalar multiplier for
+            PBR light intensity. Defaults to 1.0.
+    """
+
+    def __init__(
+        self,
+        render_items: list[RenderItems],
+        camera_params: CameraSetting,
+        recompute_vtx_normal: bool = True,
+        with_mtl: bool = False,
+        gen_color_gif: bool = False,
+        gen_color_mp4: bool = False,
+        gen_viewnormal_mp4: bool = False,
+        gen_glonormal_mp4: bool = False,
+        no_index_file: bool = False,
+        light_factor: float = 1.0,
+    ) -> None:
+        camera = init_kal_camera(camera_params)
+        self.camera = camera
+
+        # Setup MVP matrix and renderer.
+        mv = camera.view_matrix()  # (n 4 4) world2cam
+        p = camera.intrinsics.projection_matrix()
+        # NOTE: add a negative sign at P[0, 2] as the y axis is flipped in `nvdiffrast` output.  # noqa
+        p[:, 1, 1] = -p[:, 1, 1]
+        # mvp = torch.bmm(p, mv) # camera.view_projection_matrix()
+        self.mv = mv
+        self.p = p
+
+        renderer = DiffrastRender(
+            p_matrix=p,
+            mv_matrix=mv,
+            resolution_hw=camera_params.resolution_hw,
+            context=dr.RasterizeCudaContext(),
+            mask_thresh=0.5,
+            grad_db=False,
+            device=camera_params.device,
+            antialias_mask=True,
+        )
+        self.renderer = renderer
+        self.recompute_vtx_normal = recompute_vtx_normal
+        self.render_items = render_items
+        self.device = camera_params.device
+        self.with_mtl = with_mtl
+        self.gen_color_gif = gen_color_gif
+        self.gen_color_mp4 = gen_color_mp4
+        self.gen_viewnormal_mp4 = gen_viewnormal_mp4
+        self.gen_glonormal_mp4 = gen_glonormal_mp4
+        self.light_factor = light_factor
+        self.no_index_file = no_index_file
+
+    def render_mesh(
+        self,
+        mesh_path: Union[str, List[str]],
+        output_root: str,
+        uuid: Union[str, List[str]] = None,
+        prompts: List[str] = None,
+    ) -> None:
+        mesh_path = as_list(mesh_path)
+        if uuid is None:
+            uuid = [os.path.basename(p).split(".")[0] for p in mesh_path]
+        uuid = as_list(uuid)
+        assert len(mesh_path) == len(uuid)
+        os.makedirs(output_root, exist_ok=True)
+
+        meta_info = dict()
+        for idx, (path, uid) in tqdm(
+            enumerate(zip(mesh_path, uuid)), total=len(mesh_path)
+        ):
+            output_dir = os.path.join(output_root, uid)
+            os.makedirs(output_dir, exist_ok=True)
+            prompt = prompts[idx] if prompts else None
+            data_dict = self(path, output_dir, prompt)
+            meta_info[uid] = data_dict
+
+        if self.no_index_file:
+            return
+
+        index_file = os.path.join(output_root, "index.json")
+        with open(index_file, "w") as fout:
+            json.dump(meta_info, fout)
+
+        logger.info(f"Rendering meta info logged in {index_file}")
+
+    def __call__(
+        self, mesh_path: str, output_dir: str, prompt: str = None
+    ) -> dict[str, str]:
+        """Render a single mesh and return paths to the rendered outputs.
+
+        Processes the input mesh, renders multiple modalities (e.g., normals,
+        depth, albedo), and optionally saves video or image sequences.
+
+        Args:
+            mesh_path (str): Path to the mesh file (.obj/.glb).
+            output_dir (str): Directory to save rendered outputs.
+            prompt (str, optional): Optional caption prompt for MP4 metadata.
+
+        Returns:
+            dict[str, str]: A mapping render types to the saved image paths.
+        """
+        try:
+            mesh = import_kaolin_mesh(mesh_path, self.with_mtl)
+        except Exception as e:
+            logger.error(f"[ERROR MESH LOAD]: {e}, skip {mesh_path}")
+            return
+
+        mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
+        if self.recompute_vtx_normal:
+            mesh.vertex_normals = calc_vertex_normals(
+                mesh.vertices, mesh.faces
+            )
+
+        mesh = mesh.to(self.device)
+        vertices, faces, vertex_normals = (
+            mesh.vertices,
+            mesh.faces,
+            mesh.vertex_normals,
+        )
+
+        # Perform rendering.
+        data_dict = defaultdict(list)
+        if RenderItems.ALPHA.value in self.render_items:
+            masks, _ = self.renderer.render_rast_alpha(vertices, faces)
+            render_paths = save_images(
+                masks, f"{output_dir}/{RenderItems.ALPHA}"
+            )
+            data_dict[RenderItems.ALPHA.value] = render_paths
+
+        if RenderItems.GLOBAL_NORMAL.value in self.render_items:
+            rendered_normals, masks = self.renderer.render_global_normal(
+                vertices, faces, vertex_normals
+            )
+            if self.gen_glonormal_mp4:
+                if isinstance(rendered_normals, torch.Tensor):
+                    rendered_normals = rendered_normals.detach().cpu().numpy()
+                create_mp4_from_images(
+                    rendered_normals,
+                    output_path=f"{output_dir}/normal.mp4",
+                    fps=15,
+                    prompt=prompt,
+                )
+            else:
+                render_paths = save_images(
+                    rendered_normals,
+                    f"{output_dir}/{RenderItems.GLOBAL_NORMAL}",
+                    cvt_color=cv2.COLOR_BGR2RGB,
+                )
+                data_dict[RenderItems.GLOBAL_NORMAL.value] = render_paths
+
+            if RenderItems.VIEW_NORMAL.value in self.render_items:
+                assert (
+                    RenderItems.GLOBAL_NORMAL in self.render_items
+                ), f"Must render global normal firstly, got render_items: {self.render_items}."  # noqa
+                rendered_view_normals = self.renderer.transform_normal(
+                    rendered_normals, self.mv, masks, to_view=True
+                )
+
+                if self.gen_viewnormal_mp4:
+                    create_mp4_from_images(
+                        rendered_view_normals,
+                        output_path=f"{output_dir}/view_normal.mp4",
+                        fps=15,
+                        prompt=prompt,
+                    )
+                else:
+                    render_paths = save_images(
+                        rendered_view_normals,
+                        f"{output_dir}/{RenderItems.VIEW_NORMAL}",
+                        cvt_color=cv2.COLOR_BGR2RGB,
+                    )
+                    data_dict[RenderItems.VIEW_NORMAL.value] = render_paths
+
+        if RenderItems.POSITION_MAP.value in self.render_items:
+            rendered_position, masks = self.renderer.render_position(
+                vertices, faces
+            )
+            norm_position = self.renderer.normalize_map_by_mask(
+                rendered_position, masks
+            )
+            render_paths = save_images(
+                norm_position,
+                f"{output_dir}/{RenderItems.POSITION_MAP}",
+                cvt_color=cv2.COLOR_BGR2RGB,
+            )
+            data_dict[RenderItems.POSITION_MAP.value] = render_paths
+
+        if RenderItems.DEPTH.value in self.render_items:
+            rendered_depth, masks = self.renderer.render_depth(vertices, faces)
+            norm_depth = self.renderer.normalize_map_by_mask(
+                rendered_depth, masks
+            )
+            render_paths = save_images(
+                norm_depth,
+                f"{output_dir}/{RenderItems.DEPTH}",
+            )
+            data_dict[RenderItems.DEPTH.value] = render_paths
+
+            render_paths = save_images(
+                rendered_depth,
+                f"{output_dir}/{RenderItems.DEPTH}_exr",
+                to_uint8=False,
+                format=".exr",
+            )
+            data_dict[f"{RenderItems.DEPTH.value}_exr"] = render_paths
+
+        if RenderItems.IMAGE.value in self.render_items:
+            images = []
+            albedos = []
+            diffuses = []
+            masks, _ = self.renderer.render_rast_alpha(vertices, faces)
+            try:
+                for idx, cam in enumerate(self.camera):
+                    image, albedo, diffuse, _ = render_pbr(
+                        mesh, cam, light_factor=self.light_factor
+                    )
+                    image = torch.cat([image[0], masks[idx]], axis=-1)
+                    images.append(image.detach().cpu().numpy())
+
+                    if RenderItems.ALBEDO.value in self.render_items:
+                        albedo = torch.cat([albedo[0], masks[idx]], axis=-1)
+                        albedos.append(albedo.detach().cpu().numpy())
+
+                    if RenderItems.DIFFUSE.value in self.render_items:
+                        diffuse = torch.cat([diffuse[0], masks[idx]], axis=-1)
+                        diffuses.append(diffuse.detach().cpu().numpy())
+
+            except Exception as e:
+                logger.error(f"[ERROR pbr render]: {e}, skip {mesh_path}")
+                return
+
+            if self.gen_color_gif:
+                create_gif_from_images(
+                    images,
+                    output_path=f"{output_dir}/color.gif",
+                    fps=15,
+                )
+
+            if self.gen_color_mp4:
+                create_mp4_from_images(
+                    images,
+                    output_path=f"{output_dir}/color.mp4",
+                    fps=15,
+                    prompt=prompt,
+                )
+
+            if self.gen_color_mp4 or self.gen_color_gif:
+                return data_dict
+
+            render_paths = save_images(
+                images,
+                f"{output_dir}/{RenderItems.IMAGE}",
+                cvt_color=cv2.COLOR_BGRA2RGBA,
+            )
+            data_dict[RenderItems.IMAGE.value] = render_paths
+
+            render_paths = save_images(
+                albedos,
+                f"{output_dir}/{RenderItems.ALBEDO}",
+                cvt_color=cv2.COLOR_BGRA2RGBA,
+            )
+            data_dict[RenderItems.ALBEDO.value] = render_paths
+
+            render_paths = save_images(
+                diffuses,
+                f"{output_dir}/{RenderItems.DIFFUSE}",
+                cvt_color=cv2.COLOR_BGRA2RGBA,
+            )
+            data_dict[RenderItems.DIFFUSE.value] = render_paths
+
+        data_dict["status"] = "success"
+
+        logger.info(f"Finish rendering in {output_dir}")
+
+        return data_dict
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Render settings")
+
+    parser.add_argument(
+        "--mesh_path",
+        type=str,
+        nargs="+",
+        help="Paths to the mesh files for rendering.",
+    )
+    parser.add_argument(
+        "--output_root",
+        type=str,
+        help="Root directory for output",
+    )
+    parser.add_argument(
+        "--uuid",
+        type=str,
+        nargs="+",
+        default=None,
+        help="uuid for rendering saving.",
+    )
+    parser.add_argument(
+        "--num_images", type=int, default=6, help="Number of images to render."
+    )
+    parser.add_argument(
+        "--elevation",
+        type=float,
+        nargs="+",
+        default=[20.0, -10.0],
+        help="Elevation angles for the camera (default: [20.0, -10.0])",
+    )
+    parser.add_argument(
+        "--distance",
+        type=float,
+        default=5,
+        help="Camera distance (default: 5)",
+    )
+    parser.add_argument(
+        "--resolution_hw",
+        type=int,
+        nargs=2,
+        default=(512, 512),
+        help="Resolution of the output images (default: (512, 512))",
+    )
+    parser.add_argument(
+        "--fov",
+        type=float,
+        default=30,
+        help="Field of view in degrees (default: 30)",
+    )
+    parser.add_argument(
+        "--pbr_light_factor",
+        type=float,
+        default=1.0,
+        help="Light factor for mesh PBR rendering (default: 1.)",
+    )
+    parser.add_argument(
+        "--with_mtl",
+        action="store_true",
+        help="Whether to render with mesh material.",
+    )
+    parser.add_argument(
+        "--gen_color_gif",
+        action="store_true",
+        help="Whether to generate color .gif rendering file.",
+    )
+    parser.add_argument(
+        "--no_index_file",
+        action="store_true",
+        help="Whether skip the index file saving.",
+    )
+    parser.add_argument(
+        "--gen_color_mp4",
+        action="store_true",
+        help="Whether to generate color .mp4 rendering file.",
+    )
+    parser.add_argument(
+        "--gen_viewnormal_mp4",
+        action="store_true",
+        help="Whether to generate view normal .mp4 rendering file.",
+    )
+    parser.add_argument(
+        "--gen_glonormal_mp4",
+        action="store_true",
+        help="Whether to generate global normal .mp4 rendering file.",
+    )
+    parser.add_argument(
+        "--video_prompts",
+        type=str,
+        nargs="+",
+        default=None,
+        help="Text prompts for the rendering.",
+    )
+
+    args, unknown = parser.parse_known_args()
+
+    if args.uuid is None and args.mesh_path is not None:
+        args.uuid = []
+        for path in args.mesh_path:
+            uuid = os.path.basename(path).split(".")[0]
+            args.uuid.append(uuid)
+
+    return args
+
+
+def entrypoint(**kwargs) -> None:
+    args = parse_args()
+    for k, v in kwargs.items():
+        if hasattr(args, k) and v is not None:
+            setattr(args, k, v)
+
+    camera_settings = CameraSetting(
+        num_images=args.num_images,
+        elevation=args.elevation,
+        distance=args.distance,
+        resolution_hw=args.resolution_hw,
+        fov=math.radians(args.fov),
+        device="cuda",
+    )
+
+    render_items = [
+        RenderItems.ALPHA.value,
+        RenderItems.GLOBAL_NORMAL.value,
+        RenderItems.VIEW_NORMAL.value,
+        RenderItems.POSITION_MAP.value,
+        RenderItems.IMAGE.value,
+        RenderItems.DEPTH.value,
+        # RenderItems.ALBEDO.value,
+        # RenderItems.DIFFUSE.value,
+    ]
+
+    gen_video = (
+        args.gen_color_gif
+        or args.gen_color_mp4
+        or args.gen_viewnormal_mp4
+        or args.gen_glonormal_mp4
+    )
+    if gen_video:
+        render_items = []
+        if args.gen_color_gif or args.gen_color_mp4:
+            render_items.append(RenderItems.IMAGE.value)
+        if args.gen_glonormal_mp4:
+            render_items.append(RenderItems.GLOBAL_NORMAL.value)
+        if args.gen_viewnormal_mp4:
+            render_items.append(RenderItems.VIEW_NORMAL.value)
+            if RenderItems.GLOBAL_NORMAL.value not in render_items:
+                render_items.append(RenderItems.GLOBAL_NORMAL.value)
+
+    image_render = ImageRender(
+        render_items=render_items,
+        camera_params=camera_settings,
+        with_mtl=args.with_mtl,
+        gen_color_gif=args.gen_color_gif,
+        gen_color_mp4=args.gen_color_mp4,
+        gen_viewnormal_mp4=args.gen_viewnormal_mp4,
+        gen_glonormal_mp4=args.gen_glonormal_mp4,
+        light_factor=args.pbr_light_factor,
+        no_index_file=gen_video or args.no_index_file,
+    )
+    image_render.render_mesh(
+        mesh_path=args.mesh_path,
+        output_root=args.output_root,
+        uuid=args.uuid,
+        prompts=args.video_prompts,
+    )
+
+    return
+
+
+if __name__ == "__main__":
+    entrypoint()

embodied_gen/data/mesh_operator.py

@@ -0,0 +1,461 @@
+# 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.
+
+
+import logging
+import multiprocessing as mp
+import os
+from typing import Tuple, Union
+
+import coacd
+import igraph
+import numpy as np
+import pyvista as pv
+import spaces
+import torch
+import trimesh
+import utils3d
+from pymeshfix import _meshfix
+from tqdm import tqdm
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "MeshFixer",
+]
+
+
+def _radical_inverse(base, n):
+    val = 0
+    inv_base = 1.0 / base
+    inv_base_n = inv_base
+    while n > 0:
+        digit = n % base
+        val += digit * inv_base_n
+        n //= base
+        inv_base_n *= inv_base
+    return val
+
+
+def _halton_sequence(dim, n):
+    PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]
+    return [_radical_inverse(PRIMES[dim], n) for dim in range(dim)]
+
+
+def _hammersley_sequence(dim, n, num_samples):
+    return [n / num_samples] + _halton_sequence(dim - 1, n)
+
+
+def _sphere_hammersley_seq(n, num_samples, offset=(0, 0), remap=False):
+    """Generate a point on a unit sphere using the Hammersley sequence.
+
+    Args:
+        n (int): The index of the sample.
+        num_samples (int): The total number of samples.
+        offset (tuple, optional): Offset for the u and v coordinates.
+        remap (bool, optional): Whether to remap the u coordinate.
+
+    Returns:
+        list: A list containing the spherical coordinates [phi, theta].
+    """
+    u, v = _hammersley_sequence(2, n, num_samples)
+    u += offset[0] / num_samples
+    v += offset[1]
+
+    if remap:
+        u = 2 * u if u < 0.25 else 2 / 3 * u + 1 / 3
+
+    theta = np.arccos(1 - 2 * u) - np.pi / 2
+    phi = v * 2 * np.pi
+    return [phi, theta]
+
+
+class MeshFixer(object):
+    """MeshFixer simplifies and repairs 3D triangle meshes by TSDF.
+
+    Attributes:
+        vertices (torch.Tensor): A tensor of shape (V, 3) representing vertex positions.
+        faces (torch.Tensor): A tensor of shape (F, 3) representing face indices.
+        device (str): Device to run computations on, typically "cuda" or "cpu".
+
+    Main logic reference: https://github.com/microsoft/TRELLIS/blob/main/trellis/utils/postprocessing_utils.py#L22
+    """
+
+    def __init__(
+        self,
+        vertices: Union[torch.Tensor, np.ndarray],
+        faces: Union[torch.Tensor, np.ndarray],
+        device: str = "cuda",
+    ) -> None:
+        self.device = device
+        if isinstance(vertices, np.ndarray):
+            vertices = torch.tensor(vertices)
+        self.vertices = vertices
+
+        if isinstance(faces, np.ndarray):
+            faces = torch.tensor(faces)
+        self.faces = faces
+
+    @staticmethod
+    def log_mesh_changes(method):
+        def wrapper(self, *args, **kwargs):
+            logger.info(
+                f"Before {method.__name__}: {self.vertices.shape[0]} vertices, {self.faces.shape[0]} faces"  # noqa
+            )
+            result = method(self, *args, **kwargs)
+            logger.info(
+                f"After {method.__name__}: {self.vertices.shape[0]} vertices, {self.faces.shape[0]} faces"  # noqa
+            )
+            return result
+
+        return wrapper
+
+    @log_mesh_changes
+    def fill_holes(
+        self,
+        max_hole_size: float,
+        max_hole_nbe: int,
+        resolution: int,
+        num_views: int,
+        norm_mesh_ratio: float = 1.0,
+    ) -> None:
+        self.vertices = self.vertices * norm_mesh_ratio
+        vertices, self.faces = self._fill_holes(
+            self.vertices,
+            self.faces,
+            max_hole_size,
+            max_hole_nbe,
+            resolution,
+            num_views,
+        )
+        self.vertices = vertices / norm_mesh_ratio
+
+    @staticmethod
+    @torch.no_grad()
+    def _fill_holes(
+        vertices: torch.Tensor,
+        faces: torch.Tensor,
+        max_hole_size: float,
+        max_hole_nbe: int,
+        resolution: int,
+        num_views: int,
+    ) -> Union[torch.Tensor, torch.Tensor]:
+        yaws, pitchs = [], []
+        for i in range(num_views):
+            y, p = _sphere_hammersley_seq(i, num_views)
+            yaws.append(y)
+            pitchs.append(p)
+
+        yaws, pitchs = torch.tensor(yaws).to(vertices), torch.tensor(
+            pitchs
+        ).to(vertices)
+        radius, fov = 2.0, torch.deg2rad(torch.tensor(40)).to(vertices)
+        projection = utils3d.torch.perspective_from_fov_xy(fov, fov, 1, 3)
+
+        views = []
+        for yaw, pitch in zip(yaws, pitchs):
+            orig = (
+                torch.tensor(
+                    [
+                        torch.sin(yaw) * torch.cos(pitch),
+                        torch.cos(yaw) * torch.cos(pitch),
+                        torch.sin(pitch),
+                    ]
+                ).to(vertices)
+                * radius
+            )
+            view = utils3d.torch.view_look_at(
+                orig,
+                torch.tensor([0, 0, 0]).to(vertices),
+                torch.tensor([0, 0, 1]).to(vertices),
+            )
+            views.append(view)
+        views = torch.stack(views, dim=0)
+
+        # Rasterize the mesh
+        visibility = torch.zeros(
+            faces.shape[0], dtype=torch.int32, device=faces.device
+        )
+        rastctx = utils3d.torch.RastContext(backend="cuda")
+
+        for i in tqdm(
+            range(views.shape[0]), total=views.shape[0], desc="Rasterizing"
+        ):
+            view = views[i]
+            buffers = utils3d.torch.rasterize_triangle_faces(
+                rastctx,
+                vertices[None],
+                faces,
+                resolution,
+                resolution,
+                view=view,
+                projection=projection,
+            )
+            face_id = buffers["face_id"][0][buffers["mask"][0] > 0.95] - 1
+            face_id = torch.unique(face_id).long()
+            visibility[face_id] += 1
+
+        # Normalize visibility by the number of views
+        visibility = visibility.float() / num_views
+
+        # Mincut: Identify outer and inner faces
+        edges, face2edge, edge_degrees = utils3d.torch.compute_edges(faces)
+        boundary_edge_indices = torch.nonzero(edge_degrees == 1).reshape(-1)
+        connected_components = utils3d.torch.compute_connected_components(
+            faces, edges, face2edge
+        )
+
+        outer_face_indices = torch.zeros(
+            faces.shape[0], dtype=torch.bool, device=faces.device
+        )
+        for i in range(len(connected_components)):
+            outer_face_indices[connected_components[i]] = visibility[
+                connected_components[i]
+            ] > min(
+                max(
+                    visibility[connected_components[i]].quantile(0.75).item(),
+                    0.25,
+                ),
+                0.5,
+            )
+
+        outer_face_indices = outer_face_indices.nonzero().reshape(-1)
+        inner_face_indices = torch.nonzero(visibility == 0).reshape(-1)
+
+        if inner_face_indices.shape[0] == 0:
+            return vertices, faces
+
+        # Construct dual graph (faces as nodes, edges as edges)
+        dual_edges, dual_edge2edge = utils3d.torch.compute_dual_graph(
+            face2edge
+        )
+        dual_edge2edge = edges[dual_edge2edge]
+        dual_edges_weights = torch.norm(
+            vertices[dual_edge2edge[:, 0]] - vertices[dual_edge2edge[:, 1]],
+            dim=1,
+        )
+
+        # Mincut: Construct main graph and solve the mincut problem
+        g = igraph.Graph()
+        g.add_vertices(faces.shape[0])
+        g.add_edges(dual_edges.cpu().numpy())
+        g.es["weight"] = dual_edges_weights.cpu().numpy()
+
+        g.add_vertex("s")  # source
+        g.add_vertex("t")  # target
+
+        g.add_edges(
+            [(f, "s") for f in inner_face_indices],
+            attributes={
+                "weight": torch.ones(
+                    inner_face_indices.shape[0], dtype=torch.float32
+                )
+                .cpu()
+                .numpy()
+            },
+        )
+        g.add_edges(
+            [(f, "t") for f in outer_face_indices],
+            attributes={
+                "weight": torch.ones(
+                    outer_face_indices.shape[0], dtype=torch.float32
+                )
+                .cpu()
+                .numpy()
+            },
+        )
+
+        cut = g.mincut("s", "t", (np.array(g.es["weight"]) * 1000).tolist())
+        remove_face_indices = torch.tensor(
+            [v for v in cut.partition[0] if v < faces.shape[0]],
+            dtype=torch.long,
+            device=faces.device,
+        )
+
+        # Check if the cut is valid with each connected component
+        to_remove_cc = utils3d.torch.compute_connected_components(
+            faces[remove_face_indices]
+        )
+        valid_remove_cc = []
+        cutting_edges = []
+        for cc in to_remove_cc:
+            # Check visibility median for connected component
+            visibility_median = visibility[remove_face_indices[cc]].median()
+            if visibility_median > 0.25:
+                continue
+
+            # Check if the cutting loop is small enough
+            cc_edge_indices, cc_edges_degree = torch.unique(
+                face2edge[remove_face_indices[cc]], return_counts=True
+            )
+            cc_boundary_edge_indices = cc_edge_indices[cc_edges_degree == 1]
+            cc_new_boundary_edge_indices = cc_boundary_edge_indices[
+                ~torch.isin(cc_boundary_edge_indices, boundary_edge_indices)
+            ]
+            if len(cc_new_boundary_edge_indices) > 0:
+                cc_new_boundary_edge_cc = (
+                    utils3d.torch.compute_edge_connected_components(
+                        edges[cc_new_boundary_edge_indices]
+                    )
+                )
+                cc_new_boundary_edges_cc_center = [
+                    vertices[edges[cc_new_boundary_edge_indices[edge_cc]]]
+                    .mean(dim=1)
+                    .mean(dim=0)
+                    for edge_cc in cc_new_boundary_edge_cc
+                ]
+                cc_new_boundary_edges_cc_area = []
+                for i, edge_cc in enumerate(cc_new_boundary_edge_cc):
+                    _e1 = (
+                        vertices[
+                            edges[cc_new_boundary_edge_indices[edge_cc]][:, 0]
+                        ]
+                        - cc_new_boundary_edges_cc_center[i]
+                    )
+                    _e2 = (
+                        vertices[
+                            edges[cc_new_boundary_edge_indices[edge_cc]][:, 1]
+                        ]
+                        - cc_new_boundary_edges_cc_center[i]
+                    )
+                    cc_new_boundary_edges_cc_area.append(
+                        torch.norm(torch.cross(_e1, _e2, dim=-1), dim=1).sum()
+                        * 0.5
+                    )
+                cutting_edges.append(cc_new_boundary_edge_indices)
+                if any(
+                    [
+                        _l > max_hole_size
+                        for _l in cc_new_boundary_edges_cc_area
+                    ]
+                ):
+                    continue
+
+            valid_remove_cc.append(cc)
+
+        if len(valid_remove_cc) > 0:
+            remove_face_indices = remove_face_indices[
+                torch.cat(valid_remove_cc)
+            ]
+            mask = torch.ones(
+                faces.shape[0], dtype=torch.bool, device=faces.device
+            )
+            mask[remove_face_indices] = 0
+            faces = faces[mask]
+            faces, vertices = utils3d.torch.remove_unreferenced_vertices(
+                faces, vertices
+            )
+
+            tqdm.write(f"Removed {(~mask).sum()} faces by mincut")
+        else:
+            tqdm.write(f"Removed 0 faces by mincut")
+
+        # Fill small boundaries (holes)
+        mesh = _meshfix.PyTMesh()
+        mesh.load_array(vertices.cpu().numpy(), faces.cpu().numpy())
+        mesh.fill_small_boundaries(nbe=max_hole_nbe, refine=True)
+
+        _vertices, _faces = mesh.return_arrays()
+        vertices = torch.tensor(_vertices).to(vertices)
+        faces = torch.tensor(_faces).to(faces)
+
+        return vertices, faces
+
+    @property
+    def vertices_np(self) -> np.ndarray:
+        return self.vertices.cpu().numpy()
+
+    @property
+    def faces_np(self) -> np.ndarray:
+        return self.faces.cpu().numpy()
+
+    @log_mesh_changes
+    def simplify(self, ratio: float) -> None:
+        """Simplify the mesh using quadric edge collapse decimation.
+
+        Args:
+            ratio (float): Ratio of faces to filter out.
+        """
+        if ratio <= 0 or ratio >= 1:
+            raise ValueError("Simplify ratio must be between 0 and 1.")
+
+        # Convert to PyVista format for simplification
+        mesh = pv.PolyData(
+            self.vertices_np,
+            np.hstack([np.full((self.faces.shape[0], 1), 3), self.faces_np]),
+        )
+        mesh.clean(inplace=True)
+        mesh.clear_data()
+        mesh = mesh.triangulate()
+        mesh = mesh.decimate(ratio, progress_bar=True)
+
+        # Update vertices and faces
+        self.vertices = torch.tensor(
+            mesh.points, device=self.device, dtype=torch.float32
+        )
+        self.faces = torch.tensor(
+            mesh.faces.reshape(-1, 4)[:, 1:],
+            device=self.device,
+            dtype=torch.int32,
+        )
+
+    @spaces.GPU
+    def __call__(
+        self,
+        filter_ratio: float,
+        max_hole_size: float,
+        resolution: int,
+        num_views: int,
+        norm_mesh_ratio: float = 1.0,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """Post-process the mesh by simplifying and filling holes.
+
+        This method performs a two-step process:
+        1. Simplifies mesh by reducing faces using quadric edge decimation.
+        2. Fills holes by removing invisible faces, repairing small boundaries.
+
+        Args:
+            filter_ratio (float): Ratio of faces to simplify out.
+                Must be in the range (0, 1).
+            max_hole_size (float): Maximum area of a hole to fill. Connected
+                components of holes larger than this size will not be repaired.
+            resolution (int): Resolution of the rasterization buffer.
+            num_views (int): Number of viewpoints to sample for rasterization.
+            norm_mesh_ratio (float, optional): A scaling factor applied to the
+                vertices of the mesh during processing.
+
+        Returns:
+            Tuple[np.ndarray, np.ndarray]:
+                - vertices: Simplified and repaired vertex array of (V, 3).
+                - faces: Simplified and repaired face array of (F, 3).
+        """
+        self.vertices = self.vertices.to(self.device)
+        self.faces = self.faces.to(self.device)
+
+        self.simplify(ratio=filter_ratio)
+        self.fill_holes(
+            max_hole_size=max_hole_size,
+            max_hole_nbe=int(250 * np.sqrt(1 - filter_ratio)),
+            resolution=resolution,
+            num_views=num_views,
+            norm_mesh_ratio=norm_mesh_ratio,
+        )
+
+        return self.vertices_np, self.faces_np

embodied_gen/data/utils.py

@@ -0,0 +1,1039 @@
+# 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.
+
+
+import math
+import os
+import random
+import zipfile
+from shutil import rmtree
+from typing import List, Tuple, Union
+
+import cv2
+import kaolin as kal
+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
+
+logger = logging.getLogger(__name__)
+
+
+__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",
+    "save_mesh_with_mtl",
+    "get_images_from_grid",
+    "post_process_texture",
+    "quat_mult",
+    "quat_to_rotmat",
+    "gamma_shs",
+    "resize_pil",
+    "trellis_preprocess",
+    "delete_dir",
+]
+
+
+class DiffrastRender(object):
+    """A class to handle differentiable rendering using nvdiffrast.
+
+    This class provides methods to render position, depth, and normal maps
+    with optional anti-aliasing and gradient disabling for rasterization.
+
+    Attributes:
+        p_mtx (torch.Tensor): Projection matrix.
+        mv_mtx (torch.Tensor): Model-view matrix.
+        mvp_mtx (torch.Tensor): Model-view-projection matrix, calculated as
+            p_mtx @ mv_mtx if not provided.
+        resolution_hw (Tuple[int, int]): Height and width of the rendering resolution.  # noqa
+        _ctx (Union[dr.RasterizeCudaContext, dr.RasterizeGLContext]): Rasterization context.  # noqa
+        mask_thresh (float): Threshold for mask creation.
+        grad_db (bool): Whether to disable gradients during rasterization.
+        antialias_mask (bool): Whether to apply anti-aliasing to the mask.
+        device (str): Device used for rendering ('cuda' or 'cpu').
+    """
+
+    def __init__(
+        self,
+        p_matrix: torch.Tensor,
+        mv_matrix: torch.Tensor,
+        resolution_hw: Tuple[int, int],
+        context: Union[dr.RasterizeCudaContext, dr.RasterizeGLContext] = None,
+        mvp_matrix: torch.Tensor = None,
+        mask_thresh: float = 0.5,
+        grad_db: bool = False,
+        antialias_mask: bool = True,
+        align_coordinate: bool = True,
+        device: str = "cuda",
+    ) -> None:
+        self.p_mtx = p_matrix
+        self.mv_mtx = mv_matrix
+        if mvp_matrix is None:
+            self.mvp_mtx = torch.bmm(p_matrix, mv_matrix)
+
+        self.resolution_hw = resolution_hw
+        if context is None:
+            context = dr.RasterizeCudaContext(device=device)
+        self._ctx = context
+        self.mask_thresh = mask_thresh
+        self.grad_db = grad_db
+        self.antialias_mask = antialias_mask
+        self.align_coordinate = align_coordinate
+        self.device = device
+
+    def compute_dr_raster(
+        self,
+        vertices: torch.Tensor,
+        faces: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        vertices_clip = self.transform_vertices(vertices, matrix=self.mvp_mtx)
+        rast, _ = dr.rasterize(
+            self._ctx,
+            vertices_clip,
+            faces.int(),
+            resolution=self.resolution_hw,
+            grad_db=self.grad_db,
+        )
+
+        return rast, vertices_clip
+
+    def transform_vertices(
+        self,
+        vertices: torch.Tensor,
+        matrix: torch.Tensor,
+    ) -> torch.Tensor:
+        verts_ones = torch.ones(
+            (len(vertices), 1), device=vertices.device, dtype=vertices.dtype
+        )
+        verts_homo = torch.cat([vertices, verts_ones], dim=-1)
+        trans_vertices = torch.matmul(verts_homo, matrix.permute(0, 2, 1))
+
+        return trans_vertices
+
+    def normalize_map_by_mask_separately(
+        self, map: torch.Tensor, mask: torch.Tensor
+    ) -> torch.Tensor:
+        # Normalize each map separately by mask, normalized map in [0, 1].
+        normalized_maps = []
+        for map_item, mask_item in zip(map, mask):
+            normalized_map = self.normalize_map_by_mask(map_item, mask_item)
+            normalized_maps.append(normalized_map)
+
+        normalized_maps = torch.stack(normalized_maps, dim=0)
+
+        return normalized_maps
+
+    @staticmethod
+    def normalize_map_by_mask(
+        map: torch.Tensor, mask: torch.Tensor
+    ) -> torch.Tensor:
+        # Normalize all maps in total by mask, normalized map in [0, 1].
+        foreground = (mask == 1).squeeze(dim=-1)
+        foreground_elements = map[foreground]
+        if len(foreground_elements) == 0:
+            return map
+
+        min_val, _ = foreground_elements.min(dim=0)
+        max_val, _ = foreground_elements.max(dim=0)
+        val_range = (max_val - min_val).clip(min=1e-6)
+
+        normalized_map = (map - min_val) / val_range
+        normalized_map = torch.lerp(
+            torch.zeros_like(normalized_map), normalized_map, mask
+        )
+        normalized_map[normalized_map < 0] = 0
+
+        return normalized_map
+
+    def _compute_mask(
+        self,
+        rast: torch.Tensor,
+        vertices_clip: torch.Tensor,
+        faces: torch.Tensor,
+    ) -> torch.Tensor:
+        mask = (rast[..., 3:] > 0).float()
+        mask = mask.clip(min=0, max=1)
+
+        if self.antialias_mask is True:
+            mask = dr.antialias(mask, rast, vertices_clip, faces)
+        else:
+            foreground = mask > self.mask_thresh
+            mask[foreground] = 1
+            mask[~foreground] = 0
+
+        return mask
+
+    def render_rast_alpha(
+        self,
+        vertices: torch.Tensor,
+        faces: torch.Tensor,
+    ):
+        faces = faces.to(torch.int32)
+        rast, vertices_clip = self.compute_dr_raster(vertices, faces)
+        mask = self._compute_mask(rast, vertices_clip, faces)
+
+        return mask, rast
+
+    def render_position(
+        self,
+        vertices: torch.Tensor,
+        faces: torch.Tensor,
+    ) -> Union[torch.Tensor, torch.Tensor]:
+        # Vertices in model coordinate system, real position coordinate number.
+        faces = faces.to(torch.int32)
+        mask, rast = self.render_rast_alpha(vertices, faces)
+
+        vertices_model = vertices[None, ...].contiguous().float()
+        position_map, _ = dr.interpolate(vertices_model, rast, faces)
+        # Align with blender.
+        if self.align_coordinate:
+            position_map = position_map[..., [0, 2, 1]]
+            position_map[..., 1] = -position_map[..., 1]
+
+        position_map = torch.lerp(
+            torch.zeros_like(position_map), position_map, mask
+        )
+
+        return position_map, mask
+
+    def render_uv(
+        self,
+        vertices: torch.Tensor,
+        faces: torch.Tensor,
+        vtx_uv: torch.Tensor,
+    ) -> Union[torch.Tensor, torch.Tensor]:
+        faces = faces.to(torch.int32)
+        mask, rast = self.render_rast_alpha(vertices, faces)
+        uv_map, _ = dr.interpolate(vtx_uv, rast, faces)
+        uv_map = torch.lerp(torch.zeros_like(uv_map), uv_map, mask)
+
+        return uv_map, mask
+
+    def render_depth(
+        self,
+        vertices: torch.Tensor,
+        faces: torch.Tensor,
+    ) -> Union[torch.Tensor, torch.Tensor]:
+        # Vertices in model coordinate system, real depth coordinate number.
+        faces = faces.to(torch.int32)
+        mask, rast = self.render_rast_alpha(vertices, faces)
+
+        vertices_camera = self.transform_vertices(vertices, matrix=self.mv_mtx)
+        vertices_camera = vertices_camera[..., 2:3].contiguous().float()
+        depth_map, _ = dr.interpolate(vertices_camera, rast, faces)
+        # Change camera depth minus to positive.
+        if self.align_coordinate:
+            depth_map = -depth_map
+        depth_map = torch.lerp(torch.zeros_like(depth_map), depth_map, mask)
+
+        return depth_map, mask
+
+    def render_global_normal(
+        self,
+        vertices: torch.Tensor,
+        faces: torch.Tensor,
+        vertice_normals: torch.Tensor,
+    ) -> Union[torch.Tensor, torch.Tensor]:
+        # NOTE: vertice_normals in [-1, 1],  return normal in [0, 1].
+        # vertices / vertice_normals in model coordinate system.
+        faces = faces.to(torch.int32)
+        mask, rast = self.render_rast_alpha(vertices, faces)
+        im_base_normals, _ = dr.interpolate(
+            vertice_normals[None, ...].float(), rast, faces
+        )
+
+        if im_base_normals is not None:
+            faces = faces.to(torch.int64)
+            vertices_cam = self.transform_vertices(
+                vertices, matrix=self.mv_mtx
+            )
+            face_vertices_ndc = kal.ops.mesh.index_vertices_by_faces(
+                vertices_cam[..., :3], faces
+            )
+            face_normal_sign = kal.ops.mesh.face_normals(face_vertices_ndc)[
+                ..., 2
+            ]
+            for idx in range(len(im_base_normals)):
+                face_idx = (rast[idx, ..., -1].long() - 1).contiguous()
+                im_normal_sign = torch.sign(face_normal_sign[idx, face_idx])
+                im_normal_sign[face_idx == -1] = 0
+                im_base_normals[idx] *= im_normal_sign.unsqueeze(-1)
+
+        normal = (im_base_normals + 1) / 2
+        normal = normal.clip(min=0, max=1)
+        normal = torch.lerp(torch.zeros_like(normal), normal, mask)
+
+        return normal, mask
+
+    def transform_normal(
+        self,
+        normals: torch.Tensor,
+        trans_matrix: torch.Tensor,
+        masks: torch.Tensor,
+        to_view: bool,
+    ) -> torch.Tensor:
+        # NOTE: input normals in [0, 1], output normals in [0, 1].
+        normals = normals.clone()
+        assert len(normals) == len(trans_matrix)
+
+        if not to_view:
+            # Flip the sign on the x-axis to match inv bae system for global transformation.  # noqa
+            normals[..., 0] = 1 - normals[..., 0]
+
+        normals = 2 * normals - 1
+        b, h, w, c = normals.shape
+
+        transformed_normals = []
+        for normal, matrix in zip(normals, trans_matrix):
+            # Transform normals using the transformation matrix (4x4).
+            reshaped_normals = normal.view(-1, c)  # (h w 3) -> (hw 3)
+            padded_vectors = torch.nn.functional.pad(
+                reshaped_normals, pad=(0, 1), mode="constant", value=0.0
+            )
+            transformed_normal = torch.matmul(
+                padded_vectors, matrix.transpose(0, 1)
+            )[..., :3]
+
+            # Normalize and clip the normals to [0, 1] range.
+            transformed_normal = F.normalize(transformed_normal, p=2, dim=-1)
+            transformed_normal = (transformed_normal + 1) / 2
+
+            if to_view:
+                # Flip the sign on the x-axis to match bae system for view transformation.  # noqa
+                transformed_normal[..., 0] = 1 - transformed_normal[..., 0]
+
+            transformed_normals.append(transformed_normal.view(h, w, c))
+
+        transformed_normals = torch.stack(transformed_normals, dim=0)
+
+        if masks is not None:
+            transformed_normals = torch.lerp(
+                torch.zeros_like(transformed_normals),
+                transformed_normals,
+                masks,
+            )
+
+        return transformed_normals
+
+
+def _az_el_to_points(
+    azimuths: np.ndarray, elevations: np.ndarray
+) -> np.ndarray:
+    x = np.cos(azimuths) * np.cos(elevations)
+    y = np.sin(azimuths) * np.cos(elevations)
+    z = np.sin(elevations)
+
+    return np.stack([x, y, z], axis=-1)
+
+
+def _compute_az_el_by_views(
+    num_view: int, el: float
+) -> Tuple[np.ndarray, np.ndarray]:
+    azimuths = np.arange(num_view) / num_view * np.pi * 2
+    elevations = np.deg2rad(np.array([el] * num_view))
+
+    return azimuths, elevations
+
+
+def _compute_cam_pts_by_az_el(
+    azs: np.ndarray,
+    els: np.ndarray,
+    distance: float,
+    extra_pts: np.ndarray = None,
+) -> np.ndarray:
+    distances = np.array([distance for _ in range(len(azs))])
+    cam_pts = _az_el_to_points(azs, els) * distances[:, None]
+
+    if extra_pts is not None:
+        cam_pts = np.concatenate([cam_pts, extra_pts], axis=0)
+
+    # Align coordinate system.
+    cam_pts = cam_pts[:, [0, 2, 1]]  # xyz -> xzy
+    cam_pts[..., 2] = -cam_pts[..., 2]
+
+    return cam_pts
+
+
+def compute_cam_pts_by_views(
+    num_view: int, el: float, distance: float, extra_pts: np.ndarray = None
+) -> torch.Tensor:
+    """Computes object-center camera points for a given number of views.
+
+    Args:
+        num_view (int): The number of views (camera positions) to compute.
+        el (float): The elevation angle in degrees.
+        distance (float): The distance from the origin to the camera.
+        extra_pts (np.ndarray): Extra camera points postion.
+
+    Returns:
+        torch.Tensor: A tensor containing the camera points for each view, with shape `(num_view, 3)`. # noqa
+    """
+    azimuths, elevations = _compute_az_el_by_views(num_view, el)
+    cam_pts = _compute_cam_pts_by_az_el(
+        azimuths, elevations, distance, extra_pts
+    )
+
+    return cam_pts
+
+
+def save_images(
+    images: Union[list[np.ndarray], list[torch.Tensor]],
+    output_dir: str,
+    cvt_color: str = None,
+    format: str = ".png",
+    to_uint8: bool = True,
+    verbose: bool = False,
+) -> List[str]:
+    # NOTE: images in [0, 1]
+    os.makedirs(output_dir, exist_ok=True)
+    save_paths = []
+    for idx, image in enumerate(images):
+        if isinstance(image, torch.Tensor):
+            image = image.detach().cpu().numpy()
+        if to_uint8:
+            image = image.clip(min=0, max=1)
+            image = (255.0 * image).astype(np.uint8)
+        if cvt_color is not None:
+            image = cv2.cvtColor(image, cvt_color)
+        save_path = os.path.join(output_dir, f"{idx:04d}{format}")
+        save_paths.append(save_path)
+
+        cv2.imwrite(save_path, image)
+
+    if verbose:
+        logger.info(f"Images saved in {output_dir}")
+
+    return save_paths
+
+
+def _current_lighting(
+    azimuths: List[float],
+    elevations: List[float],
+    light_factor: float = 1.0,
+    device: str = "cuda",
+):
+    # azimuths, elevations in degress.
+    directions = []
+    for az, el in zip(azimuths, elevations):
+        az, el = math.radians(az), math.radians(el)
+        direction = kal.render.lighting.sg_direction_from_azimuth_elevation(
+            az, el
+        )
+        directions.append(direction)
+    directions = torch.cat(directions, dim=0)
+
+    amplitude = torch.ones_like(directions) * light_factor
+    light_condition = kal.render.lighting.SgLightingParameters(
+        amplitude=amplitude,
+        direction=directions,
+        sharpness=3,
+    ).to(device)
+
+    # light_condition = kal.render.lighting.SgLightingParameters.from_sun(
+    #     directions, strength=1, angle=90, color=None
+    # ).to(device)
+
+    return light_condition
+
+
+def render_pbr(
+    mesh,
+    camera,
+    device="cuda",
+    cxt=None,
+    custom_materials=None,
+    light_factor=1.0,
+):
+    if cxt is None:
+        cxt = dr.RasterizeCudaContext()
+
+    light_condition = _current_lighting(
+        azimuths=[0, 90, 180, 270],
+        elevations=[90, 60, 30, 20],
+        light_factor=light_factor,
+        device=device,
+    )
+    render_res = kal.render.easy_render.render_mesh(
+        camera,
+        mesh,
+        lighting=light_condition,
+        nvdiffrast_context=cxt,
+        custom_materials=custom_materials,
+    )
+
+    image = render_res[kal.render.easy_render.RenderPass.render]
+    image = image.clip(0, 1)
+
+    albedo = render_res[kal.render.easy_render.RenderPass.albedo]
+    albedo = albedo.clip(0, 1)
+
+    diffuse = render_res[kal.render.easy_render.RenderPass.diffuse]
+    diffuse = diffuse.clip(0, 1)
+
+    normal = render_res[kal.render.easy_render.RenderPass.normals]
+    normal = normal.clip(-1, 1)
+
+    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
+    normalize: bool = False,
+) -> torch.Tensor:  # F,3
+    full_vertices = vertices[faces]  # F,C=3,3
+    v0, v1, v2 = full_vertices.unbind(dim=1)  # F,3
+    face_normals = torch.cross(v1 - v0, v2 - v0, dim=1)  # F,3
+    if normalize:
+        face_normals = F.normalize(
+            face_normals, eps=1e-6, dim=1
+        )  # TODO inplace?
+    return face_normals  # F,3
+
+
+def calc_vertex_normals(
+    vertices: torch.Tensor,  # V,3 first vertex may be unreferenced
+    faces: torch.Tensor,  # F,3 long, first face may be all zero
+    face_normals: torch.Tensor = None,  # F,3, not normalized
+) -> torch.Tensor:  # F,3
+    _F = faces.shape[0]
+
+    if face_normals is None:
+        face_normals = _calc_face_normals(vertices, faces)
+
+    vertex_normals = torch.zeros(
+        (vertices.shape[0], 3, 3), dtype=vertices.dtype, device=vertices.device
+    )  # V,C=3,3
+    vertex_normals.scatter_add_(
+        dim=0,
+        index=faces[:, :, None].expand(_F, 3, 3),
+        src=face_normals[:, None, :].expand(_F, 3, 3),
+    )
+    vertex_normals = vertex_normals.sum(dim=1)  # V,3
+    return F.normalize(vertex_normals, eps=1e-6, dim=1)
+
+
+def normalize_vertices_array(
+    vertices: Union[torch.Tensor, np.ndarray],
+    mesh_scale: float = 1.0,
+    exec_norm: bool = True,
+):
+    if isinstance(vertices, torch.Tensor):
+        bbmin, bbmax = vertices.min(0)[0], vertices.max(0)[0]
+    else:
+        bbmin, bbmax = vertices.min(0), vertices.max(0)  # (3,)
+    center = (bbmin + bbmax) * 0.5
+    bbsize = bbmax - bbmin
+    scale = 2 * mesh_scale / bbsize.max()
+    if exec_norm:
+        vertices = (vertices - center) * scale
+
+    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
+    elif isinstance(obj, set):
+        return list(obj)
+    elif obj is None:
+        return obj
+    else:
+        return [obj]
+
+
+@dataclass
+class CameraSetting:
+    """Camera settings for images rendering."""
+
+    num_images: int
+    elevation: list[float]
+    distance: float
+    resolution_hw: tuple[int, int]
+    fov: float
+    at: tuple[float, float, float] = field(
+        default_factory=lambda: (0.0, 0.0, 0.0)
+    )
+    up: tuple[float, float, float] = field(
+        default_factory=lambda: (0.0, 1.0, 0.0)
+    )
+    device: str = "cuda"
+    near: float = 1e-2
+    far: float = 1e2
+
+    def __post_init__(
+        self,
+    ):
+        h = self.resolution_hw[0]
+        f = (h / 2) / math.tan(self.fov / 2)
+        cx = self.resolution_hw[1] / 2
+        cy = self.resolution_hw[0] / 2
+        Ks = [
+            [f, 0, cx],
+            [0, f, cy],
+            [0, 0, 1],
+        ]
+
+        self.Ks = Ks
+
+
+def _compute_az_el_by_camera_params(
+    camera_params: CameraSetting, flip_az: bool = False
+):
+    num_view = camera_params.num_images // len(camera_params.elevation)
+    view_interval = 2 * np.pi / num_view / 2
+    if num_view == 1:
+        view_interval = np.pi / 2
+    azimuths = []
+    elevations = []
+    for idx, el in enumerate(camera_params.elevation):
+        azs = np.arange(num_view) / num_view * np.pi * 2 + idx * view_interval
+        if flip_az:
+            azs *= -1
+        els = np.deg2rad(np.array([el] * num_view))
+        azimuths.append(azs)
+        elevations.append(els)
+
+    azimuths = np.concatenate(azimuths, axis=0)
+    elevations = np.concatenate(elevations, axis=0)
+
+    return azimuths, elevations
+
+
+def init_kal_camera(
+    camera_params: CameraSetting,
+    flip_az: bool = False,
+) -> Camera:
+    azimuths, elevations = _compute_az_el_by_camera_params(
+        camera_params, flip_az
+    )
+    cam_pts = _compute_cam_pts_by_az_el(
+        azimuths, elevations, camera_params.distance
+    )
+
+    up = torch.cat(
+        [
+            torch.tensor(camera_params.up).repeat(camera_params.num_images, 1),
+        ],
+        dim=0,
+    )
+
+    camera = Camera.from_args(
+        eye=torch.tensor(cam_pts),
+        at=torch.tensor(camera_params.at),
+        up=up,
+        fov=camera_params.fov,
+        height=camera_params.resolution_hw[0],
+        width=camera_params.resolution_hw[1],
+        near=camera_params.near,
+        far=camera_params.far,
+        device=camera_params.device,
+    )
+
+    return camera
+
+
+def import_kaolin_mesh(mesh_path: str, with_mtl: bool = False):
+    if mesh_path.endswith(".glb"):
+        mesh = kal.io.gltf.import_mesh(mesh_path)
+    elif mesh_path.endswith(".obj"):
+        with_material = True if with_mtl else False
+        mesh = kal.io.obj.import_mesh(mesh_path, with_materials=with_material)
+        if with_mtl and mesh.materials and len(mesh.materials) > 0:
+            material = kal.render.materials.PBRMaterial()
+            assert (
+                "map_Kd" in mesh.materials[0]
+            ), "'map_Kd' not found in materials."
+            material.diffuse_texture = mesh.materials[0]["map_Kd"] / 255.0
+            mesh.materials = [material]
+    elif mesh_path.endswith(".ply"):
+        mesh = trimesh.load(mesh_path)
+        mesh_path = mesh_path.replace(".ply", ".obj")
+        mesh.export(mesh_path)
+        mesh = kal.io.obj.import_mesh(mesh_path)
+    elif mesh_path.endswith(".off"):
+        mesh = kal.io.off.import_mesh(mesh_path)
+    else:
+        raise RuntimeError(
+            f"{mesh_path} mesh type not supported, "
+            "supported mesh type `.glb`, `.obj`, `.ply`, `.off`."
+        )
+
+    return mesh
+
+
+def save_mesh_with_mtl(
+    vertices: np.ndarray,
+    faces: np.ndarray,
+    uvs: np.ndarray,
+    texture: Union[Image.Image, np.ndarray],
+    output_path: str,
+    material_base=(250, 250, 250, 255),
+) -> trimesh.Trimesh:
+    if isinstance(texture, np.ndarray):
+        texture = Image.fromarray(texture)
+
+    mesh = trimesh.Trimesh(
+        vertices,
+        faces,
+        visual=trimesh.visual.TextureVisuals(uv=uvs, image=texture),
+    )
+    mesh.visual.material = trimesh.visual.material.SimpleMaterial(
+        image=texture,
+        diffuse=material_base,
+        ambient=material_base,
+        specular=material_base,
+    )
+
+    dir_name = os.path.dirname(output_path)
+    os.makedirs(dir_name, exist_ok=True)
+
+    _ = mesh.export(output_path)
+    # texture.save(os.path.join(dir_name, f"{file_name}_texture.png"))
+
+    logger.info(f"Saved mesh with texture to {output_path}")
+
+    return mesh
+
+
+def get_images_from_grid(
+    image: Union[str, Image.Image], img_size: int
+) -> list[Image.Image]:
+    if isinstance(image, str):
+        image = Image.open(image)
+
+    view_images = np.array(image)
+    height, width, _ = view_images.shape
+    rows = height // img_size
+    cols = width // img_size
+    blocks = []
+    for i in range(rows):
+        for j in range(cols):
+            block = view_images[
+                i * img_size : (i + 1) * img_size,
+                j * img_size : (j + 1) * img_size,
+                :,
+            ]
+            blocks.append(Image.fromarray(block))
+
+    return blocks
+
+
+def enhance_image(
+    image: Image.Image,
+    contrast_factor: float = 1.3,
+    color_factor: float = 1.2,
+    brightness_factor: float = 0.95,
+) -> Image.Image:
+    enhancer_contrast = ImageEnhance.Contrast(image)
+    img_contrasted = enhancer_contrast.enhance(contrast_factor)
+
+    enhancer_color = ImageEnhance.Color(img_contrasted)
+    img_colored = enhancer_color.enhance(color_factor)
+
+    enhancer_brightness = ImageEnhance.Brightness(img_colored)
+    enhanced_image = enhancer_brightness.enhance(brightness_factor)
+
+    return enhanced_image
+
+
+def post_process_texture(texture: np.ndarray, iter: int = 1) -> np.ndarray:
+    for _ in range(iter):
+        texture = cv2.fastNlMeansDenoisingColored(texture, None, 2, 2, 7, 15)
+        texture = cv2.bilateralFilter(
+            texture, d=5, sigmaColor=20, sigmaSpace=20
+        )
+
+    texture = enhance_image(
+        image=Image.fromarray(texture),
+        contrast_factor=1.3,
+        color_factor=1.2,
+        brightness_factor=0.95,
+    )
+
+    return np.array(texture)
+
+
+def quat_mult(q1, q2):
+    # NOTE:
+    # Q1 is the quaternion that rotates the vector from the original position to the final position  # noqa
+    # Q2 is the quaternion that been rotated
+    w1, x1, y1, z1 = q1.T
+    w2, x2, y2, z2 = q2.T
+    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
+    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
+    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
+    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
+    return torch.stack([w, x, y, z]).T
+
+
+def quat_to_rotmat(quats: torch.Tensor, mode="wxyz") -> torch.Tensor:
+    """Convert quaternion to rotation matrix."""
+    quats = F.normalize(quats, p=2, dim=-1)
+
+    if mode == "xyzw":
+        x, y, z, w = torch.unbind(quats, dim=-1)
+    elif mode == "wxyz":
+        w, x, y, z = torch.unbind(quats, dim=-1)
+    else:
+        raise ValueError(f"Invalid mode: {mode}.")
+
+    R = torch.stack(
+        [
+            1 - 2 * (y**2 + z**2),
+            2 * (x * y - w * z),
+            2 * (x * z + w * y),
+            2 * (x * y + w * z),
+            1 - 2 * (x**2 + z**2),
+            2 * (y * z - w * x),
+            2 * (x * z - w * y),
+            2 * (y * z + w * x),
+            1 - 2 * (x**2 + y**2),
+        ],
+        dim=-1,
+    )
+
+    return R.reshape(quats.shape[:-1] + (3, 3))
+
+
+def gamma_shs(shs: torch.Tensor, gamma: float) -> torch.Tensor:
+    C0 = 0.28209479177387814  # Constant for normalization in spherical harmonics  # noqa
+    # Clip to the range [0.0, 1.0], apply gamma correction, and then un-clip back  # noqa
+    new_shs = torch.clip(shs * C0 + 0.5, 0.0, 1.0)
+    new_shs = (torch.pow(new_shs, gamma) - 0.5) / C0
+    return new_shs
+
+
+def resize_pil(image: Image.Image, max_size: int = 1024) -> Image.Image:
+    max_size = max(image.size)
+    scale = min(1, 1024 / max_size)
+    if scale < 1:
+        new_size = (int(image.width * scale), int(image.height * scale))
+        image = image.resize(new_size, Image.Resampling.LANCZOS)
+
+    return image
+
+
+def trellis_preprocess(image: Image.Image) -> Image.Image:
+    """Process the input image as trellis done."""
+    image_np = np.array(image)
+    alpha = image_np[:, :, 3]
+    bbox = np.argwhere(alpha > 0.8 * 255)
+    bbox = (
+        np.min(bbox[:, 1]),
+        np.min(bbox[:, 0]),
+        np.max(bbox[:, 1]),
+        np.max(bbox[:, 0]),
+    )
+    center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
+    size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
+    size = int(size * 1.2)
+    bbox = (
+        center[0] - size // 2,
+        center[1] - size // 2,
+        center[0] + size // 2,
+        center[1] + size // 2,
+    )
+    image = image.crop(bbox)
+    image = image.resize((518, 518), Image.Resampling.LANCZOS)
+    image = np.array(image).astype(np.float32) / 255
+    image = image[:, :, :3] * image[:, :, 3:4]
+    image = Image.fromarray((image * 255).astype(np.uint8))
+
+    return image
+
+
+def zip_files(input_paths: list[str], output_zip: str) -> str:
+    with zipfile.ZipFile(output_zip, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for input_path in input_paths:
+            if not os.path.exists(input_path):
+                raise FileNotFoundError(f"File not found: {input_path}")
+
+            if os.path.isdir(input_path):
+                for root, _, files in os.walk(input_path):
+                    for file in files:
+                        file_path = os.path.join(root, file)
+                        arcname = os.path.relpath(
+                            file_path, start=os.path.commonpath(input_paths)
+                        )
+                        zipf.write(file_path, arcname=arcname)
+            else:
+                arcname = os.path.relpath(
+                    input_path, start=os.path.commonpath(input_paths)
+                )
+                zipf.write(input_path, arcname=arcname)
+
+    return output_zip
+
+
+def delete_dir(folder_path: str, keep_subs: list[str] = None) -> None:
+    for item in os.listdir(folder_path):
+        if keep_subs is not None and item in keep_subs:
+            continue
+        item_path = os.path.join(folder_path, item)
+        if os.path.isdir(item_path):
+            rmtree(item_path)
+        else:
+            os.remove(item_path)

embodied_gen/envs/pick_embodiedgen.py

@@ -0,0 +1,420 @@
+# 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.
+
+import json
+import os
+
+import numpy as np
+import sapien
+import torch
+import torchvision.transforms as transforms
+from mani_skill.envs.sapien_env import BaseEnv
+from mani_skill.sensors.camera import CameraConfig
+from mani_skill.utils import sapien_utils
+from mani_skill.utils.building import actors
+from mani_skill.utils.building.ground import build_ground
+from mani_skill.utils.registration import register_env
+from mani_skill.utils.structs.actor import Actor
+from mani_skill.utils.structs.pose import Pose
+from mani_skill.utils.structs.types import (
+    GPUMemoryConfig,
+    SceneConfig,
+    SimConfig,
+)
+from mani_skill.utils.visualization.misc import tile_images
+from tqdm import tqdm
+from embodied_gen.models.gs_model import GaussianOperator
+from embodied_gen.utils.enum import LayoutInfo, Scene3DItemEnum
+from embodied_gen.utils.geometry import bfs_placement, quaternion_multiply
+from embodied_gen.utils.log import logger
+from embodied_gen.utils.process_media import alpha_blend_rgba
+from embodied_gen.utils.simulation import (
+    SIM_COORD_ALIGN,
+    load_assets_from_layout_file,
+)
+
+__all__ = ["PickEmbodiedGen"]
+
+
+@register_env("PickEmbodiedGen-v1", max_episode_steps=100)
+class PickEmbodiedGen(BaseEnv):
+    SUPPORTED_ROBOTS = ["panda", "panda_wristcam", "fetch"]
+    goal_thresh = 0.0
+
+    def __init__(
+        self,
+        *args,
+        robot_uids: str | list[str] = "panda",
+        robot_init_qpos_noise: float = 0.02,
+        num_envs: int = 1,
+        reconfiguration_freq: int = None,
+        **kwargs,
+    ):
+        self.robot_init_qpos_noise = robot_init_qpos_noise
+        if reconfiguration_freq is None:
+            if num_envs == 1:
+                reconfiguration_freq = 1
+            else:
+                reconfiguration_freq = 0
+
+        # Init params from kwargs.
+        layout_file = kwargs.pop("layout_file", None)
+        replace_objs = kwargs.pop("replace_objs", True)
+        self.enable_grasp = kwargs.pop("enable_grasp", False)
+        self.init_3dgs_quat = kwargs.pop(
+            "init_3dgs_quat", [0.7071, 0, 0, 0.7071]
+        )
+        # Add small offset in z-axis to avoid collision.
+        self.objs_z_offset = kwargs.pop("objs_z_offset", 0.002)
+        self.robot_z_offset = kwargs.pop("robot_z_offset", 0.002)
+        self.camera_cfg = kwargs.pop("camera_cfg", None)
+        if self.camera_cfg is None:
+            self.camera_cfg = dict(
+                camera_eye=[0.9, 0.0, 1.1],
+                camera_target_pt=[0.0, 0.0, 0.9],
+                image_hw=[256, 256],
+                fovy_deg=75,
+            )
+
+        self.layouts = self.init_env_layouts(
+            layout_file, num_envs, replace_objs
+        )
+        self.robot_pose = self.compute_robot_init_pose(
+            self.layouts, num_envs, self.robot_z_offset
+        )
+        self.env_actors = dict()
+        self.image_transform = transforms.PILToTensor()
+
+        super().__init__(
+            *args,
+            robot_uids=robot_uids,
+            reconfiguration_freq=reconfiguration_freq,
+            num_envs=num_envs,
+            **kwargs,
+        )
+
+        self.bg_images = dict()
+        if self.render_mode == "hybrid":
+            self.bg_images = self.render_gs3d_images(
+                self.layouts, num_envs, self.init_3dgs_quat
+            )
+
+    @staticmethod
+    def init_env_layouts(
+        layout_file: str, num_envs: int, replace_objs: bool
+    ) -> list[LayoutInfo]:
+        layouts = []
+        for env_idx in range(num_envs):
+            if replace_objs and env_idx > 0:
+                layout_info = bfs_placement(layout_file)
+            else:
+                layout_info = json.load(open(layout_file, "r"))
+                layout_info = LayoutInfo.from_dict(layout_info)
+
+            layout_path = layout_file.replace(".json", f"_env{env_idx}.json")
+            with open(layout_path, "w") as f:
+                json.dump(layout_info.to_dict(), f, indent=4)
+
+            layouts.append(layout_path)
+
+        return layouts
+
+    @staticmethod
+    def compute_robot_init_pose(
+        layouts: list[str], num_envs: int, z_offset: float = 0.0
+    ) -> list[list[float]]:
+        robot_pose = []
+        for env_idx in range(num_envs):
+            layout = json.load(open(layouts[env_idx], "r"))
+            layout = LayoutInfo.from_dict(layout)
+            robot_node = layout.relation[Scene3DItemEnum.ROBOT.value]
+            x, y, z, qx, qy, qz, qw = layout.position[robot_node]
+            robot_pose.append([x, y, z + z_offset, qw, qx, qy, qz])
+
+        return robot_pose
+
+    @property
+    def _default_sim_config(self):
+        return SimConfig(
+            scene_config=SceneConfig(
+                solver_position_iterations=30,
+                # contact_offset=0.04,
+                # rest_offset=0.001,
+            ),
+            # sim_freq=200,
+            control_freq=50,
+            gpu_memory_config=GPUMemoryConfig(
+                max_rigid_contact_count=2**20, max_rigid_patch_count=2**19
+            ),
+        )
+
+    @property
+    def _default_sensor_configs(self):
+        pose = sapien_utils.look_at(eye=[0.3, 0, 0.6], target=[-0.1, 0, 0.1])
+
+        return [
+            CameraConfig("base_camera", pose, 128, 128, np.pi / 2, 0.01, 100)
+        ]
+
+    @property
+    def _default_human_render_camera_configs(self):
+        pose = sapien_utils.look_at(
+            eye=self.camera_cfg["camera_eye"],
+            target=self.camera_cfg["camera_target_pt"],
+        )
+
+        return CameraConfig(
+            "render_camera",
+            pose,
+            self.camera_cfg["image_hw"][1],
+            self.camera_cfg["image_hw"][0],
+            np.deg2rad(self.camera_cfg["fovy_deg"]),
+            0.01,
+            100,
+        )
+
+    def _load_agent(self, options: dict):
+        self.ground = build_ground(self.scene)
+        super()._load_agent(options, sapien.Pose(p=[-10, 0, 10]))
+
+    def _load_scene(self, options: dict):
+        all_objects = []
+        logger.info(f"Loading EmbodiedGen assets...")
+        for env_idx in range(self.num_envs):
+            env_actors = load_assets_from_layout_file(
+                self.scene,
+                self.layouts[env_idx],
+                z_offset=self.objs_z_offset,
+                env_idx=env_idx,
+            )
+            self.env_actors[f"env{env_idx}"] = env_actors
+            all_objects.extend(env_actors.values())
+
+        self.obj = all_objects[-1]
+        for obj in all_objects:
+            self.remove_from_state_dict_registry(obj)
+
+        self.all_objects = Actor.merge(all_objects, name="all_objects")
+        self.add_to_state_dict_registry(self.all_objects)
+
+        self.goal_site = actors.build_sphere(
+            self.scene,
+            radius=self.goal_thresh,
+            color=[0, 1, 0, 0],
+            name="goal_site",
+            body_type="kinematic",
+            add_collision=False,
+            initial_pose=sapien.Pose(),
+        )
+        self._hidden_objects.append(self.goal_site)
+
+    def _initialize_episode(self, env_idx: torch.Tensor, options: dict):
+        with torch.device(self.device):
+            b = len(env_idx)
+            goal_xyz = torch.zeros((b, 3))
+            goal_xyz[:, :2] = torch.rand((b, 2)) * 0.2 - 0.1
+            self.goal_site.set_pose(Pose.create_from_pq(goal_xyz))
+
+            qpos = np.array(
+                [
+                    0.0,
+                    np.pi / 8,
+                    0,
+                    -np.pi * 3 / 8,
+                    0,
+                    np.pi * 3 / 4,
+                    np.pi / 4,
+                    0.04,
+                    0.04,
+                ]
+            )
+            qpos = (
+                np.random.normal(
+                    0, self.robot_init_qpos_noise, (self.num_envs, len(qpos))
+                )
+                + qpos
+            )
+            qpos[:, -2:] = 0.04
+            self.agent.robot.set_root_pose(np.array(self.robot_pose))
+            self.agent.reset(qpos)
+            self.agent.init_qpos = qpos
+            self.agent.controller.controllers["gripper"].reset()
+
+    def render_gs3d_images(
+        self, layouts: list[str], num_envs: int, init_quat: list[float]
+    ) -> dict[str, np.ndarray]:
+        sim_coord_align = (
+            torch.tensor(SIM_COORD_ALIGN).to(torch.float32).to(self.device)
+        )
+        cameras = self.scene.sensors.copy()
+        cameras.update(self.scene.human_render_cameras)
+
+        # Preload the background Gaussian Splatting model.
+        asset_root = os.path.dirname(layouts[0])
+        layout = LayoutInfo.from_dict(json.load(open(layouts[0], "r")))
+        bg_node = layout.relation[Scene3DItemEnum.BACKGROUND.value]
+        gs_path = os.path.join(
+            asset_root, layout.assets[bg_node], "gs_model.ply"
+        )
+        raw_gs: GaussianOperator = GaussianOperator.load_from_ply(gs_path)
+        bg_images = dict()
+        for env_idx in tqdm(range(num_envs), desc="Pre-rendering Background"):
+            layout = json.load(open(layouts[env_idx], "r"))
+            layout = LayoutInfo.from_dict(layout)
+            x, y, z, qx, qy, qz, qw = layout.position[bg_node]
+            qx, qy, qz, qw = quaternion_multiply([qx, qy, qz, qw], init_quat)
+            init_pose = torch.tensor([x, y, z, qx, qy, qz, qw])
+            gs_model = raw_gs.get_gaussians(instance_pose=init_pose)
+            for key in cameras:
+                camera = cameras[key]
+                Ks = camera.camera.get_intrinsic_matrix()  # (n_env, 3, 3)
+                c2w = camera.camera.get_model_matrix()  # (n_env, 4, 4)
+                result = gs_model.render(
+                    c2w[env_idx] @ sim_coord_align,
+                    Ks[env_idx],
+                    image_width=camera.config.width,
+                    image_height=camera.config.height,
+                )
+                bg_images[f"{key}-env{env_idx}"] = result.rgb[..., ::-1]
+
+        return bg_images
+
+    def render(self):
+        if self.render_mode is None:
+            raise RuntimeError("render_mode is not set.")
+        if self.render_mode == "human":
+            return self.render_human()
+        elif self.render_mode == "rgb_array":
+            res = self.render_rgb_array()
+            return res
+        elif self.render_mode == "sensors":
+            res = self.render_sensors()
+            return res
+        elif self.render_mode == "all":
+            return self.render_all()
+        elif self.render_mode == "hybrid":
+            return self.hybrid_render()
+        else:
+            raise NotImplementedError(
+                f"Unsupported render mode {self.render_mode}."
+            )
+
+    def render_rgb_array(
+        self, camera_name: str = None, return_alpha: bool = False
+    ):
+        for obj in self._hidden_objects:
+            obj.show_visual()
+        self.scene.update_render(
+            update_sensors=False, update_human_render_cameras=True
+        )
+        images = []
+        render_images = self.scene.get_human_render_camera_images(
+            camera_name, return_alpha
+        )
+        for image in render_images.values():
+            images.append(image)
+        if len(images) == 0:
+            return None
+        if len(images) == 1:
+            return images[0]
+        for obj in self._hidden_objects:
+            obj.hide_visual()
+        return tile_images(images)
+
+    def render_sensors(self):
+        images = []
+        sensor_images = self.get_sensor_images()
+        for image in sensor_images.values():
+            for img in image.values():
+                images.append(img)
+        return tile_images(images)
+
+    def hybrid_render(self):
+        fg_images = self.render_rgb_array(
+            return_alpha=True
+        )  # (n_env, h, w, 3)
+        images = []
+        for key in self.bg_images:
+            if "render_camera" not in key:
+                continue
+            env_idx = int(key.split("-env")[-1])
+            rgba = alpha_blend_rgba(
+                fg_images[env_idx].cpu().numpy(), self.bg_images[key]
+            )
+            images.append(self.image_transform(rgba))
+
+        images = torch.stack(images, dim=0)
+        images = images.permute(0, 2, 3, 1)
+
+        return images[..., :3]
+
+    def evaluate(self):
+        obj_to_goal_pos = (
+            self.obj.pose.p
+        )  # self.goal_site.pose.p - self.obj.pose.p
+        is_obj_placed = (
+            torch.linalg.norm(obj_to_goal_pos, axis=1) <= self.goal_thresh
+        )
+        is_grasped = self.agent.is_grasping(self.obj)
+        is_robot_static = self.agent.is_static(0.2)
+
+        return dict(
+            is_grasped=is_grasped,
+            obj_to_goal_pos=obj_to_goal_pos,
+            is_obj_placed=is_obj_placed,
+            is_robot_static=is_robot_static,
+            is_grasping=self.agent.is_grasping(self.obj),
+            success=torch.logical_and(is_obj_placed, is_robot_static),
+        )
+
+    def _get_obs_extra(self, info: dict):
+
+        return dict()
+
+    def compute_dense_reward(self, obs: any, action: torch.Tensor, info: dict):
+        tcp_to_obj_dist = torch.linalg.norm(
+            self.obj.pose.p - self.agent.tcp.pose.p, axis=1
+        )
+        reaching_reward = 1 - torch.tanh(5 * tcp_to_obj_dist)
+        reward = reaching_reward
+
+        is_grasped = info["is_grasped"]
+        reward += is_grasped
+
+        # obj_to_goal_dist = torch.linalg.norm(
+        #     self.goal_site.pose.p - self.obj.pose.p, axis=1
+        # )
+        obj_to_goal_dist = torch.linalg.norm(
+            self.obj.pose.p - self.obj.pose.p, axis=1
+        )
+        place_reward = 1 - torch.tanh(5 * obj_to_goal_dist)
+        reward += place_reward * is_grasped
+
+        reward += info["is_obj_placed"] * is_grasped
+
+        static_reward = 1 - torch.tanh(
+            5
+            * torch.linalg.norm(self.agent.robot.get_qvel()[..., :-2], axis=1)
+        )
+        reward += static_reward * info["is_obj_placed"] * is_grasped
+
+        reward[info["success"]] = 6
+        return reward
+
+    def compute_normalized_dense_reward(
+        self, obs: any, action: torch.Tensor, info: dict
+    ):
+        return self.compute_dense_reward(obs=obs, action=action, info=info) / 6

embodied_gen/models/delight_model.py

@@ -0,0 +1,202 @@
+# 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.
+
+
+import os
+from typing import Union
+
+import cv2
+import numpy as np
+import spaces
+import torch
+from diffusers import (
+    EulerAncestralDiscreteScheduler,
+    StableDiffusionInstructPix2PixPipeline,
+)
+from huggingface_hub import snapshot_download
+from PIL import Image
+from embodied_gen.models.segment_model import RembgRemover
+from embodied_gen.utils.log import logger
+
+__all__ = [
+    "DelightingModel",
+]
+
+
+class DelightingModel(object):
+    """A model to remove the lighting in image space.
+
+    This model is encapsulated based on the Hunyuan3D-Delight model
+    from https://huggingface.co/tencent/Hunyuan3D-2/tree/main/hunyuan3d-delight-v2-0 # noqa
+
+    Attributes:
+        image_guide_scale (float): Weight of image guidance in diffusion process.
+        text_guide_scale (float): Weight of text (prompt) guidance in diffusion process.
+        num_infer_step (int): Number of inference steps for diffusion model.
+        mask_erosion_size (int): Size of erosion kernel for alpha mask cleanup.
+        device (str): Device used for inference, e.g., 'cuda' or 'cpu'.
+        seed (int): Random seed for diffusion model reproducibility.
+        model_path (str): Filesystem path to pretrained model weights.
+        pipeline: Lazy-loaded diffusion pipeline instance.
+    """
+
+    def __init__(
+        self,
+        model_path: str = None,
+        num_infer_step: int = 50,
+        mask_erosion_size: int = 3,
+        image_guide_scale: float = 1.5,
+        text_guide_scale: float = 1.0,
+        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
+        self.mask_erosion_size = mask_erosion_size
+        self.kernel = np.ones(
+            (self.mask_erosion_size, self.mask_erosion_size), np.uint8
+        )
+        self.seed = seed
+        self.device = device
+        self.pipeline = None  # lazy load model adapt to @spaces.GPU
+
+        if model_path is None:
+            suffix = "hunyuan3d-delight-v2-0"
+            model_path = snapshot_download(
+                repo_id="tencent/Hunyuan3D-2", allow_patterns=f"{suffix}/*"
+            )
+            model_path = os.path.join(model_path, suffix)
+
+        self.model_path = model_path
+
+    def _lazy_init_pipeline(self):
+        if self.pipeline is None:
+            logger.info("Loading Delighting Model...")
+            pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+                self.model_path,
+                torch_dtype=torch.float16,
+                safety_checker=None,
+            )
+            pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+                pipeline.scheduler.config
+            )
+            pipeline.set_progress_bar_config(disable=True)
+
+            pipeline.to(self.device, torch.float16)
+            self.pipeline = pipeline
+
+    def recenter_image(
+        self, image: Image.Image, border_ratio: float = 0.2
+    ) -> Image.Image:
+        if image.mode == "RGB":
+            return image
+        elif image.mode == "L":
+            image = image.convert("RGB")
+            return image
+
+        alpha_channel = np.array(image)[:, :, 3]
+        non_zero_indices = np.argwhere(alpha_channel > 0)
+        if non_zero_indices.size == 0:
+            raise ValueError("Image is fully transparent")
+
+        min_row, min_col = non_zero_indices.min(axis=0)
+        max_row, max_col = non_zero_indices.max(axis=0)
+
+        cropped_image = image.crop(
+            (min_col, min_row, max_col + 1, max_row + 1)
+        )
+
+        width, height = cropped_image.size
+        border_width = int(width * border_ratio)
+        border_height = int(height * border_ratio)
+
+        new_width = width + 2 * border_width
+        new_height = height + 2 * border_height
+
+        square_size = max(new_width, new_height)
+
+        new_image = Image.new(
+            "RGBA", (square_size, square_size), (255, 255, 255, 0)
+        )
+
+        paste_x = (square_size - new_width) // 2 + border_width
+        paste_y = (square_size - new_height) // 2 + border_height
+
+        new_image.paste(cropped_image, (paste_x, paste_y))
+
+        return new_image
+
+    @spaces.GPU
+    @torch.no_grad()
+    def __call__(
+        self,
+        image: Union[str, np.ndarray, Image.Image],
+        preprocess: bool = False,
+        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):
+            image = Image.fromarray(image)
+
+        if preprocess:
+            bg_remover = RembgRemover()
+            image = bg_remover(image)
+            image = self.recenter_image(image)
+
+        if target_wh is not None:
+            image = image.resize(target_wh)
+        else:
+            target_wh = image.size
+
+        image_array = np.array(image)
+        assert image_array.shape[-1] == 4, "Image must have alpha channel"
+
+        raw_alpha_channel = image_array[:, :, 3]
+        alpha_channel = cv2.erode(raw_alpha_channel, self.kernel, iterations=1)
+        image_array[alpha_channel == 0, :3] = 255  # must be white background
+        image_array[:, :, 3] = alpha_channel
+
+        image = self.pipeline(
+            prompt="",
+            image=Image.fromarray(image_array).convert("RGB"),
+            generator=torch.manual_seed(self.seed),
+            num_inference_steps=self.num_infer_step,
+            image_guidance_scale=self.image_guide_scale,
+            guidance_scale=self.text_guide_scale,
+        ).images[0]
+
+        alpha_channel = Image.fromarray(alpha_channel)
+        rgba_image = image.convert("RGBA").resize(target_wh)
+        rgba_image.putalpha(alpha_channel)
+
+        return rgba_image
+
+
+if __name__ == "__main__":
+    delighting_model = DelightingModel()
+    image_path = "apps/assets/example_image/sample_12.jpg"
+    image = delighting_model(
+        image_path, preprocess=True, target_wh=(512, 512)
+    )  # noqa
+    image.save("delight.png")
+
+    # image_path = "embodied_gen/scripts/test_robot.png"
+    # image = delighting_model(image_path)
+    # image.save("delighting_image_a2.png")

embodied_gen/models/gs_model.py

@@ -0,0 +1,511 @@
+# 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.
+
+
+import logging
+import os
+import struct
+from dataclasses import dataclass
+from typing import Optional
+
+import cv2
+import numpy as np
+import torch
+from gsplat.cuda._wrapper import spherical_harmonics
+from gsplat.rendering import rasterization
+from plyfile import PlyData
+from scipy.spatial.transform import Rotation
+from embodied_gen.data.utils import gamma_shs, quat_mult, quat_to_rotmat
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "RenderResult",
+    "GaussianOperator",
+]
+
+SH_C0 = 0.2820947917738781
+
+
+@dataclass
+class RenderResult:
+    rgb: np.ndarray
+    depth: np.ndarray
+    opacity: np.ndarray
+    mask_threshold: float = 10
+    mask: Optional[np.ndarray] = None
+    rgba: Optional[np.ndarray] = None
+
+    def __post_init__(self):
+        if isinstance(self.rgb, torch.Tensor):
+            rgb = (self.rgb * 255).to(torch.uint8)
+            self.rgb = rgb.cpu().numpy()[..., ::-1]
+        if isinstance(self.depth, torch.Tensor):
+            self.depth = self.depth.cpu().numpy()
+        if isinstance(self.opacity, torch.Tensor):
+            opacity = (self.opacity * 255).to(torch.uint8)
+            self.opacity = opacity.cpu().numpy()
+            mask = np.where(self.opacity > self.mask_threshold, 255, 0)
+            self.mask = mask.astype(np.uint8)
+            self.rgba = np.concatenate([self.rgb, self.mask], axis=-1)
+
+
+@dataclass
+class GaussianBase:
+    _opacities: torch.Tensor
+    _means: torch.Tensor
+    _scales: torch.Tensor
+    _quats: torch.Tensor
+    _rgbs: Optional[torch.Tensor] = None
+    _features_dc: Optional[torch.Tensor] = None
+    _features_rest: Optional[torch.Tensor] = None
+    sh_degree: Optional[int] = 0
+    device: str = "cuda"
+
+    def __post_init__(self):
+        self.active_sh_degree: int = self.sh_degree
+        self.to(self.device)
+
+    def to(self, device: str) -> None:
+        for k, v in self.__dict__.items():
+            if not isinstance(v, torch.Tensor):
+                continue
+            self.__dict__[k] = v.to(device)
+
+    def get_numpy_data(self):
+        data = {}
+        for k, v in self.__dict__.items():
+            if not isinstance(v, torch.Tensor):
+                continue
+            data[k] = v.detach().cpu().numpy()
+
+        return data
+
+    def quat_norm(self, x: torch.Tensor) -> torch.Tensor:
+        return x / x.norm(dim=-1, keepdim=True)
+
+    @classmethod
+    def load_from_ply(
+        cls,
+        path: str,
+        gamma: float = 1.0,
+        device: str = "cuda",
+    ) -> "GaussianBase":
+        plydata = PlyData.read(path)
+        xyz = torch.stack(
+            (
+                torch.tensor(plydata.elements[0]["x"], dtype=torch.float32),
+                torch.tensor(plydata.elements[0]["y"], dtype=torch.float32),
+                torch.tensor(plydata.elements[0]["z"], dtype=torch.float32),
+            ),
+            dim=1,
+        )
+
+        opacities = torch.tensor(
+            plydata.elements[0]["opacity"], dtype=torch.float32
+        ).unsqueeze(-1)
+        features_dc = torch.zeros((xyz.shape[0], 3), dtype=torch.float32)
+        features_dc[:, 0] = torch.tensor(
+            plydata.elements[0]["f_dc_0"], dtype=torch.float32
+        )
+        features_dc[:, 1] = torch.tensor(
+            plydata.elements[0]["f_dc_1"], dtype=torch.float32
+        )
+        features_dc[:, 2] = torch.tensor(
+            plydata.elements[0]["f_dc_2"], dtype=torch.float32
+        )
+
+        scale_names = [
+            p.name
+            for p in plydata.elements[0].properties
+            if p.name.startswith("scale_")
+        ]
+        scale_names = sorted(scale_names, key=lambda x: int(x.split("_")[-1]))
+        scales = torch.zeros(
+            (xyz.shape[0], len(scale_names)), dtype=torch.float32
+        )
+        for idx, attr_name in enumerate(scale_names):
+            scales[:, idx] = torch.tensor(
+                plydata.elements[0][attr_name], dtype=torch.float32
+            )
+
+        rot_names = [
+            p.name
+            for p in plydata.elements[0].properties
+            if p.name.startswith("rot_")
+        ]
+        rot_names = sorted(rot_names, key=lambda x: int(x.split("_")[-1]))
+        rots = torch.zeros((xyz.shape[0], len(rot_names)), dtype=torch.float32)
+        for idx, attr_name in enumerate(rot_names):
+            rots[:, idx] = torch.tensor(
+                plydata.elements[0][attr_name], dtype=torch.float32
+            )
+
+        rots = rots / torch.norm(rots, dim=-1, keepdim=True)
+
+        # extra features
+        extra_f_names = [
+            p.name
+            for p in plydata.elements[0].properties
+            if p.name.startswith("f_rest_")
+        ]
+        extra_f_names = sorted(
+            extra_f_names, key=lambda x: int(x.split("_")[-1])
+        )
+
+        max_sh_degree = int(np.sqrt((len(extra_f_names) + 3) / 3) - 1)
+        if max_sh_degree != 0:
+            features_extra = torch.zeros(
+                (xyz.shape[0], len(extra_f_names)), dtype=torch.float32
+            )
+            for idx, attr_name in enumerate(extra_f_names):
+                features_extra[:, idx] = torch.tensor(
+                    plydata.elements[0][attr_name], dtype=torch.float32
+                )
+
+            features_extra = features_extra.view(
+                (features_extra.shape[0], 3, (max_sh_degree + 1) ** 2 - 1)
+            )
+            features_extra = features_extra.permute(0, 2, 1)
+
+            if abs(gamma - 1.0) > 1e-3:
+                features_dc = gamma_shs(features_dc, gamma)
+                features_extra[..., :] = 0.0
+                opacities *= 0.8
+
+            shs = torch.cat(
+                [
+                    features_dc.reshape(-1, 3),
+                    features_extra.reshape(len(features_dc), -1),
+                ],
+                dim=-1,
+            )
+        else:
+            # sh_dim is 0, only dc features
+            shs = features_dc
+            features_extra = None
+
+        return cls(
+            sh_degree=max_sh_degree,
+            _means=xyz,
+            _opacities=opacities,
+            _rgbs=shs,
+            _scales=scales,
+            _quats=rots,
+            _features_dc=features_dc,
+            _features_rest=features_extra,
+            device=device,
+        )
+
+    def save_to_ply(self, path: str, enable_mask: bool = False) -> None:
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        numpy_data = self.get_numpy_data()
+        means = numpy_data["_means"]
+        scales = numpy_data["_scales"]
+        quats = numpy_data["_quats"]
+        opacities = numpy_data["_opacities"]
+        sh0 = numpy_data["_features_dc"]
+        shN = numpy_data.get("_features_rest", np.zeros((means.shape[0], 0)))
+        shN = shN.reshape(means.shape[0], -1)
+
+        # Create a mask to identify rows with NaN or Inf in any of the numpy_data arrays  # noqa
+        if enable_mask:
+            invalid_mask = (
+                np.isnan(means).any(axis=1)
+                | np.isinf(means).any(axis=1)
+                | np.isnan(scales).any(axis=1)
+                | np.isinf(scales).any(axis=1)
+                | np.isnan(quats).any(axis=1)
+                | np.isinf(quats).any(axis=1)
+                | np.isnan(opacities).any(axis=0)
+                | np.isinf(opacities).any(axis=0)
+                | np.isnan(sh0).any(axis=1)
+                | np.isinf(sh0).any(axis=1)
+                | np.isnan(shN).any(axis=1)
+                | np.isinf(shN).any(axis=1)
+            )
+
+            # Filter out rows with NaNs or Infs from all data arrays
+            means = means[~invalid_mask]
+            scales = scales[~invalid_mask]
+            quats = quats[~invalid_mask]
+            opacities = opacities[~invalid_mask]
+            sh0 = sh0[~invalid_mask]
+            shN = shN[~invalid_mask]
+
+        num_points = means.shape[0]
+        with open(path, "wb") as f:
+            # Write PLY header
+            f.write(b"ply\n")
+            f.write(b"format binary_little_endian 1.0\n")
+            f.write(f"element vertex {num_points}\n".encode())
+            f.write(b"property float x\n")
+            f.write(b"property float y\n")
+            f.write(b"property float z\n")
+
+            for i, data in enumerate([sh0, shN]):
+                prefix = "f_dc" if i == 0 else "f_rest"
+                for j in range(data.shape[1]):
+                    f.write(f"property float {prefix}_{j}\n".encode())
+
+            f.write(b"property float opacity\n")
+
+            for i in range(scales.shape[1]):
+                f.write(f"property float scale_{i}\n".encode())
+            for i in range(quats.shape[1]):
+                f.write(f"property float rot_{i}\n".encode())
+
+            f.write(b"end_header\n")
+
+            # Write vertex data
+            for i in range(num_points):
+                f.write(struct.pack("<fff", *means[i]))  # x, y, z
+
+                for data in [sh0, shN]:
+                    for j in range(data.shape[1]):
+                        f.write(struct.pack("<f", data[i, j]))
+
+                f.write(struct.pack("<f", opacities[i].item()))  # opacity
+
+                for data in [scales, quats]:
+                    for j in range(data.shape[1]):
+                        f.write(struct.pack("<f", data[i, j]))
+
+        return
+
+
+@dataclass
+class GaussianOperator(GaussianBase):
+    """Gaussian Splatting operator.
+
+    Supports transformation, scaling, color computation, and
+    rasterization-based rendering.
+
+    Inherits:
+        GaussianBase: Base class with Gaussian params (means, scales, etc.)
+
+    Functionality includes:
+    - Applying instance poses to transform Gaussian means and quaternions.
+    - Scaling Gaussians to a real-world size.
+    - Computing colors using spherical harmonics.
+    - Rendering images via differentiable rasterization.
+    - Exporting transformed and rescaled models to .ply format.
+    """
+
+    def _compute_transform(
+        self,
+        means: torch.Tensor,
+        quats: torch.Tensor,
+        instance_pose: torch.Tensor,
+    ):
+        """Compute the transform of the GS models.
+
+        Args:
+            means: tensor of gs means.
+            quats: tensor of gs quaternions.
+            instance_pose: instances poses in [x y z qx qy qz qw] format.
+
+        """
+        # (x y z qx qy qz qw) -> (x y z qw qx qy qz)
+        instance_pose = instance_pose[[0, 1, 2, 6, 3, 4, 5]]
+        cur_instances_quats = self.quat_norm(instance_pose[3:])
+        rot_cur = quat_to_rotmat(cur_instances_quats, mode="wxyz")
+
+        # update the means
+        num_gs = means.shape[0]
+        trans_per_pts = torch.stack([instance_pose[:3]] * num_gs, dim=0)
+        quat_per_pts = torch.stack([instance_pose[3:]] * num_gs, dim=0)
+        rot_per_pts = torch.stack([rot_cur] * num_gs, dim=0)  # (num_gs, 3, 3)
+
+        # update the means
+        cur_means = (
+            torch.bmm(rot_per_pts, means.unsqueeze(-1)).squeeze(-1)
+            + trans_per_pts
+        )
+
+        # update the quats
+        _quats = self.quat_norm(quats)
+        cur_quats = quat_mult(quat_per_pts, _quats)
+
+        return cur_means, cur_quats
+
+    def get_gaussians(
+        self,
+        c2w: torch.Tensor = None,
+        instance_pose: torch.Tensor = None,
+        apply_activate: bool = False,
+    ) -> "GaussianBase":
+        """Get Gaussian data under the given instance_pose."""
+        if c2w is None:
+            c2w = torch.eye(4).to(self.device)
+
+        if instance_pose is not None:
+            # compute the transformed gs means and quats
+            world_means, world_quats = self._compute_transform(
+                self._means, self._quats, instance_pose.float().to(self.device)
+            )
+        else:
+            world_means, world_quats = self._means, self._quats
+
+        # get colors of gaussians
+        if self._features_rest is not None:
+            colors = torch.cat(
+                (self._features_dc[:, None, :], self._features_rest), dim=1
+            )
+        else:
+            colors = self._features_dc[:, None, :]
+
+        if self.sh_degree > 0:
+            viewdirs = world_means.detach() - c2w[..., :3, 3]  # (N, 3)
+            viewdirs = viewdirs / viewdirs.norm(dim=-1, keepdim=True)
+            rgbs = spherical_harmonics(self.sh_degree, viewdirs, colors)
+            rgbs = torch.clamp(rgbs + 0.5, 0.0, 1.0)
+        else:
+            rgbs = torch.sigmoid(colors[:, 0, :])
+
+        gs_dict = dict(
+            _means=world_means,
+            _opacities=(
+                torch.sigmoid(self._opacities)
+                if apply_activate
+                else self._opacities
+            ),
+            _rgbs=rgbs,
+            _scales=(
+                torch.exp(self._scales) if apply_activate else self._scales
+            ),
+            _quats=self.quat_norm(world_quats),
+            _features_dc=self._features_dc,
+            _features_rest=self._features_rest,
+            sh_degree=self.sh_degree,
+            device=self.device,
+        )
+
+        return GaussianOperator(**gs_dict)
+
+    def rescale(self, scale: float):
+        if scale != 1.0:
+            self._means *= scale
+            self._scales += torch.log(self._scales.new_tensor(scale))
+
+    def set_scale_by_height(self, real_height: float) -> None:
+        def _ptp(tensor, dim):
+            val = tensor.max(dim=dim).values - tensor.min(dim=dim).values
+            return val.tolist()
+
+        xyz_scale = max(_ptp(self._means, dim=0))
+        self.rescale(1 / (xyz_scale + 1e-6))  # Normalize to [-0.5, 0.5]
+        raw_height = _ptp(self._means, dim=0)[1]
+        scale = real_height / raw_height
+
+        self.rescale(scale)
+
+        return
+
+    @staticmethod
+    def resave_ply(
+        in_ply: str,
+        out_ply: str,
+        real_height: float = None,
+        instance_pose: np.ndarray = None,
+        device: str = "cuda",
+    ) -> None:
+        gs_model = GaussianOperator.load_from_ply(in_ply, device=device)
+
+        if instance_pose is not None:
+            gs_model = gs_model.get_gaussians(instance_pose=instance_pose)
+
+        if real_height is not None:
+            gs_model.set_scale_by_height(real_height)
+
+        gs_model.save_to_ply(out_ply)
+
+        return
+
+    @staticmethod
+    def trans_to_quatpose(
+        rot_matrix: list[list[float]],
+        trans_matrix: list[float] = [0, 0, 0],
+    ) -> torch.Tensor:
+        if isinstance(rot_matrix, list):
+            rot_matrix = np.array(rot_matrix)
+
+        rot = Rotation.from_matrix(rot_matrix)
+        qx, qy, qz, qw = rot.as_quat()
+        instance_pose = torch.tensor([*trans_matrix, qx, qy, qz, qw])
+
+        return instance_pose
+
+    def render(
+        self,
+        c2w: torch.Tensor,
+        Ks: torch.Tensor,
+        image_width: int,
+        image_height: int,
+    ) -> RenderResult:
+        gs = self.get_gaussians(c2w, apply_activate=True)
+        renders, alphas, _ = rasterization(
+            means=gs._means,
+            quats=gs._quats,
+            scales=gs._scales,
+            opacities=gs._opacities.squeeze(),
+            colors=gs._rgbs,
+            viewmats=torch.linalg.inv(c2w)[None, ...],
+            Ks=Ks[None, ...],
+            width=image_width,
+            height=image_height,
+            packed=False,
+            absgrad=True,
+            sparse_grad=False,
+            # rasterize_mode="classic",
+            rasterize_mode="antialiased",
+            **{
+                "near_plane": 0.01,
+                "far_plane": 1000000000,
+                "radius_clip": 0.0,
+                "render_mode": "RGB+ED",
+            },
+        )
+        renders = renders[0]
+        alphas = alphas[0].squeeze(-1)
+
+        assert renders.shape[-1] == 4, f"Must render rgb, depth and alpha"
+        rendered_rgb, rendered_depth = torch.split(renders, [3, 1], dim=-1)
+
+        return RenderResult(
+            torch.clamp(rendered_rgb, min=0, max=1),
+            rendered_depth,
+            alphas[..., None],
+        )
+
+
+if __name__ == "__main__":
+    input_gs = "outputs/layouts_gens_demo/task_0000/background/gs_model.ply"
+    output_gs = "./gs_model.ply"
+    gs_model: GaussianOperator = GaussianOperator.load_from_ply(input_gs)
+
+    # 绕 x 轴旋转 180°
+    R_x = [[1, 0, 0], [0, -1, 0], [0, 0, -1]]
+    instance_pose = gs_model.trans_to_quatpose(R_x)
+    gs_model = gs_model.get_gaussians(instance_pose=instance_pose)
+
+    gs_model.rescale(2)
+
+    gs_model.set_scale_by_height(1.3)
+
+    gs_model.save_to_ply(output_gs)

embodied_gen/models/image_comm_model.py

@@ -0,0 +1,236 @@
+# 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.
+# Text-to-Image generation models from Hugging Face community.
+
+import os
+from abc import ABC, abstractmethod
+
+import torch
+from diffusers import (
+    ChromaPipeline,
+    Cosmos2TextToImagePipeline,
+    DPMSolverMultistepScheduler,
+    FluxPipeline,
+    KolorsPipeline,
+    StableDiffusion3Pipeline,
+)
+from diffusers.quantizers import PipelineQuantizationConfig
+from huggingface_hub import snapshot_download
+from PIL import Image
+from transformers import AutoModelForCausalLM, SiglipProcessor
+
+__all__ = [
+    "build_hf_image_pipeline",
+]
+
+
+class BasePipelineLoader(ABC):
+    def __init__(self, device="cuda"):
+        self.device = device
+
+    @abstractmethod
+    def load(self):
+        pass
+
+
+class BasePipelineRunner(ABC):
+    def __init__(self, pipe):
+        self.pipe = pipe
+
+    @abstractmethod
+    def run(self, prompt: str, **kwargs) -> Image.Image:
+        pass
+
+
+# ===== SD3.5-medium =====
+class SD35Loader(BasePipelineLoader):
+    def load(self):
+        pipe = StableDiffusion3Pipeline.from_pretrained(
+            "stabilityai/stable-diffusion-3.5-medium",
+            torch_dtype=torch.float16,
+        )
+        pipe = pipe.to(self.device)
+        pipe.enable_model_cpu_offload()
+        pipe.enable_xformers_memory_efficient_attention()
+        pipe.enable_attention_slicing()
+        return pipe
+
+
+class SD35Runner(BasePipelineRunner):
+    def run(self, prompt: str, **kwargs) -> Image.Image:
+        return self.pipe(prompt=prompt, **kwargs).images
+
+
+# ===== Cosmos2 =====
+class CosmosLoader(BasePipelineLoader):
+    def __init__(
+        self,
+        model_id="nvidia/Cosmos-Predict2-2B-Text2Image",
+        local_dir="weights/cosmos2",
+        device="cuda",
+    ):
+        super().__init__(device)
+        self.model_id = model_id
+        self.local_dir = local_dir
+
+    def _patch(self):
+        def patch_model(cls):
+            orig = cls.from_pretrained
+
+            def new(*args, **kwargs):
+                kwargs.setdefault("attn_implementation", "flash_attention_2")
+                kwargs.setdefault("torch_dtype", torch.bfloat16)
+                return orig(*args, **kwargs)
+
+            cls.from_pretrained = new
+
+        def patch_processor(cls):
+            orig = cls.from_pretrained
+
+            def new(*args, **kwargs):
+                kwargs.setdefault("use_fast", True)
+                return orig(*args, **kwargs)
+
+            cls.from_pretrained = new
+
+        patch_model(AutoModelForCausalLM)
+        patch_processor(SiglipProcessor)
+
+    def load(self):
+        self._patch()
+        snapshot_download(
+            repo_id=self.model_id,
+            local_dir=self.local_dir,
+            local_dir_use_symlinks=False,
+            resume_download=True,
+        )
+
+        config = PipelineQuantizationConfig(
+            quant_backend="bitsandbytes_4bit",
+            quant_kwargs={
+                "load_in_4bit": True,
+                "bnb_4bit_quant_type": "nf4",
+                "bnb_4bit_compute_dtype": torch.bfloat16,
+                "bnb_4bit_use_double_quant": True,
+            },
+            components_to_quantize=["text_encoder", "transformer", "unet"],
+        )
+
+        pipe = Cosmos2TextToImagePipeline.from_pretrained(
+            self.model_id,
+            torch_dtype=torch.bfloat16,
+            quantization_config=config,
+            use_safetensors=True,
+            safety_checker=None,
+            requires_safety_checker=False,
+        ).to(self.device)
+        return pipe
+
+
+class CosmosRunner(BasePipelineRunner):
+    def run(self, prompt: str, negative_prompt=None, **kwargs) -> Image.Image:
+        return self.pipe(
+            prompt=prompt, negative_prompt=negative_prompt, **kwargs
+        ).images
+
+
+# ===== Kolors =====
+class KolorsLoader(BasePipelineLoader):
+    def load(self):
+        pipe = KolorsPipeline.from_pretrained(
+            "Kwai-Kolors/Kolors-diffusers",
+            torch_dtype=torch.float16,
+            variant="fp16",
+        ).to(self.device)
+        pipe.enable_model_cpu_offload()
+        pipe.enable_xformers_memory_efficient_attention()
+        pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+            pipe.scheduler.config, use_karras_sigmas=True
+        )
+        return pipe
+
+
+class KolorsRunner(BasePipelineRunner):
+    def run(self, prompt: str, **kwargs) -> Image.Image:
+        return self.pipe(prompt=prompt, **kwargs).images
+
+
+# ===== Flux =====
+class FluxLoader(BasePipelineLoader):
+    def load(self):
+        os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+        pipe = FluxPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
+        )
+        pipe.enable_model_cpu_offload()
+        pipe.enable_xformers_memory_efficient_attention()
+        pipe.enable_attention_slicing()
+        return pipe.to(self.device)
+
+
+class FluxRunner(BasePipelineRunner):
+    def run(self, prompt: str, **kwargs) -> Image.Image:
+        return self.pipe(prompt=prompt, **kwargs).images
+
+
+# ===== Chroma =====
+class ChromaLoader(BasePipelineLoader):
+    def load(self):
+        return ChromaPipeline.from_pretrained(
+            "lodestones/Chroma", torch_dtype=torch.bfloat16
+        ).to(self.device)
+
+
+class ChromaRunner(BasePipelineRunner):
+    def run(self, prompt: str, negative_prompt=None, **kwargs) -> Image.Image:
+        return self.pipe(
+            prompt=prompt, negative_prompt=negative_prompt, **kwargs
+        ).images
+
+
+PIPELINE_REGISTRY = {
+    "sd35": (SD35Loader, SD35Runner),
+    "cosmos": (CosmosLoader, CosmosRunner),
+    "kolors": (KolorsLoader, KolorsRunner),
+    "flux": (FluxLoader, FluxRunner),
+    "chroma": (ChromaLoader, ChromaRunner),
+}
+
+
+def build_hf_image_pipeline(name: str, device="cuda") -> BasePipelineRunner:
+    if name not in PIPELINE_REGISTRY:
+        raise ValueError(f"Unsupported model: {name}")
+    loader_cls, runner_cls = PIPELINE_REGISTRY[name]
+    pipe = loader_cls(device=device).load()
+
+    return runner_cls(pipe)
+
+
+if __name__ == "__main__":
+    model_name = "sd35"
+    runner = build_hf_image_pipeline(model_name)
+    # NOTE: Just for pipeline testing, generation quality at low resolution is poor.
+    images = runner.run(
+        prompt="A robot holding a sign that says 'Hello'",
+        height=512,
+        width=512,
+        num_inference_steps=10,
+        guidance_scale=6,
+        num_images_per_prompt=1,
+    )
+
+    for i, img in enumerate(images):
+        img.save(f"image_{model_name}_{i}.jpg")

embodied_gen/models/layout.py

@@ -0,0 +1,510 @@
+# 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.
+
+
+import argparse
+import json
+import logging
+import os
+import re
+
+import json_repair
+from embodied_gen.utils.enum import (
+    LayoutInfo,
+    RobotItemEnum,
+    Scene3DItemEnum,
+    SpatialRelationEnum,
+)
+from embodied_gen.utils.gpt_clients import GPT_CLIENT, GPTclient
+from embodied_gen.utils.process_media import SceneTreeVisualizer
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "LayoutDesigner",
+    "LAYOUT_DISASSEMBLER",
+    "LAYOUT_GRAPHER",
+    "LAYOUT_DESCRIBER",
+]
+
+
+DISTRACTOR_NUM = 2  # Maximum number of distractor objects allowed
+LAYOUT_DISASSEMBLE_PROMPT = f"""
+    You are an intelligent 3D scene planner. Given a natural language
+    description of a robotic task, output a structured description of
+    an interactive 3D scene.
+
+    The output must include the following fields:
+    - task: A high-level task type (e.g., "single-arm pick",
+        "dual-arm grasping", "pick and place", "object sorting").
+    - {Scene3DItemEnum.ROBOT}: The name or type of robot involved. If not mentioned,
+        use {RobotItemEnum.FRANKA} as default.
+    - {Scene3DItemEnum.BACKGROUND}: The room or indoor environment where the task happens
+        (e.g., Kitchen, Bedroom, Living Room, Workshop, Office).
+    - {Scene3DItemEnum.CONTEXT}: A indoor object involved in the manipulation
+        (e.g., Table, Shelf, Desk, Bed, Cabinet).
+    - {Scene3DItemEnum.MANIPULATED_OBJS}: The main object(s) that the robot directly interacts with.
+    - {Scene3DItemEnum.DISTRACTOR_OBJS}: Other objects that naturally belong to the scene but are not part of the main task.
+
+    Constraints:
+    - The {Scene3DItemEnum.BACKGROUND} must logically match the described task.
+    - The {Scene3DItemEnum.CONTEXT} must fit within the {Scene3DItemEnum.BACKGROUND}. (e.g., a bedroom may include a table or bed, but not a workbench.)
+    - The {Scene3DItemEnum.CONTEXT} must be a concrete indoor object, such as a "table",
+        "shelf", "desk", or "bed". It must not be an abstract concept (e.g., "area", "space", "zone")
+        or structural surface (e.g., "floor", "ground"). If the input describes an interaction near
+        the floor or vague space, you must infer a plausible object like a "table", "cabinet", or "storage box" instead.
+    - {Scene3DItemEnum.MANIPULATED_OBJS} and {Scene3DItemEnum.DISTRACTOR_OBJS} objects must be plausible,
+        and semantically compatible with the {Scene3DItemEnum.CONTEXT} and {Scene3DItemEnum.BACKGROUND}.
+    - {Scene3DItemEnum.DISTRACTOR_OBJS} must not confuse or overlap with the manipulated objects.
+    - {Scene3DItemEnum.DISTRACTOR_OBJS} number limit: {DISTRACTOR_NUM} distractors maximum.
+    - All {Scene3DItemEnum.BACKGROUND} are limited to indoor environments.
+    - {Scene3DItemEnum.MANIPULATED_OBJS} and {Scene3DItemEnum.DISTRACTOR_OBJS} are rigid bodies and not include flexible objects.
+    - {Scene3DItemEnum.MANIPULATED_OBJS} and {Scene3DItemEnum.DISTRACTOR_OBJS} must be common
+        household or office items or furniture, not abstract concepts, not too small like needle.
+    - If the input includes a plural or grouped object (e.g., "pens", "bottles", "plates", "fruit"),
+        you must decompose it into multiple individual instances (e.g., ["pen1", "pen2"], ["apple", "pear"]).
+    - Containers that hold objects (e.g., "bowl of apples", "box of tools") must
+        be separated into individual items (e.g., ["bowl", "apple1", "apple2"]).
+    - Do not include transparent objects such as "glass", "plastic", etc.
+    - All {Scene3DItemEnum.MANIPULATED_OBJS} and {Scene3DItemEnum.DISTRACTOR_OBJS} must be child node of {Scene3DItemEnum.CONTEXT}.
+    - The output must be in compact JSON format and use Markdown syntax, just like the output in the example below.
+
+    Examples:
+
+    Input:
+    "Pick up the marker from the table and put it in the bowl robot {RobotItemEnum.UR5}."
+    Output:
+    ```json
+    {{
+        "task_desc": "Pick up the marker from the table and put it in the bowl.",
+        "task": "pick and place",
+        "{Scene3DItemEnum.ROBOT}": "{RobotItemEnum.UR5}",
+        "{Scene3DItemEnum.BACKGROUND}": "kitchen",
+        "{Scene3DItemEnum.CONTEXT}": "table",
+        "{Scene3DItemEnum.MANIPULATED_OBJS}": ["marker"],
+        "{Scene3DItemEnum.DISTRACTOR_OBJS}": ["mug", "notebook", "bowl"]
+    }}
+    ```
+
+    Input:
+    "Put the rubik's cube on the top of the shelf."
+    Output:
+    ```json
+    {{
+        "task_desc": "Put the rubik's cube on the top of the shelf.",
+        "task": "pick and place",
+        "{Scene3DItemEnum.ROBOT}": "{RobotItemEnum.FRANKA}",
+        "{Scene3DItemEnum.BACKGROUND}": "bedroom",
+        "{Scene3DItemEnum.CONTEXT}": "shelf",
+        "{Scene3DItemEnum.MANIPULATED_OBJS}": ["rubik's cube"],
+        "{Scene3DItemEnum.DISTRACTOR_OBJS}": ["pen", "cup", "toy car"]
+    }}
+    ```
+
+    Input:
+    "Remove all the objects from the white basket and put them on the table."
+    Output:
+    ```json
+    {{
+        "task_desc": "Remove all the objects from the white basket and put them on the table, robot {RobotItemEnum.PIPER}.",
+        "task": "pick and place",
+        "{Scene3DItemEnum.ROBOT}": "{RobotItemEnum.PIPER}",
+        "{Scene3DItemEnum.BACKGROUND}": "office",
+        "{Scene3DItemEnum.CONTEXT}": "table",
+        "{Scene3DItemEnum.MANIPULATED_OBJS}": ["banana", "mobile phone"],
+        "{Scene3DItemEnum.DISTRACTOR_OBJS}": ["plate", "white basket"]
+    }}
+    ```
+
+    Input:
+    "Pick up the rope on the chair and put it in the box."
+    Output:
+    ```json
+    {{
+        "task_desc": "Pick up the rope on the chair and put it in the box, robot {RobotItemEnum.FRANKA}.",
+        "task": "pick and place",
+        "{Scene3DItemEnum.ROBOT}": "{RobotItemEnum.FRANKA}",
+        "{Scene3DItemEnum.BACKGROUND}": "living room",
+        "{Scene3DItemEnum.CONTEXT}": "chair",
+        "{Scene3DItemEnum.MANIPULATED_OBJS}": ["rope", "box"],
+        "{Scene3DItemEnum.DISTRACTOR_OBJS}": ["magazine"]
+    }}
+    ```
+
+    Input:
+    "Pick up the seal tape and plastic from the counter and put them in the open drawer and close it."
+    Output:
+    ```json
+    {{
+        "task_desc": "Pick up the seal tape and plastic from the counter and put them in the open drawer and close it.",
+        "task": "pick and place",
+        "robot": "franka",
+        "background": "kitchen",
+        "context": "counter",
+        "manipulated_objs": ["seal tape", "plastic", "opened drawer"],
+        "distractor_objs": ["scissors"]
+    }}
+    ```
+
+    Input:
+    "Put the pens in the grey bowl."
+    Output:
+    ```json
+    {{
+        "task_desc": "Put the pens in the grey bowl.",
+        "task": "pick and place",
+        "robot": "franka",
+        "background": "office",
+        "context": "table",
+        "manipulated_objs": ["pen1", "pen2", "grey bowl"],
+        "distractor_objs": ["notepad", "cup"]
+    }}
+    ```
+
+"""
+
+
+LAYOUT_HIERARCHY_PROMPT = f"""
+    You are a 3D scene layout reasoning expert.
+    Your task is to generate a spatial relationship dictionary in multiway tree
+    that describes how objects are arranged in a 3D environment
+    based on a given task description and object list.
+
+    Input in JSON format containing the task description, task type,
+    {Scene3DItemEnum.ROBOT}, {Scene3DItemEnum.BACKGROUND}, {Scene3DItemEnum.CONTEXT},
+    and a list of objects, including {Scene3DItemEnum.MANIPULATED_OBJS} and {Scene3DItemEnum.DISTRACTOR_OBJS}.
+
+    ### Supported Spatial Relations:
+    - "{SpatialRelationEnum.ON}": The child object bottom is directly on top of the parent object top.
+    - "{SpatialRelationEnum.INSIDE}": The child object is inside the context object.
+    - "{SpatialRelationEnum.IN}": The {Scene3DItemEnum.ROBOT} in the {Scene3DItemEnum.BACKGROUND}.
+    - "{SpatialRelationEnum.FLOOR}": The child object bottom is on the floor of the {Scene3DItemEnum.BACKGROUND}.
+
+    ### Rules:
+    - The {Scene3DItemEnum.CONTEXT} object must be "{SpatialRelationEnum.FLOOR}" the {Scene3DItemEnum.BACKGROUND}.
+    - {Scene3DItemEnum.MANIPULATED_OBJS} and {Scene3DItemEnum.DISTRACTOR_OBJS} must be either
+        "{SpatialRelationEnum.ON}" or "{SpatialRelationEnum.INSIDE}" the {Scene3DItemEnum.CONTEXT}
+    - Or "{SpatialRelationEnum.FLOOR}" {Scene3DItemEnum.BACKGROUND}.
+    - Use "{SpatialRelationEnum.INSIDE}" only if the parent is a container-like object (e.g., shelf, rack, cabinet).
+    - Do not define relationship edges between objects, only for the child and parent nodes.
+    - {Scene3DItemEnum.ROBOT} must "{SpatialRelationEnum.IN}" the {Scene3DItemEnum.BACKGROUND}.
+    - Ensure that each object appears only once in the layout tree, and its spatial relationship is defined with only one parent.
+    - Ensure a valid multiway tree structure with a maximum depth of 2 levels suitable for a 3D scene layout representation.
+    - Only output the final output in JSON format, using Markdown syntax as in examples.
+
+    ### Example
+    Input:
+    {{
+        "task_desc": "Pick up the marker from the table and put it in the bowl.",
+        "task": "pick and place",
+        "{Scene3DItemEnum.ROBOT}": "{RobotItemEnum.FRANKA}",
+        "{Scene3DItemEnum.BACKGROUND}": "kitchen",
+        "{Scene3DItemEnum.CONTEXT}": "table",
+        "{Scene3DItemEnum.MANIPULATED_OBJS}": ["marker", "bowl"],
+        "{Scene3DItemEnum.DISTRACTOR_OBJS}": ["mug", "chair"]
+    }}
+    Intermediate Think:
+        table {SpatialRelationEnum.FLOOR} kitchen
+        chair {SpatialRelationEnum.FLOOR} kitchen
+        {RobotItemEnum.FRANKA} {SpatialRelationEnum.IN} kitchen
+        marker {SpatialRelationEnum.ON} table
+        bowl {SpatialRelationEnum.ON} table
+        mug {SpatialRelationEnum.ON} table
+    Final Output:
+    ```json
+    {{
+        "kitchen": [
+            ["table", "{SpatialRelationEnum.FLOOR}"],
+            ["chair", "{SpatialRelationEnum.FLOOR}"],
+            ["{RobotItemEnum.FRANKA}", "{SpatialRelationEnum.IN}"]
+        ],
+        "table": [
+            ["marker", "{SpatialRelationEnum.ON}"],
+            ["bowl", "{SpatialRelationEnum.ON}"],
+            ["mug", "{SpatialRelationEnum.ON}"]
+        ]
+    }}
+    ```
+
+    Input:
+    {{
+        "task_desc": "Put the marker on top of the book.",
+        "task": "pick and place",
+        "{Scene3DItemEnum.ROBOT}": "{RobotItemEnum.UR5}",
+        "{Scene3DItemEnum.BACKGROUND}": "office",
+        "{Scene3DItemEnum.CONTEXT}": "desk",
+        "{Scene3DItemEnum.MANIPULATED_OBJS}": ["marker", "book"],
+        "{Scene3DItemEnum.DISTRACTOR_OBJS}": ["pen holder", "notepad"]
+    }}
+    Intermediate Think:
+        desk {SpatialRelationEnum.FLOOR} office
+        {RobotItemEnum.UR5} {SpatialRelationEnum.IN} office
+        marker {SpatialRelationEnum.ON} desk
+        book {SpatialRelationEnum.ON} desk
+        pen holder {SpatialRelationEnum.ON} desk
+        notepad {SpatialRelationEnum.ON} desk
+    Final Output:
+    ```json
+    {{
+        "office": [
+            ["desk", "{SpatialRelationEnum.FLOOR}"],
+            ["{RobotItemEnum.UR5}", "{SpatialRelationEnum.IN}"]
+        ],
+        "desk": [
+            ["marker", "{SpatialRelationEnum.ON}"],
+            ["book", "{SpatialRelationEnum.ON}"],
+            ["pen holder", "{SpatialRelationEnum.ON}"],
+            ["notepad", "{SpatialRelationEnum.ON}"]
+        ]
+    }}
+    ```
+
+    Input:
+    {{
+        "task_desc": "Put the rubik's cube on the top of the shelf.",
+        "task": "pick and place",
+        "{Scene3DItemEnum.ROBOT}": "{RobotItemEnum.UR5}",
+        "{Scene3DItemEnum.BACKGROUND}": "bedroom",
+        "{Scene3DItemEnum.CONTEXT}": "shelf",
+        "{Scene3DItemEnum.MANIPULATED_OBJS}": ["rubik's cube"],
+        "{Scene3DItemEnum.DISTRACTOR_OBJS}": ["toy car", "pen"]
+    }}
+    Intermediate Think:
+        shelf {SpatialRelationEnum.FLOOR} bedroom
+        {RobotItemEnum.UR5} {SpatialRelationEnum.IN} bedroom
+        rubik's cube {SpatialRelationEnum.INSIDE} shelf
+        toy car {SpatialRelationEnum.INSIDE} shelf
+        pen {SpatialRelationEnum.INSIDE} shelf
+    Final Output:
+    ```json
+    {{
+        "bedroom": [
+            ["shelf", "{SpatialRelationEnum.FLOOR}"],
+            ["{RobotItemEnum.UR5}", "{SpatialRelationEnum.IN}"]
+        ],
+        "shelf": [
+            ["rubik's cube", "{SpatialRelationEnum.INSIDE}"],
+            ["toy car", "{SpatialRelationEnum.INSIDE}"],
+            ["pen", "{SpatialRelationEnum.INSIDE}"]
+        ]
+    }}
+    ```
+
+    Input:
+    {{
+        "task_desc": "Put the marker in the cup on the counter.",
+        "task": "pick and place",
+        "robot": "franka",
+        "background": "kitchen",
+        "context": "counter",
+        "manipulated_objs": ["marker", "cup"],
+        "distractor_objs": ["plate", "spoon"]
+    }}
+    Intermediate Think:
+        counter {SpatialRelationEnum.FLOOR} kitchen
+        {RobotItemEnum.FRANKA} {SpatialRelationEnum.IN} kitchen
+        marker {SpatialRelationEnum.ON} counter
+        cup {SpatialRelationEnum.ON} counter
+        plate {SpatialRelationEnum.ON} counter
+        spoon {SpatialRelationEnum.ON} counter
+    Final Output:
+    ```json
+    {{
+        "kitchen": [
+            ["counter", "{SpatialRelationEnum.FLOOR}"],
+            ["{RobotItemEnum.FRANKA}", "{SpatialRelationEnum.IN}"]
+        ],
+        "counter": [
+            ["marker", "{SpatialRelationEnum.ON}"],
+            ["cup", "{SpatialRelationEnum.ON}"],
+            ["plate", "{SpatialRelationEnum.ON}"],
+            ["spoon", "{SpatialRelationEnum.ON}"]
+        ]
+    }}
+    ```
+"""
+
+
+LAYOUT_DESCRIBER_PROMPT = """
+    You are a 3D asset style descriptor.
+
+    Given a task description and a dictionary where the key is the object content and
+    the value is the object type, output a JSON dictionary with each object paired
+    with a concise, styled visual description suitable for 3D asset generation.
+
+    Generation Guidelines:
+    - For each object, brainstorm multiple style candidates before selecting the final
+        description. Vary phrasing, material, texture, color, and spatial details.
+    - Each description must be a maximum of 15 words, including color, style, materials.
+    - Descriptions should be visually grounded, specific, and reflect surface texture and structure.
+    - For objects marked as "context", explicitly mention the object is standalone, has an empty top.
+    - Use rich style descriptors: e.g., "scratched brown wooden desk" etc.
+    - Ensure all object styles align with the task's overall context and environment.
+
+    Format your output in JSON like the example below.
+
+    Example Input:
+    "Pick up the rope on the chair and put it in the box. {'living room': 'background', 'chair': 'context',
+        'rope': 'manipulated_objs', 'box': 'manipulated_objs', 'magazine': 'distractor_objs'}"
+
+    Example Output:
+    ```json
+    {
+        "living room": "modern cozy living room with soft sunlight and light grey carpet",
+        "chair": "standalone dark oak chair with no surroundings and clean empty seat",
+        "rope": "twisted hemp rope with rough fibers and dusty beige texture",
+        "box": "slightly crumpled cardboard box with open flaps and brown textured surface",
+        "magazine": "celebrity magazine with glossy red cover and large bold title"
+    }
+    ```
+"""
+
+
+class LayoutDesigner(object):
+    def __init__(
+        self,
+        gpt_client: GPTclient,
+        system_prompt: str,
+        verbose: bool = False,
+    ) -> None:
+        self.prompt = system_prompt.strip()
+        self.verbose = verbose
+        self.gpt_client = gpt_client
+
+    def query(self, prompt: str, params: dict = None) -> str:
+        full_prompt = self.prompt + f"\n\nInput:\n\"{prompt}\""
+
+        response = self.gpt_client.query(
+            text_prompt=full_prompt,
+            params=params,
+        )
+
+        if self.verbose:
+            logger.info(f"Response: {response}")
+
+        return response
+
+    def format_response(self, response: str) -> dict:
+        cleaned = re.sub(r"^```json\s*|\s*```$", "", response.strip())
+        try:
+            output = json.loads(cleaned)
+        except json.JSONDecodeError as e:
+            raise json.JSONDecodeError(
+                f"Error: {e}, failed to parse JSON response: {response}"
+            )
+
+        return output
+
+    def format_response_repair(self, response: str) -> dict:
+        return json_repair.loads(response)
+
+    def save_output(self, output: dict, save_path: str) -> None:
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        with open(save_path, 'w') as f:
+            json.dump(output, f, indent=4)
+
+    def __call__(
+        self, prompt: str, save_path: str = None, params: dict = None
+    ) -> dict | str:
+        response = self.query(prompt, params=params)
+        output = self.format_response_repair(response)
+        self.save_output(output, save_path) if save_path else None
+
+        return output
+
+
+LAYOUT_DISASSEMBLER = LayoutDesigner(
+    gpt_client=GPT_CLIENT, system_prompt=LAYOUT_DISASSEMBLE_PROMPT
+)
+LAYOUT_GRAPHER = LayoutDesigner(
+    gpt_client=GPT_CLIENT, system_prompt=LAYOUT_HIERARCHY_PROMPT
+)
+LAYOUT_DESCRIBER = LayoutDesigner(
+    gpt_client=GPT_CLIENT, system_prompt=LAYOUT_DESCRIBER_PROMPT
+)
+
+
+def build_scene_layout(
+    task_desc: str, output_path: str = None, gpt_params: dict = None
+) -> LayoutInfo:
+    layout_relation = LAYOUT_DISASSEMBLER(task_desc, params=gpt_params)
+    layout_tree = LAYOUT_GRAPHER(layout_relation, params=gpt_params)
+    object_mapping = Scene3DItemEnum.object_mapping(layout_relation)
+    obj_prompt = f'{layout_relation["task_desc"]} {object_mapping}'
+    objs_desc = LAYOUT_DESCRIBER(obj_prompt, params=gpt_params)
+    layout_info = LayoutInfo(
+        layout_tree, layout_relation, objs_desc, object_mapping
+    )
+
+    if output_path is not None:
+        visualizer = SceneTreeVisualizer(layout_info)
+        visualizer.render(save_path=output_path)
+        logger.info(f"Scene hierarchy tree saved to {output_path}")
+
+    return layout_info
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="3D Scene Layout Designer")
+    parser.add_argument(
+        "--task_desc",
+        type=str,
+        default="Put the apples on the table on the plate",
+        help="Natural language description of the robotic task",
+    )
+    parser.add_argument(
+        "--save_root",
+        type=str,
+        default="outputs/layout_tree",
+        help="Path to save the layout output",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    from embodied_gen.utils.enum import LayoutInfo
+    from embodied_gen.utils.process_media import SceneTreeVisualizer
+
+    args = parse_args()
+    params = {
+        "temperature": 1.0,
+        "top_p": 0.95,
+        "frequency_penalty": 0.3,
+        "presence_penalty": 0.5,
+    }
+    layout_relation = LAYOUT_DISASSEMBLER(args.task_desc, params=params)
+    layout_tree = LAYOUT_GRAPHER(layout_relation, params=params)
+
+    object_mapping = Scene3DItemEnum.object_mapping(layout_relation)
+    obj_prompt = f'{layout_relation["task_desc"]} {object_mapping}'
+
+    objs_desc = LAYOUT_DESCRIBER(obj_prompt, params=params)
+
+    layout_info = LayoutInfo(layout_tree, layout_relation, objs_desc)
+
+    visualizer = SceneTreeVisualizer(layout_info)
+    os.makedirs(args.save_root, exist_ok=True)
+    scene_graph_path = f"{args.save_root}/scene_tree.jpg"
+    visualizer.render(save_path=scene_graph_path)
+    with open(f"{args.save_root}/layout.json", "w") as f:
+        json.dump(layout_info.to_dict(), f, indent=4)
+
+    print(f"Scene hierarchy tree saved to {scene_graph_path}")
+    print(f"Disassembled Layout: {layout_relation}")
+    print(f"Layout Graph: {layout_tree}")
+    print(f"Layout Descriptions: {objs_desc}")

embodied_gen/models/segment_model.py

@@ -0,0 +1,379 @@
+# 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.
+
+
+import logging
+import os
+from typing import Literal, Union
+
+import cv2
+import numpy as np
+import rembg
+import torch
+from huggingface_hub import snapshot_download
+from PIL import Image
+from segment_anything import (
+    SamAutomaticMaskGenerator,
+    SamPredictor,
+    sam_model_registry,
+)
+from transformers import pipeline
+from embodied_gen.data.utils import resize_pil, trellis_preprocess
+from embodied_gen.utils.process_media import filter_small_connected_components
+from embodied_gen.validators.quality_checkers import ImageSegChecker
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "SAMRemover",
+    "SAMPredictor",
+    "RembgRemover",
+    "get_segmented_image_by_agent",
+]
+
+
+class SAMRemover(object):
+    """Loading SAM models and performing background removal on images.
+
+    Attributes:
+        checkpoint (str): Path to the model checkpoint.
+        model_type (str): Type of the SAM model to load (default: "vit_h").
+        area_ratio (float): Area ratio filtering small connected components.
+    """
+
+    def __init__(
+        self,
+        checkpoint: str = None,
+        model_type: str = "vit_h",
+        area_ratio: float = 15,
+    ):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model_type = model_type
+        self.area_ratio = area_ratio
+
+        if checkpoint is None:
+            suffix = "sam"
+            model_path = snapshot_download(
+                repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
+            )
+            checkpoint = os.path.join(
+                model_path, suffix, "sam_vit_h_4b8939.pth"
+            )
+
+        self.mask_generator = self._load_sam_model(checkpoint)
+
+    def _load_sam_model(self, checkpoint: str) -> SamAutomaticMaskGenerator:
+        sam = sam_model_registry[self.model_type](checkpoint=checkpoint)
+        sam.to(device=self.device)
+
+        return SamAutomaticMaskGenerator(sam)
+
+    def __call__(
+        self, image: Union[str, Image.Image, np.ndarray], save_path: str = None
+    ) -> Image.Image:
+        """Removes the background from an image using the SAM model.
+
+        Args:
+            image (Union[str, Image.Image, np.ndarray]): Input image,
+                can be a file path, PIL Image, or numpy array.
+            save_path (str): Path to save the output image (default: None).
+
+        Returns:
+            Image.Image: The image with background removed,
+                including an alpha channel.
+        """
+        # Convert input to numpy array
+        if isinstance(image, str):
+            image = Image.open(image)
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image).convert("RGB")
+        image = resize_pil(image)
+        image = np.array(image.convert("RGB"))
+
+        # Generate masks
+        masks = self.mask_generator.generate(image)
+        masks = sorted(masks, key=lambda x: x["area"], reverse=True)
+
+        if not masks:
+            logger.warning(
+                "Segmentation failed: No mask generated, return raw image."
+            )
+            output_image = Image.fromarray(image, mode="RGB")
+        else:
+            # Use the largest mask
+            best_mask = masks[0]["segmentation"]
+            mask = (best_mask * 255).astype(np.uint8)
+            mask = filter_small_connected_components(
+                mask, area_ratio=self.area_ratio
+            )
+            # Apply the mask to remove the background
+            background_removed = cv2.bitwise_and(image, image, mask=mask)
+            output_image = np.dstack((background_removed, mask))
+            output_image = Image.fromarray(output_image, mode="RGBA")
+
+        if save_path is not None:
+            os.makedirs(os.path.dirname(save_path), exist_ok=True)
+            output_image.save(save_path)
+
+        return output_image
+
+
+class SAMPredictor(object):
+    def __init__(
+        self,
+        checkpoint: str = None,
+        model_type: str = "vit_h",
+        binary_thresh: float = 0.1,
+        device: str = "cuda",
+    ):
+        self.device = device
+        self.model_type = model_type
+
+        if checkpoint is None:
+            suffix = "sam"
+            model_path = snapshot_download(
+                repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
+            )
+            checkpoint = os.path.join(
+                model_path, suffix, "sam_vit_h_4b8939.pth"
+            )
+
+        self.predictor = self._load_sam_model(checkpoint)
+        self.binary_thresh = binary_thresh
+
+    def _load_sam_model(self, checkpoint: str) -> SamPredictor:
+        sam = sam_model_registry[self.model_type](checkpoint=checkpoint)
+        sam.to(device=self.device)
+
+        return SamPredictor(sam)
+
+    def preprocess_image(self, image: Image.Image) -> np.ndarray:
+        if isinstance(image, str):
+            image = Image.open(image)
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image).convert("RGB")
+
+        image = resize_pil(image)
+        image = np.array(image.convert("RGB"))
+
+        return image
+
+    def generate_masks(
+        self,
+        image: np.ndarray,
+        selected_points: list[list[int]],
+    ) -> np.ndarray:
+        if len(selected_points) == 0:
+            return []
+
+        points = (
+            torch.Tensor([p for p, _ in selected_points])
+            .to(self.predictor.device)
+            .unsqueeze(1)
+        )
+
+        labels = (
+            torch.Tensor([int(l) for _, l in selected_points])
+            .to(self.predictor.device)
+            .unsqueeze(1)
+        )
+
+        transformed_points = self.predictor.transform.apply_coords_torch(
+            points, image.shape[:2]
+        )
+
+        masks, scores, _ = self.predictor.predict_torch(
+            point_coords=transformed_points,
+            point_labels=labels,
+            multimask_output=True,
+        )
+        valid_mask = masks[:, torch.argmax(scores, dim=1)]
+        masks_pos = valid_mask[labels[:, 0] == 1, 0].cpu().detach().numpy()
+        masks_neg = valid_mask[labels[:, 0] == 0, 0].cpu().detach().numpy()
+        if len(masks_neg) == 0:
+            masks_neg = np.zeros_like(masks_pos)
+        if len(masks_pos) == 0:
+            masks_pos = np.zeros_like(masks_neg)
+        masks_neg = masks_neg.max(axis=0, keepdims=True)
+        masks_pos = masks_pos.max(axis=0, keepdims=True)
+        valid_mask = (masks_pos.astype(int) - masks_neg.astype(int)).clip(0, 1)
+
+        binary_mask = (valid_mask > self.binary_thresh).astype(np.int32)
+
+        return [(mask, f"mask_{i}") for i, mask in enumerate(binary_mask)]
+
+    def get_segmented_image(
+        self, image: np.ndarray, masks: list[tuple[np.ndarray, str]]
+    ) -> Image.Image:
+        seg_image = Image.fromarray(image, mode="RGB")
+        alpha_channel = np.zeros(
+            (seg_image.height, seg_image.width), dtype=np.uint8
+        )
+        for mask, _ in masks:
+            # Use the maximum to combine multiple masks
+            alpha_channel = np.maximum(alpha_channel, mask)
+
+        alpha_channel = np.clip(alpha_channel, 0, 1)
+        alpha_channel = (alpha_channel * 255).astype(np.uint8)
+        alpha_image = Image.fromarray(alpha_channel, mode="L")
+        r, g, b = seg_image.split()
+        seg_image = Image.merge("RGBA", (r, g, b, alpha_image))
+
+        return seg_image
+
+    def __call__(
+        self,
+        image: Union[str, Image.Image, np.ndarray],
+        selected_points: list[list[int]],
+    ) -> Image.Image:
+        image = self.preprocess_image(image)
+        self.predictor.set_image(image)
+        masks = self.generate_masks(image, selected_points)
+
+        return self.get_segmented_image(image, masks)
+
+
+class RembgRemover(object):
+    def __init__(self):
+        self.rembg_session = rembg.new_session("u2net")
+
+    def __call__(
+        self, image: Union[str, Image.Image, np.ndarray], save_path: str = None
+    ) -> Image.Image:
+        if isinstance(image, str):
+            image = Image.open(image)
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+
+        image = resize_pil(image)
+        output_image = rembg.remove(image, session=self.rembg_session)
+
+        if save_path is not None:
+            os.makedirs(os.path.dirname(save_path), exist_ok=True)
+            output_image.save(save_path)
+
+        return output_image
+
+
+class BMGG14Remover(object):
+    def __init__(self) -> None:
+        self.model = pipeline(
+            "image-segmentation",
+            model="briaai/RMBG-1.4",
+            trust_remote_code=True,
+        )
+
+    def __call__(
+        self, image: Union[str, Image.Image, np.ndarray], save_path: str = None
+    ):
+        if isinstance(image, str):
+            image = Image.open(image)
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+
+        image = resize_pil(image)
+        output_image = self.model(image)
+
+        if save_path is not None:
+            os.makedirs(os.path.dirname(save_path), exist_ok=True)
+            output_image.save(save_path)
+
+        return output_image
+
+
+def invert_rgba_pil(
+    image: Image.Image, mask: Image.Image, save_path: str = None
+) -> Image.Image:
+    mask = (255 - np.array(mask))[..., None]
+    image_array = np.concatenate([np.array(image), mask], axis=-1)
+    inverted_image = Image.fromarray(image_array, "RGBA")
+
+    if save_path is not None:
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        inverted_image.save(save_path)
+
+    return inverted_image
+
+
+def get_segmented_image_by_agent(
+    image: Image.Image,
+    sam_remover: SAMRemover,
+    rbg_remover: RembgRemover,
+    seg_checker: ImageSegChecker = None,
+    save_path: str = None,
+    mode: Literal["loose", "strict"] = "loose",
+) -> Image.Image:
+    def _is_valid_seg(raw_img: Image.Image, seg_img: Image.Image) -> bool:
+        if seg_checker is None:
+            return True
+        return raw_img.mode == "RGBA" and seg_checker([raw_img, seg_img])[0]
+
+    out_sam = f"{save_path}_sam.png" if save_path else None
+    out_sam_inv = f"{save_path}_sam_inv.png" if save_path else None
+    out_rbg = f"{save_path}_rbg.png" if save_path else None
+
+    seg_image = sam_remover(image, out_sam)
+    seg_image = seg_image.convert("RGBA")
+    _, _, _, alpha = seg_image.split()
+    seg_image_inv = invert_rgba_pil(image.convert("RGB"), alpha, out_sam_inv)
+    seg_image_rbg = rbg_remover(image, out_rbg)
+
+    final_image = None
+    if _is_valid_seg(image, seg_image):
+        final_image = seg_image
+    elif _is_valid_seg(image, seg_image_inv):
+        final_image = seg_image_inv
+    elif _is_valid_seg(image, seg_image_rbg):
+        logger.warning(f"Failed to segment by `SAM`, retry with `rembg`.")
+        final_image = seg_image_rbg
+    else:
+        if mode == "strict":
+            raise RuntimeError(
+                f"Failed to segment by `SAM` or `rembg`, abort."
+            )
+        logger.warning("Failed to segment by SAM or rembg, use raw image.")
+        final_image = image.convert("RGBA")
+
+    if save_path:
+        final_image.save(save_path)
+
+    final_image = trellis_preprocess(final_image)
+
+    return final_image
+
+
+if __name__ == "__main__":
+    input_image = "outputs/text2image/demo_objects/electrical/sample_0.jpg"
+    output_image = "sample_0_seg2.png"
+
+    # input_image = "outputs/text2image/tmp/coffee_machine.jpeg"
+    # output_image = "outputs/text2image/tmp/coffee_machine_seg.png"
+
+    # 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 = BMGG14Remover()
+    remover("embodied_gen/models/test_seg.jpg", "./seg.png")

embodied_gen/models/sr_model.py

@@ -0,0 +1,174 @@
+# 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.
+
+
+import logging
+import os
+from typing import Union
+
+import numpy as np
+import spaces
+import torch
+from huggingface_hub import snapshot_download
+from PIL import Image
+from embodied_gen.data.utils import get_images_from_grid
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "ImageStableSR",
+    "ImageRealESRGAN",
+]
+
+
+class ImageStableSR:
+    """Super-resolution image upscaler using Stable Diffusion x4 upscaling model from StabilityAI."""
+
+    def __init__(
+        self,
+        model_path: str = "stabilityai/stable-diffusion-x4-upscaler",
+        device="cuda",
+    ) -> None:
+        from diffusers import StableDiffusionUpscalePipeline
+
+        self.up_pipeline_x4 = StableDiffusionUpscalePipeline.from_pretrained(
+            model_path,
+            torch_dtype=torch.float16,
+        ).to(device)
+        self.up_pipeline_x4.set_progress_bar_config(disable=True)
+        # self.up_pipeline_x4.enable_model_cpu_offload()
+
+    @spaces.GPU
+    def __call__(
+        self,
+        image: Union[Image.Image, np.ndarray],
+        prompt: str = "",
+        infer_step: int = 20,
+    ) -> Image.Image:
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+
+        image = image.convert("RGB")
+
+        with torch.no_grad():
+            upscaled_image = self.up_pipeline_x4(
+                image=image,
+                prompt=[prompt],
+                num_inference_steps=infer_step,
+            ).images[0]
+
+        return upscaled_image
+
+
+class ImageRealESRGAN:
+    """A wrapper for Real-ESRGAN-based image super-resolution.
+
+    This class uses the RealESRGAN model to perform image upscaling,
+    typically by a factor of 4.
+
+    Attributes:
+        outscale (int): The output image scale factor (e.g., 2, 4).
+        model_path (str): Path to the pre-trained model weights.
+    """
+
+    def __init__(self, outscale: int, model_path: str = None) -> None:
+        # monkey patch to support torchvision>=0.16
+        import torchvision
+        from packaging import version
+
+        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.rgb_to_grayscale = TF.rgb_to_grayscale
+            sys.modules["torchvision.transforms.functional_tensor"] = (
+                functional_tensor
+            )
+
+        self.outscale = outscale
+        self.upsampler = None
+
+        if model_path is None:
+            suffix = "super_resolution"
+            model_path = snapshot_download(
+                repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
+            )
+            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
+
+            model = RRDBNet(
+                num_in_ch=3,
+                num_out_ch=3,
+                num_feat=64,
+                num_block=23,
+                num_grow_ch=32,
+                scale=4,
+            )
+
+            self.upsampler = RealESRGANer(
+                scale=4,
+                model_path=self.model_path,
+                model=model,
+                pre_pad=0,
+                half=True,
+            )
+
+    @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)
+
+        with torch.no_grad():
+            output, _ = self.upsampler.enhance(image, outscale=self.outscale)
+
+        return Image.fromarray(output)
+
+
+if __name__ == "__main__":
+    color_path = "outputs/texture_mesh_gen/multi_view/color_sample0.png"
+
+    # Use RealESRGAN_x4plus for x4 (512->2048) image super resolution.
+    super_model = ImageRealESRGAN(outscale=4)
+    multiviews = get_images_from_grid(color_path, img_size=512)
+    multiviews = [super_model(img.convert("RGB")) for img in multiviews]
+    for idx, img in enumerate(multiviews):
+        img.save(f"sr{idx}.png")
+
+    # # Use stable diffusion for x4 (512->2048) image super resolution.
+    # super_model = ImageStableSR()
+    # multiviews = get_images_from_grid(color_path, img_size=512)
+    # multiviews = [super_model(img) for img in multiviews]
+    # for idx, img in enumerate(multiviews):
+    #     img.save(f"sr_stable{idx}.png")

embodied_gen/models/text_model.py

@@ -0,0 +1,213 @@
+# 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.
+
+
+import logging
+import os
+import random
+import subprocess
+
+import numpy as np
+import torch
+from diffusers import (
+    AutoencoderKL,
+    EulerDiscreteScheduler,
+    UNet2DConditionModel,
+)
+from kolors.models.modeling_chatglm import ChatGLMModel
+from kolors.models.tokenization_chatglm import ChatGLMTokenizer
+from kolors.models.unet_2d_condition import (
+    UNet2DConditionModel as UNet2DConditionModelIP,
+)
+from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256 import (
+    StableDiffusionXLPipeline,
+)
+from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256_ipadapter import (  # noqa
+    StableDiffusionXLPipeline as StableDiffusionXLPipelineIP,
+)
+from PIL import Image
+from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "build_text2img_ip_pipeline",
+    "build_text2img_pipeline",
+    "text2img_gen",
+    "download_kolors_weights",
+]
+
+PROMPT_APPEND = (
+    "Angled 3D view of one {object}, centered, no cropping, no occlusion, isolated product photo, "
+    "no surroundings, high-quality appearance, vivid colors, on a plain clean surface, 3D style revealing multiple surfaces"
+)
+PROMPT_KAPPEND = "Single {object}, in the center of the image, white background, 3D style, best quality"
+
+
+def download_kolors_weights(local_dir: str = "weights/Kolors") -> None:
+    logger.info(f"Download kolors weights from huggingface...")
+    os.makedirs(local_dir, exist_ok=True)
+    subprocess.run(
+        [
+            "huggingface-cli",
+            "download",
+            "--resume-download",
+            "Kwai-Kolors/Kolors",
+            "--local-dir",
+            local_dir,
+        ],
+        check=True,
+    )
+
+    ip_adapter_path = f"{local_dir}/../Kolors-IP-Adapter-Plus"
+    subprocess.run(
+        [
+            "huggingface-cli",
+            "download",
+            "--resume-download",
+            "Kwai-Kolors/Kolors-IP-Adapter-Plus",
+            "--local-dir",
+            ip_adapter_path,
+        ],
+        check=True,
+    )
+
+
+def build_text2img_ip_pipeline(
+    ckpt_dir: str,
+    ref_scale: float,
+    device: str = "cuda",
+) -> StableDiffusionXLPipelineIP:
+    download_kolors_weights(ckpt_dir)
+
+    text_encoder = ChatGLMModel.from_pretrained(
+        f"{ckpt_dir}/text_encoder", torch_dtype=torch.float16
+    ).half()
+    tokenizer = ChatGLMTokenizer.from_pretrained(f"{ckpt_dir}/text_encoder")
+    vae = AutoencoderKL.from_pretrained(
+        f"{ckpt_dir}/vae", revision=None
+    ).half()
+    scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler")
+    unet = UNet2DConditionModelIP.from_pretrained(
+        f"{ckpt_dir}/unet", revision=None
+    ).half()
+    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+        f"{ckpt_dir}/../Kolors-IP-Adapter-Plus/image_encoder",
+        ignore_mismatched_sizes=True,
+    ).to(dtype=torch.float16)
+    clip_image_processor = CLIPImageProcessor(size=336, crop_size=336)
+
+    pipe = StableDiffusionXLPipelineIP(
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        unet=unet,
+        scheduler=scheduler,
+        image_encoder=image_encoder,
+        feature_extractor=clip_image_processor,
+        force_zeros_for_empty_prompt=False,
+    )
+
+    if hasattr(pipe.unet, "encoder_hid_proj"):
+        pipe.unet.text_encoder_hid_proj = pipe.unet.encoder_hid_proj
+
+    pipe.load_ip_adapter(
+        f"{ckpt_dir}/../Kolors-IP-Adapter-Plus",
+        subfolder="",
+        weight_name=["ip_adapter_plus_general.bin"],
+    )
+    pipe.set_ip_adapter_scale([ref_scale])
+
+    pipe = pipe.to(device)
+    pipe.image_encoder = pipe.image_encoder.to(device)
+    # pipe.enable_model_cpu_offload()
+    # # pipe.enable_xformers_memory_efficient_attention()
+    # pipe.enable_vae_slicing()
+
+    return pipe
+
+
+def build_text2img_pipeline(
+    ckpt_dir: str,
+    device: str = "cuda",
+) -> StableDiffusionXLPipeline:
+    download_kolors_weights(ckpt_dir)
+
+    text_encoder = ChatGLMModel.from_pretrained(
+        f"{ckpt_dir}/text_encoder", torch_dtype=torch.float16
+    ).half()
+    tokenizer = ChatGLMTokenizer.from_pretrained(f"{ckpt_dir}/text_encoder")
+    vae = AutoencoderKL.from_pretrained(
+        f"{ckpt_dir}/vae", revision=None
+    ).half()
+    scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler")
+    unet = UNet2DConditionModel.from_pretrained(
+        f"{ckpt_dir}/unet", revision=None
+    ).half()
+    pipe = StableDiffusionXLPipeline(
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        unet=unet,
+        scheduler=scheduler,
+        force_zeros_for_empty_prompt=False,
+    )
+    pipe = pipe.to(device)
+    # pipe.enable_model_cpu_offload()
+    # pipe.enable_xformers_memory_efficient_attention()
+
+    return pipe
+
+
+def text2img_gen(
+    prompt: str,
+    n_sample: int,
+    guidance_scale: float,
+    pipeline: StableDiffusionXLPipeline | StableDiffusionXLPipelineIP,
+    ip_image: Image.Image | str = None,
+    image_wh: tuple[int, int] = [1024, 1024],
+    infer_step: int = 50,
+    ip_image_size: int = 512,
+    seed: int = None,
+) -> list[Image.Image]:
+    prompt = PROMPT_KAPPEND.format(object=prompt.strip())
+    logger.info(f"Processing prompt: {prompt}")
+
+    generator = None
+    if seed is not None:
+        generator = torch.Generator(pipeline.device).manual_seed(seed)
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+        random.seed(seed)
+
+    kwargs = dict(
+        prompt=prompt,
+        height=image_wh[1],
+        width=image_wh[0],
+        num_inference_steps=infer_step,
+        guidance_scale=guidance_scale,
+        num_images_per_prompt=n_sample,
+        generator=generator,
+    )
+    if ip_image is not None:
+        if isinstance(ip_image, str):
+            ip_image = Image.open(ip_image)
+        ip_image = ip_image.resize((ip_image_size, ip_image_size))
+        kwargs.update(ip_adapter_image=[ip_image])
+
+    return pipeline(**kwargs).images

embodied_gen/models/texture_model.py

@@ -0,0 +1,112 @@
+# 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.
+
+
+import os
+
+import torch
+from diffusers import AutoencoderKL, DiffusionPipeline, EulerDiscreteScheduler
+from huggingface_hub import snapshot_download
+from kolors.models.controlnet import ControlNetModel
+from kolors.models.modeling_chatglm import ChatGLMModel
+from kolors.models.tokenization_chatglm import ChatGLMTokenizer
+from kolors.models.unet_2d_condition import UNet2DConditionModel
+from kolors.pipelines.pipeline_controlnet_xl_kolors_img2img import (
+    StableDiffusionXLControlNetImg2ImgPipeline,
+)
+from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
+from embodied_gen.models.text_model import download_kolors_weights
+from embodied_gen.utils.log import logger
+
+__all__ = [
+    "build_texture_gen_pipe",
+]
+
+
+def build_texture_gen_pipe(
+    base_ckpt_dir: str,
+    controlnet_ckpt: str = None,
+    ip_adapt_scale: float = 0,
+    device: str = "cuda",
+) -> DiffusionPipeline:
+    download_kolors_weights(f"{base_ckpt_dir}/Kolors")
+    logger.info(f"Load Kolors weights...")
+    tokenizer = ChatGLMTokenizer.from_pretrained(
+        f"{base_ckpt_dir}/Kolors/text_encoder"
+    )
+    text_encoder = ChatGLMModel.from_pretrained(
+        f"{base_ckpt_dir}/Kolors/text_encoder", torch_dtype=torch.float16
+    ).half()
+    vae = AutoencoderKL.from_pretrained(
+        f"{base_ckpt_dir}/Kolors/vae", revision=None
+    ).half()
+    unet = UNet2DConditionModel.from_pretrained(
+        f"{base_ckpt_dir}/Kolors/unet", revision=None
+    ).half()
+    scheduler = EulerDiscreteScheduler.from_pretrained(
+        f"{base_ckpt_dir}/Kolors/scheduler"
+    )
+
+    if controlnet_ckpt is None:
+        suffix = "texture_gen_mv_v1"  # "geo_cond_mv"
+        model_path = snapshot_download(
+            repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
+        )
+        controlnet_ckpt = os.path.join(model_path, suffix)
+
+    controlnet = ControlNetModel.from_pretrained(
+        controlnet_ckpt, use_safetensors=True
+    ).half()
+
+    # IP-Adapter model
+    image_encoder = None
+    clip_image_processor = None
+    if ip_adapt_scale > 0:
+        image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+            f"{base_ckpt_dir}/Kolors-IP-Adapter-Plus/image_encoder",
+            # ignore_mismatched_sizes=True,
+        ).to(dtype=torch.float16)
+        ip_img_size = 336
+        clip_image_processor = CLIPImageProcessor(
+            size=ip_img_size, crop_size=ip_img_size
+        )
+
+    pipe = StableDiffusionXLControlNetImg2ImgPipeline(
+        vae=vae,
+        controlnet=controlnet,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        unet=unet,
+        scheduler=scheduler,
+        image_encoder=image_encoder,
+        feature_extractor=clip_image_processor,
+        force_zeros_for_empty_prompt=False,
+    )
+
+    if ip_adapt_scale > 0:
+        if hasattr(pipe.unet, "encoder_hid_proj"):
+            pipe.unet.text_encoder_hid_proj = pipe.unet.encoder_hid_proj
+        pipe.load_ip_adapter(
+            f"{base_ckpt_dir}/Kolors-IP-Adapter-Plus",
+            subfolder="",
+            weight_name=["ip_adapter_plus_general.bin"],
+        )
+        pipe.set_ip_adapter_scale([ip_adapt_scale])
+
+    pipe = pipe.to(device)
+    # pipe.enable_model_cpu_offload()
+
+    return pipe

embodied_gen/scripts/compose_layout.py

@@ -0,0 +1,79 @@
+# 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.
+
+import json
+import os
+import shutil
+from dataclasses import dataclass
+
+import tyro
+from embodied_gen.scripts.simulate_sapien import entrypoint as sim_cli
+from embodied_gen.utils.enum import LayoutInfo
+from embodied_gen.utils.geometry import bfs_placement, compose_mesh_scene
+from embodied_gen.utils.log import logger
+
+
+@dataclass
+class LayoutPlacementConfig:
+    layout_path: str
+    output_dir: str | None = None
+    seed: int | None = None
+    max_attempts: int = 1000
+    output_iscene: bool = False
+    insert_robot: bool = False
+
+
+def entrypoint(**kwargs):
+    if kwargs is None or len(kwargs) == 0:
+        args = tyro.cli(LayoutPlacementConfig)
+    else:
+        args = LayoutPlacementConfig(**kwargs)
+
+    output_dir = (
+        args.output_dir
+        if args.output_dir is not None
+        else os.path.dirname(args.layout_path)
+    )
+    os.makedirs(output_dir, exist_ok=True)
+    out_scene_path = f"{output_dir}/Iscene.glb"
+    out_layout_path = f"{output_dir}/layout.json"
+
+    layout_info = bfs_placement(args.layout_path, seed=args.seed)
+    origin_dir = os.path.dirname(args.layout_path)
+    for key in layout_info.assets:
+        src = f"{origin_dir}/{layout_info.assets[key]}"
+        dst = f"{output_dir}/{layout_info.assets[key]}"
+        if src == dst:
+            continue
+        shutil.copytree(src, dst, dirs_exist_ok=True)
+
+    with open(out_layout_path, "w") as f:
+        json.dump(layout_info.to_dict(), f, indent=4)
+
+    if args.output_iscene:
+        compose_mesh_scene(layout_info, out_scene_path)
+
+    sim_cli(
+        layout_path=out_layout_path,
+        output_dir=output_dir,
+        insert_robot=args.insert_robot,
+    )
+
+    logger.info(f"Layout placement completed in {output_dir}")
+
+
+if __name__ == "__main__":
+    entrypoint()

embodied_gen/scripts/gen_layout.py

@@ -0,0 +1,171 @@
+# 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.
+
+import gc
+import json
+import os
+from dataclasses import dataclass, field
+from shutil import copytree
+from time import time
+from typing import Optional
+
+import torch
+import tyro
+from embodied_gen.models.layout import build_scene_layout
+from embodied_gen.scripts.simulate_sapien import entrypoint as sim_cli
+from embodied_gen.scripts.textto3d import text_to_3d
+from embodied_gen.utils.config import GptParamsConfig
+from embodied_gen.utils.enum import LayoutInfo, Scene3DItemEnum
+from embodied_gen.utils.geometry import bfs_placement, compose_mesh_scene
+from embodied_gen.utils.gpt_clients import GPT_CLIENT
+from embodied_gen.utils.log import logger
+from embodied_gen.utils.process_media import (
+    load_scene_dict,
+    parse_text_prompts,
+)
+from embodied_gen.validators.quality_checkers import SemanticMatcher
+
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+
+@dataclass
+class LayoutGenConfig:
+    task_descs: list[str]
+    output_root: str
+    bg_list: str = "outputs/bg_scenes/scene_list.txt"
+    n_img_sample: int = 3
+    text_guidance_scale: float = 7.0
+    img_denoise_step: int = 25
+    n_image_retry: int = 4
+    n_asset_retry: int = 3
+    n_pipe_retry: int = 2
+    seed_img: Optional[int] = None
+    seed_3d: Optional[int] = None
+    seed_layout: Optional[int] = None
+    keep_intermediate: bool = False
+    output_iscene: bool = False
+    insert_robot: bool = False
+    gpt_params: GptParamsConfig = field(
+        default_factory=lambda: GptParamsConfig(
+            temperature=1.0,
+            top_p=0.95,
+            frequency_penalty=0.3,
+            presence_penalty=0.5,
+        )
+    )
+
+
+def entrypoint() -> None:
+    args = tyro.cli(LayoutGenConfig)
+    SCENE_MATCHER = SemanticMatcher(GPT_CLIENT)
+    task_descs = parse_text_prompts(args.task_descs)
+    scene_dict = load_scene_dict(args.bg_list)
+    gpt_params = args.gpt_params.to_dict()
+    for idx, task_desc in enumerate(task_descs):
+        logger.info(f"Generate Layout and 3D scene for task: {task_desc}")
+        output_root = f"{args.output_root}/task_{idx:04d}"
+        scene_graph_path = f"{output_root}/scene_tree.jpg"
+        start_time = time()
+        layout_info: LayoutInfo = build_scene_layout(
+            task_desc, scene_graph_path, gpt_params
+        )
+        prompts_mapping = {v: k for k, v in layout_info.objs_desc.items()}
+        prompts = [
+            v
+            for k, v in layout_info.objs_desc.items()
+            if layout_info.objs_mapping[k] != Scene3DItemEnum.BACKGROUND.value
+        ]
+
+        for prompt in prompts:
+            node = prompts_mapping[prompt]
+            generation_log = text_to_3d(
+                prompts=[
+                    prompt,
+                ],
+                output_root=output_root,
+                asset_names=[
+                    node,
+                ],
+                n_img_sample=args.n_img_sample,
+                text_guidance_scale=args.text_guidance_scale,
+                img_denoise_step=args.img_denoise_step,
+                n_image_retry=args.n_image_retry,
+                n_asset_retry=args.n_asset_retry,
+                n_pipe_retry=args.n_pipe_retry,
+                seed_img=args.seed_img,
+                seed_3d=args.seed_3d,
+                keep_intermediate=args.keep_intermediate,
+            )
+            layout_info.assets.update(generation_log["assets"])
+            layout_info.quality.update(generation_log["quality"])
+
+        # Background GEN (for efficiency, temp use retrieval instead)
+        bg_node = layout_info.relation[Scene3DItemEnum.BACKGROUND.value]
+        text = layout_info.objs_desc[bg_node]
+        match_key = SCENE_MATCHER.query(
+            text, str(scene_dict), params=gpt_params
+        )
+        n_max_attempt = 10
+        while match_key not in scene_dict and n_max_attempt > 0:
+            logger.error(
+                f"Cannot find matched scene {match_key}, retrying left {n_max_attempt}..."
+            )
+            match_key = SCENE_MATCHER.query(
+                text, str(scene_dict), params=gpt_params
+            )
+            n_max_attempt -= 1
+
+        match_scene_path = f"{os.path.dirname(args.bg_list)}/{match_key}"
+        bg_save_dir = os.path.join(output_root, "background")
+        copytree(match_scene_path, bg_save_dir, dirs_exist_ok=True)
+        layout_info.assets[bg_node] = "background"
+
+        # BFS layout placement.
+        layout_path = f"{output_root}/layout.json"
+        with open(layout_path, "w") as f:
+            json.dump(layout_info.to_dict(), f, indent=4)
+
+        layout_info = bfs_placement(
+            layout_path,
+            seed=args.seed_layout,
+        )
+        layout_path = f"{output_root}/layout.json"
+        with open(layout_path, "w") as f:
+            json.dump(layout_info.to_dict(), f, indent=4)
+
+        if args.output_iscene:
+            compose_mesh_scene(layout_info, f"{output_root}/Iscene.glb")
+
+        sim_cli(
+            layout_path=layout_path,
+            output_dir=output_root,
+            insert_robot=args.insert_robot,
+        )
+
+        torch.cuda.empty_cache()
+        gc.collect()
+
+        elapsed_time = (time() - start_time) / 60
+        logger.info(
+            f"Layout generation done for {scene_graph_path}, layout result "
+            f"in {layout_path}, finished in {elapsed_time:.2f} mins."
+        )
+
+    logger.info(f"All tasks completed in {args.output_root}")
+
+
+if __name__ == "__main__":
+    entrypoint()

embodied_gen/scripts/gen_scene3d.py

@@ -0,0 +1,191 @@
+import logging
+import os
+import random
+import time
+import warnings
+from dataclasses import dataclass, field
+from shutil import copy, rmtree
+
+import torch
+import tyro
+from huggingface_hub import snapshot_download
+from packaging import version
+
+# Suppress warnings
+warnings.filterwarnings("ignore", category=FutureWarning)
+logging.getLogger("transformers").setLevel(logging.ERROR)
+logging.getLogger("diffusers").setLevel(logging.ERROR)
+
+# TorchVision monkey patch for >0.16
+if version.parse(torch.__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.rgb_to_grayscale = TF.rgb_to_grayscale
+    sys.modules["torchvision.transforms.functional_tensor"] = functional_tensor
+
+from gsplat.distributed import cli
+from txt2panoimg import Text2360PanoramaImagePipeline
+from embodied_gen.trainer.gsplat_trainer import (
+    DefaultStrategy,
+    GsplatTrainConfig,
+)
+from embodied_gen.trainer.gsplat_trainer import entrypoint as gsplat_entrypoint
+from embodied_gen.trainer.pono2mesh_trainer import Pano2MeshSRPipeline
+from embodied_gen.utils.config import Pano2MeshSRConfig
+from embodied_gen.utils.gaussian import restore_scene_scale_and_position
+from embodied_gen.utils.gpt_clients import GPT_CLIENT
+from embodied_gen.utils.log import logger
+from embodied_gen.utils.process_media import is_image_file, parse_text_prompts
+from embodied_gen.validators.quality_checkers import (
+    PanoHeightEstimator,
+    PanoImageOccChecker,
+)
+
+__all__ = [
+    "generate_pano_image",
+    "entrypoint",
+]
+
+
+@dataclass
+class Scene3DGenConfig:
+    prompts: list[str]  # Text desc of indoor room or style reference image.
+    output_dir: str
+    seed: int | None = None
+    real_height: float | None = None  # The real height of the room in meters.
+    pano_image_only: bool = False
+    disable_pano_check: bool = False
+    keep_middle_result: bool = False
+    n_retry: int = 7
+    gs3d: GsplatTrainConfig = field(
+        default_factory=lambda: GsplatTrainConfig(
+            strategy=DefaultStrategy(verbose=True),
+            max_steps=4000,
+            init_opa=0.9,
+            opacity_reg=2e-3,
+            sh_degree=0,
+            means_lr=1e-4,
+            scales_lr=1e-3,
+        )
+    )
+
+
+def generate_pano_image(
+    prompt: str,
+    output_path: str,
+    pipeline,
+    seed: int,
+    n_retry: int,
+    checker=None,
+    num_inference_steps: int = 40,
+) -> None:
+    for i in range(n_retry):
+        logger.info(
+            f"GEN Panorama: Retry {i+1}/{n_retry} for prompt: {prompt}, seed: {seed}"
+        )
+        if is_image_file(prompt):
+            raise NotImplementedError("Image mode not implemented yet.")
+        else:
+            txt_prompt = f"{prompt}, spacious, empty, wide open, open floor, minimal furniture"
+            inputs = {
+                "prompt": txt_prompt,
+                "num_inference_steps": num_inference_steps,
+                "upscale": False,
+                "seed": seed,
+            }
+            pano_image = pipeline(inputs)
+
+        pano_image.save(output_path)
+        if checker is None:
+            break
+
+        flag, response = checker(pano_image)
+        logger.warning(f"{response}, image saved in {output_path}")
+        if flag is True or flag is None:
+            break
+
+        seed = random.randint(0, 100000)
+
+    return
+
+
+def entrypoint(*args, **kwargs):
+    cfg = tyro.cli(Scene3DGenConfig)
+
+    # Init global models.
+    model_path = snapshot_download("archerfmy0831/sd-t2i-360panoimage")
+    IMG2PANO_PIPE = Text2360PanoramaImagePipeline(
+        model_path, torch_dtype=torch.float16, device="cuda"
+    )
+    PANOMESH_CFG = Pano2MeshSRConfig()
+    PANO2MESH_PIPE = Pano2MeshSRPipeline(PANOMESH_CFG)
+    PANO_CHECKER = PanoImageOccChecker(GPT_CLIENT, box_hw=[95, 1000])
+    PANOHEIGHT_ESTOR = PanoHeightEstimator(GPT_CLIENT)
+
+    prompts = parse_text_prompts(cfg.prompts)
+    for idx, prompt in enumerate(prompts):
+        start_time = time.time()
+        output_dir = os.path.join(cfg.output_dir, f"scene_{idx:04d}")
+        os.makedirs(output_dir, exist_ok=True)
+        pano_path = os.path.join(output_dir, "pano_image.png")
+        with open(f"{output_dir}/prompt.txt", "w") as f:
+            f.write(prompt)
+
+        generate_pano_image(
+            prompt,
+            pano_path,
+            IMG2PANO_PIPE,
+            cfg.seed if cfg.seed is not None else random.randint(0, 100000),
+            cfg.n_retry,
+            checker=None if cfg.disable_pano_check else PANO_CHECKER,
+        )
+
+        if cfg.pano_image_only:
+            continue
+
+        logger.info("GEN and REPAIR Mesh from Panorama...")
+        PANO2MESH_PIPE(pano_path, output_dir)
+
+        logger.info("TRAIN 3DGS from Mesh Init and Cube Image...")
+        cfg.gs3d.data_dir = output_dir
+        cfg.gs3d.result_dir = f"{output_dir}/gaussian"
+        cfg.gs3d.adjust_steps(cfg.gs3d.steps_scaler)
+        torch.set_default_device("cpu")  # recover default setting.
+        cli(gsplat_entrypoint, cfg.gs3d, verbose=True)
+
+        # Clean up the middle results.
+        gs_path = (
+            f"{cfg.gs3d.result_dir}/ply/point_cloud_{cfg.gs3d.max_steps-1}.ply"
+        )
+        copy(gs_path, f"{output_dir}/gs_model.ply")
+        video_path = f"{cfg.gs3d.result_dir}/renders/video_step{cfg.gs3d.max_steps-1}.mp4"
+        copy(video_path, f"{output_dir}/video.mp4")
+        gs_cfg_path = f"{cfg.gs3d.result_dir}/cfg.yml"
+        copy(gs_cfg_path, f"{output_dir}/gsplat_cfg.yml")
+        if not cfg.keep_middle_result:
+            rmtree(cfg.gs3d.result_dir, ignore_errors=True)
+            os.remove(f"{output_dir}/{PANOMESH_CFG.gs_data_file}")
+
+        real_height = (
+            PANOHEIGHT_ESTOR(pano_path)
+            if cfg.real_height is None
+            else cfg.real_height
+        )
+        gs_path = os.path.join(output_dir, "gs_model.ply")
+        mesh_path = os.path.join(output_dir, "mesh_model.ply")
+        restore_scene_scale_and_position(real_height, mesh_path, gs_path)
+
+        elapsed_time = (time.time() - start_time) / 60
+        logger.info(
+            f"FINISHED 3D scene generation in {output_dir} in {elapsed_time:.2f} mins."
+        )
+
+
+if __name__ == "__main__":
+    entrypoint()

embodied_gen/scripts/gen_texture.py

@@ -0,0 +1,123 @@
+import os
+import shutil
+from dataclasses import dataclass
+
+import tyro
+from embodied_gen.data.backproject_v2 import entrypoint as backproject_api
+from embodied_gen.data.differentiable_render import entrypoint as drender_api
+from embodied_gen.data.utils import as_list
+from embodied_gen.models.delight_model import DelightingModel
+from embodied_gen.models.sr_model import ImageRealESRGAN
+from embodied_gen.scripts.render_mv import (
+    build_texture_gen_pipe,
+)
+from embodied_gen.scripts.render_mv import infer_pipe as render_mv_api
+from embodied_gen.utils.log import logger
+
+
+@dataclass
+class TextureGenConfig:
+    mesh_path: str | list[str]
+    prompt: str | list[str]
+    output_root: str
+    controlnet_cond_scale: float = 0.7
+    guidance_scale: float = 9
+    strength: float = 0.9
+    num_inference_steps: int = 40
+    delight: bool = True
+    seed: int = 0
+    base_ckpt_dir: str = "./weights"
+    texture_size: int = 2048
+    ip_adapt_scale: float = 0.0
+    ip_img_path: str | list[str] | None = None
+
+
+def entrypoint() -> None:
+    cfg = tyro.cli(TextureGenConfig)
+    cfg.mesh_path = as_list(cfg.mesh_path)
+    cfg.prompt = as_list(cfg.prompt)
+    cfg.ip_img_path = as_list(cfg.ip_img_path)
+    assert len(cfg.mesh_path) == len(cfg.prompt)
+
+    # Pre-load models.
+    if cfg.ip_adapt_scale > 0:
+        PIPELINE = build_texture_gen_pipe(
+            base_ckpt_dir="./weights",
+            ip_adapt_scale=cfg.ip_adapt_scale,
+            device="cuda",
+        )
+    else:
+        PIPELINE = build_texture_gen_pipe(
+            base_ckpt_dir="./weights",
+            ip_adapt_scale=0,
+            device="cuda",
+        )
+    DELIGHT = None
+    if cfg.delight:
+        DELIGHT = DelightingModel()
+    IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
+
+    for idx in range(len(cfg.mesh_path)):
+        mesh_path = cfg.mesh_path[idx]
+        prompt = cfg.prompt[idx]
+        uuid = os.path.splitext(os.path.basename(mesh_path))[0]
+        output_root = os.path.join(cfg.output_root, uuid)
+        drender_api(
+            mesh_path=mesh_path,
+            output_root=f"{output_root}/condition",
+            uuid=uuid,
+        )
+        render_mv_api(
+            index_file=f"{output_root}/condition/index.json",
+            controlnet_cond_scale=cfg.controlnet_cond_scale,
+            guidance_scale=cfg.guidance_scale,
+            strength=cfg.strength,
+            num_inference_steps=cfg.num_inference_steps,
+            ip_adapt_scale=cfg.ip_adapt_scale,
+            ip_img_path=(
+                None if cfg.ip_img_path is None else cfg.ip_img_path[idx]
+            ),
+            prompt=prompt,
+            save_dir=f"{output_root}/multi_view",
+            sub_idxs=[[0, 1, 2], [3, 4, 5]],
+            pipeline=PIPELINE,
+            seed=cfg.seed,
+        )
+        textured_mesh = backproject_api(
+            delight_model=DELIGHT,
+            imagesr_model=IMAGESR_MODEL,
+            mesh_path=mesh_path,
+            color_path=f"{output_root}/multi_view/color_sample0.png",
+            output_path=f"{output_root}/texture_mesh/{uuid}.obj",
+            save_glb_path=f"{output_root}/texture_mesh/{uuid}.glb",
+            skip_fix_mesh=True,
+            delight=cfg.delight,
+            no_save_delight_img=True,
+            texture_wh=[cfg.texture_size, cfg.texture_size],
+        )
+        drender_api(
+            mesh_path=f"{output_root}/texture_mesh/{uuid}.obj",
+            output_root=f"{output_root}/texture_mesh",
+            uuid=uuid,
+            num_images=90,
+            elevation=[20],
+            with_mtl=True,
+            gen_color_mp4=True,
+            pbr_light_factor=1.2,
+        )
+
+        # Re-organize folders
+        shutil.rmtree(f"{output_root}/condition")
+        shutil.copy(
+            f"{output_root}/texture_mesh/{uuid}/color.mp4",
+            f"{output_root}/color.mp4",
+        )
+        shutil.rmtree(f"{output_root}/texture_mesh/{uuid}")
+
+        logger.info(
+            f"Successfully generate textured mesh in {output_root}/texture_mesh"
+        )
+
+
+if __name__ == "__main__":
+    entrypoint()

embodied_gen/scripts/imageto3d.py

@@ -0,0 +1,350 @@
+# 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.
+
+
+import argparse
+import os
+import random
+import sys
+from glob import glob
+from shutil import copy, copytree, rmtree
+
+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.utils import delete_dir, trellis_preprocess
+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.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
+from embodied_gen.utils.process_media import merge_images_video
+from embodied_gen.utils.tags import VERSION
+from embodied_gen.utils.trender import render_video
+from embodied_gen.validators.quality_checkers import (
+    BaseChecker,
+    ImageAestheticChecker,
+    ImageSegChecker,
+    MeshGeoChecker,
+)
+from embodied_gen.validators.urdf_convertor import URDFGenerator
+
+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.pipelines import TrellisImageTo3DPipeline
+
+os.environ["TORCH_EXTENSIONS_DIR"] = os.path.expanduser(
+    "~/.cache/torch_extensions"
+)
+os.environ["GRADIO_ANALYTICS_ENABLED"] = "false"
+os.environ["SPCONV_ALGO"] = "native"
+random.seed(0)
+
+logger.info("Loading Image3D Models...")
+DELIGHT = DelightingModel()
+IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
+RBG_REMOVER = RembgRemover()
+PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
+    "microsoft/TRELLIS-image-large"
+)
+# PIPELINE.cuda()
+SEG_CHECKER = ImageSegChecker(GPT_CLIENT)
+GEO_CHECKER = MeshGeoChecker(GPT_CLIENT)
+AESTHETIC_CHECKER = ImageAestheticChecker()
+CHECKERS = [GEO_CHECKER, SEG_CHECKER, AESTHETIC_CHECKER]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Image to 3D pipeline args.")
+    parser.add_argument(
+        "--image_path", type=str, nargs="+", help="Path to the input images."
+    )
+    parser.add_argument(
+        "--image_root", type=str, help="Path to the input images folder."
+    )
+    parser.add_argument(
+        "--output_root",
+        type=str,
+        help="Root directory for saving outputs.",
+    )
+    parser.add_argument(
+        "--height_range",
+        type=str,
+        default=None,
+        help="The hight in meter to restore the mesh real size.",
+    )
+    parser.add_argument(
+        "--mass_range",
+        type=str,
+        default=None,
+        help="The mass in kg to restore the mesh real weight.",
+    )
+    parser.add_argument("--asset_type", type=str, nargs="+", default=None)
+    parser.add_argument("--skip_exists", action="store_true")
+    parser.add_argument("--version", type=str, default=VERSION)
+    parser.add_argument("--keep_intermediate", action="store_true")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument(
+        "--n_retry",
+        type=int,
+        default=2,
+    )
+    parser.add_argument("--disable_decompose_convex", action="store_true")
+    parser.add_argument(
+        "--texture_wh", type=int, nargs=2, default=[2048, 2048]
+    )
+    args, unknown = parser.parse_known_args()
+
+    return args
+
+
+def entrypoint(**kwargs):
+    args = parse_args()
+    for k, v in kwargs.items():
+        if hasattr(args, k) and v is not None:
+            setattr(args, k, v)
+
+    assert (
+        args.image_path or args.image_root
+    ), "Please provide either --image_path or --image_root."
+    if not args.image_path:
+        args.image_path = glob(os.path.join(args.image_root, "*.png"))
+        args.image_path += glob(os.path.join(args.image_root, "*.jpg"))
+        args.image_path += glob(os.path.join(args.image_root, "*.jpeg"))
+
+    for idx, image_path in enumerate(args.image_path):
+        try:
+            filename = os.path.basename(image_path).split(".")[0]
+            output_root = args.output_root
+            if args.image_root is not None or len(args.image_path) > 1:
+                output_root = os.path.join(output_root, filename)
+            os.makedirs(output_root, exist_ok=True)
+
+            mesh_out = f"{output_root}/{filename}.obj"
+            if args.skip_exists and os.path.exists(mesh_out):
+                logger.warning(
+                    f"Skip {image_path}, already processed in {mesh_out}"
+                )
+                continue
+
+            image = Image.open(image_path)
+            image.save(f"{output_root}/{filename}_raw.png")
+
+            # Segmentation: Get segmented image using Rembg.
+            seg_path = f"{output_root}/{filename}_cond.png"
+            seg_image = RBG_REMOVER(image) if image.mode != "RGBA" else image
+            seg_image = trellis_preprocess(seg_image)
+            seg_image.save(seg_path)
+
+            seed = args.seed
+            asset_node = "unknown"
+            if isinstance(args.asset_type, list) and args.asset_type[idx]:
+                asset_node = args.asset_type[idx]
+            for try_idx in range(args.n_retry):
+                logger.info(
+                    f"Try: {try_idx + 1}/{args.n_retry}, Seed: {seed}, Prompt: {seg_path}"
+                )
+                # Run the pipeline
+                try:
+                    PIPELINE.cuda()
+                    outputs = PIPELINE.run(
+                        seg_image,
+                        preprocess_image=False,
+                        seed=(
+                            random.randint(0, 100000) if seed is None else seed
+                        ),
+                        # Optional parameters
+                        # sparse_structure_sampler_params={
+                        #     "steps": 12,
+                        #     "cfg_strength": 7.5,
+                        # },
+                        # slat_sampler_params={
+                        #     "steps": 12,
+                        #     "cfg_strength": 3,
+                        # },
+                    )
+                    PIPELINE.cpu()
+                    torch.cuda.empty_cache()
+                except Exception as e:
+                    logger.error(
+                        f"[Pipeline Failed] process {image_path}: {e}, skip."
+                    )
+                    continue
+
+                gs_model = outputs["gaussian"][0]
+                mesh_model = outputs["mesh"][0]
+
+                # Save the raw Gaussian model
+                gs_path = mesh_out.replace(".obj", "_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]]
+                gs_add_rot = [[1, 0, 0], [0, -1, 0], [0, 0, -1]]
+                mesh_add_rot = [[1, 0, 0], [0, 0, -1], [0, 1, 0]]
+
+                # Addtional rotation for GS to align mesh.
+                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")
+                GaussianOperator.resave_ply(
+                    in_ply=gs_path,
+                    out_ply=aligned_gs_path,
+                    instance_pose=pose,
+                    device="cpu",
+                )
+                color_path = os.path.join(output_root, "color.png")
+                render_gs_api(
+                    input_gs=aligned_gs_path,
+                    output_path=color_path,
+                    elevation=[20, -10, 60, -50],
+                    num_images=12,
+                )
+
+                color_img = Image.open(color_path)
+                keep_height = int(color_img.height * 2 / 3)
+                crop_img = color_img.crop((0, 0, color_img.width, keep_height))
+                geo_flag, geo_result = GEO_CHECKER([crop_img], text=asset_node)
+                logger.warning(
+                    f"{GEO_CHECKER.__class__.__name__}: {geo_result} for {seg_path}"
+                )
+                if geo_flag is True or geo_flag is None:
+                    break
+
+                seed = random.randint(0, 100000) if seed is not None else None
+
+            # Render the video for generated 3D asset.
+            color_images = render_video(gs_model)["color"]
+            normal_images = render_video(mesh_model)["normal"]
+            video_path = os.path.join(output_root, "gs_mesh.mp4")
+            merge_images_video(color_images, normal_images, video_path)
+
+            mesh = trimesh.Trimesh(
+                vertices=mesh_model.vertices.cpu().numpy(),
+                faces=mesh_model.faces.cpu().numpy(),
+            )
+            mesh.vertices = mesh.vertices @ np.array(mesh_add_rot)
+            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 = backproject_api(
+                delight_model=DELIGHT,
+                imagesr_model=IMAGESR_MODEL,
+                color_path=color_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,
+            )
+
+            mesh_glb_path = os.path.join(output_root, f"{filename}.glb")
+            mesh.export(mesh_glb_path)
+
+            urdf_convertor = URDFGenerator(
+                GPT_CLIENT,
+                render_view_num=4,
+                decompose_convex=not args.disable_decompose_convex,
+            )
+            asset_attrs = {
+                "version": VERSION,
+                "gs_model": f"{urdf_convertor.output_mesh_dir}/{filename}_gs.ply",
+            }
+            if args.height_range:
+                min_height, max_height = map(
+                    float, args.height_range.split("-")
+                )
+                asset_attrs["min_height"] = min_height
+                asset_attrs["max_height"] = max_height
+            if args.mass_range:
+                min_mass, max_mass = map(float, args.mass_range.split("-"))
+                asset_attrs["min_mass"] = min_mass
+                asset_attrs["max_mass"] = max_mass
+            if isinstance(args.asset_type, list) and args.asset_type[idx]:
+                asset_attrs["category"] = args.asset_type[idx]
+            if args.version:
+                asset_attrs["version"] = args.version
+
+            urdf_root = f"{output_root}/URDF_{filename}"
+            urdf_path = urdf_convertor(
+                mesh_path=mesh_obj_path,
+                output_root=urdf_root,
+                **asset_attrs,
+            )
+
+            # Rescale GS and save to URDF/mesh folder.
+            real_height = urdf_convertor.get_attr_from_urdf(
+                urdf_path, attr_name="real_height"
+            )
+            out_gs = f"{urdf_root}/{urdf_convertor.output_mesh_dir}/{filename}_gs.ply"  # noqa
+            GaussianOperator.resave_ply(
+                in_ply=aligned_gs_path,
+                out_ply=out_gs,
+                real_height=real_height,
+                device="cpu",
+            )
+
+            # Quality check and update .urdf file.
+            mesh_out = f"{urdf_root}/{urdf_convertor.output_mesh_dir}/{filename}.obj"  # noqa
+            trimesh.load(mesh_out).export(mesh_out.replace(".obj", ".glb"))
+
+            image_dir = f"{urdf_root}/{urdf_convertor.output_render_dir}/image_color"  # noqa
+            image_paths = glob(f"{image_dir}/*.png")
+            images_list = []
+            for checker in CHECKERS:
+                images = image_paths
+                if isinstance(checker, ImageSegChecker):
+                    images = [
+                        f"{output_root}/{filename}_raw.png",
+                        f"{output_root}/{filename}_cond.png",
+                    ]
+                images_list.append(images)
+
+            qa_results = BaseChecker.validate(CHECKERS, images_list)
+            urdf_convertor.add_quality_tag(urdf_path, qa_results)
+
+            # Organize the final result files
+            result_dir = f"{output_root}/result"
+            if os.path.exists(result_dir):
+                rmtree(result_dir, ignore_errors=True)
+            os.makedirs(result_dir, exist_ok=True)
+            copy(urdf_path, f"{result_dir}/{os.path.basename(urdf_path)}")
+            copytree(
+                f"{urdf_root}/{urdf_convertor.output_mesh_dir}",
+                f"{result_dir}/{urdf_convertor.output_mesh_dir}",
+            )
+            copy(video_path, f"{result_dir}/video.mp4")
+            if not args.keep_intermediate:
+                delete_dir(output_root, keep_subs=["result"])
+
+        except Exception as e:
+            logger.error(f"Failed to process {image_path}: {e}, skip.")
+            continue
+
+    logger.info(f"Processing complete. Outputs saved to {args.output_root}")
+
+
+if __name__ == "__main__":
+    entrypoint()

embodied_gen/scripts/parallel_sim.py

@@ -0,0 +1,166 @@
+# 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_maniskill
+
+monkey_patch_maniskill()
+import json
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import Literal
+
+import gymnasium as gym
+import numpy as np
+import torch
+import tyro
+from mani_skill.utils.wrappers import RecordEpisode
+from tqdm import tqdm
+import embodied_gen.envs.pick_embodiedgen
+from embodied_gen.utils.enum import LayoutInfo, Scene3DItemEnum
+from embodied_gen.utils.log import logger
+from embodied_gen.utils.simulation import FrankaPandaGrasper
+
+
+@dataclass
+class ParallelSimConfig:
+    """CLI parameters for Parallel Sapien simulation."""
+
+    # Environment configuration
+    layout_file: str
+    """Path to the layout JSON file"""
+    output_dir: str
+    """Directory to save recorded videos"""
+    gym_env_name: str = "PickEmbodiedGen-v1"
+    """Name of the Gym environment to use"""
+    num_envs: int = 4
+    """Number of parallel environments"""
+    render_mode: Literal["rgb_array", "hybrid"] = "hybrid"
+    """Rendering mode: rgb_array or hybrid"""
+    enable_shadow: bool = True
+    """Whether to enable shadows in rendering"""
+    control_mode: str = "pd_joint_pos"
+    """Control mode for the agent"""
+
+    # Recording configuration
+    max_steps_per_video: int = 1000
+    """Maximum steps to record per video"""
+    save_trajectory: bool = False
+    """Whether to save trajectory data"""
+
+    # Simulation parameters
+    seed: int = 0
+    """Random seed for environment reset"""
+    warmup_steps: int = 50
+    """Number of warmup steps before action computation"""
+    reach_target_only: bool = True
+    """Whether to only reach target without full action"""
+
+    # Camera settings
+    camera_eye: list[float] = field(default_factory=lambda: [0.9, 0.0, 1.1])
+    """Camera eye position [x, y, z] in global coordiante system"""
+    camera_target_pt: list[float] = field(
+        default_factory=lambda: [0.0, 0.0, 0.9]
+    )
+    """Camera target(look-at) point [x, y, z] in global coordiante system"""
+    image_hw: list[int] = field(default_factory=lambda: [256, 256])
+    """Rendered image height and width [height, width]"""
+    fovy_deg: float = 75
+    """Camera vertical field of view in degrees"""
+
+
+def entrypoint(**kwargs):
+    if kwargs is None or len(kwargs) == 0:
+        cfg = tyro.cli(ParallelSimConfig)
+    else:
+        cfg = ParallelSimConfig(**kwargs)
+
+    env = gym.make(
+        cfg.gym_env_name,
+        num_envs=cfg.num_envs,
+        render_mode=cfg.render_mode,
+        enable_shadow=cfg.enable_shadow,
+        layout_file=cfg.layout_file,
+        control_mode=cfg.control_mode,
+        camera_cfg=dict(
+            camera_eye=cfg.camera_eye,
+            camera_target_pt=cfg.camera_target_pt,
+            image_hw=cfg.image_hw,
+            fovy_deg=cfg.fovy_deg,
+        ),
+    )
+    env = RecordEpisode(
+        env,
+        cfg.output_dir,
+        max_steps_per_video=cfg.max_steps_per_video,
+        save_trajectory=cfg.save_trajectory,
+    )
+    env.reset(seed=cfg.seed)
+
+    default_action = env.unwrapped.agent.init_qpos[:, :8]
+    for _ in tqdm(range(cfg.warmup_steps), desc="SIM Warmup"):
+        # action = env.action_space.sample() # Random action
+        obs, reward, terminated, truncated, info = env.step(default_action)
+
+    grasper = FrankaPandaGrasper(
+        env.unwrapped.agent,
+        env.unwrapped.sim_config.control_freq,
+    )
+
+    layout_data = LayoutInfo.from_dict(json.load(open(cfg.layout_file, "r")))
+    actions = defaultdict(list)
+    # Plan Grasp reach pose for each manipulated object in each env.
+    for env_idx in range(env.num_envs):
+        actors = env.unwrapped.env_actors[f"env{env_idx}"]
+        for node in layout_data.relation[
+            Scene3DItemEnum.MANIPULATED_OBJS.value
+        ]:
+            action = grasper.compute_grasp_action(
+                actor=actors[node]._objs[0],
+                reach_target_only=True,
+                env_idx=env_idx,
+            )
+            actions[node].append(action)
+
+    # Excute the planned actions for each manipulated object in each env.
+    for node in actions:
+        max_env_steps = 0
+        for env_idx in range(env.num_envs):
+            if actions[node][env_idx] is None:
+                continue
+            max_env_steps = max(max_env_steps, len(actions[node][env_idx]))
+
+        action_tensor = np.ones(
+            (max_env_steps, env.num_envs, env.action_space.shape[-1])
+        )
+        action_tensor *= default_action[None, ...]
+        for env_idx in range(env.num_envs):
+            action = actions[node][env_idx]
+            if action is None:
+                continue
+            action_tensor[: len(action), env_idx, :] = action
+
+        for step in tqdm(range(max_env_steps), desc=f"Grasping: {node}"):
+            action = torch.Tensor(action_tensor[step]).to(env.unwrapped.device)
+            env.unwrapped.agent.set_action(action)
+            obs, reward, terminated, truncated, info = env.step(action)
+
+    env.close()
+    logger.info(f"Results saved in {cfg.output_dir}")
+
+
+if __name__ == "__main__":
+    entrypoint()

embodied_gen/scripts/render_gs.py

@@ -0,0 +1,164 @@
+# 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.
+
+
+import argparse
+import logging
+import math
+
+import cv2
+import spaces
+import torch
+from PIL import Image
+from tqdm import tqdm
+from embodied_gen.data.utils import (
+    CameraSetting,
+    init_kal_camera,
+    normalize_vertices_array,
+)
+from embodied_gen.models.gs_model import GaussianOperator
+from embodied_gen.utils.process_media import combine_images_to_grid
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Render GS color images")
+
+    parser.add_argument(
+        "--input_gs", type=str, help="Input render GS.ply path."
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        help="Output grid image path for rendered GS color images.",
+    )
+    parser.add_argument(
+        "--num_images", type=int, default=6, help="Number of images to render."
+    )
+    parser.add_argument(
+        "--elevation",
+        type=float,
+        nargs="+",
+        default=[20.0, -10.0],
+        help="Elevation angles for the camera (default: [20.0, -10.0])",
+    )
+    parser.add_argument(
+        "--distance",
+        type=float,
+        default=5,
+        help="Camera distance (default: 5)",
+    )
+    parser.add_argument(
+        "--resolution_hw",
+        type=int,
+        nargs=2,
+        default=(512, 512),
+        help="Resolution of the output images (default: (512, 512))",
+    )
+    parser.add_argument(
+        "--fov",
+        type=float,
+        default=30,
+        help="Field of view in degrees (default: 30)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        choices=["cpu", "cuda"],
+        default="cuda",
+        help="Device to run on (default: `cuda`)",
+    )
+    parser.add_argument(
+        "--image_size",
+        type=int,
+        default=512,
+        help="Output image size for single view in color grid (default: 512)",
+    )
+
+    args, unknown = parser.parse_known_args()
+
+    return args
+
+
+def load_gs_model(
+    input_gs: str, pre_quat: list[float] = [0.0, 0.7071, 0.0, -0.7071]
+) -> GaussianOperator:
+    gs_model = GaussianOperator.load_from_ply(input_gs)
+    # Normalize vertices to [-1, 1], center to (0, 0, 0).
+    _, scale, center = normalize_vertices_array(gs_model._means)
+    scale, center = float(scale), center.tolist()
+    transpose = [*[v for v in center], *pre_quat]
+    instance_pose = torch.tensor(transpose).to(gs_model.device)
+    gs_model = gs_model.get_gaussians(instance_pose=instance_pose)
+    gs_model.rescale(scale)
+
+    return gs_model
+
+
+@spaces.GPU
+def entrypoint(**kwargs) -> None:
+    args = parse_args()
+    for k, v in kwargs.items():
+        if hasattr(args, k) and v is not None:
+            setattr(args, k, v)
+
+    # Setup camera parameters
+    camera_params = CameraSetting(
+        num_images=args.num_images,
+        elevation=args.elevation,
+        distance=args.distance,
+        resolution_hw=args.resolution_hw,
+        fov=math.radians(args.fov),
+        device=args.device,
+    )
+    camera = init_kal_camera(camera_params, flip_az=True)
+    matrix_mv = camera.view_matrix()  # (n_cam 4 4) world2cam
+    matrix_mv[:, :3, 3] = -matrix_mv[:, :3, 3]
+    w2cs = matrix_mv.to(camera_params.device)
+    c2ws = [torch.linalg.inv(matrix) for matrix in w2cs]
+    Ks = torch.tensor(camera_params.Ks).to(camera_params.device)
+
+    # Load GS model and normalize.
+    gs_model = load_gs_model(args.input_gs, pre_quat=[0.0, 0.0, 1.0, 0.0])
+
+    # Render GS color images.
+    images = []
+    for idx in tqdm(range(len(c2ws)), desc="Rendering GS"):
+        result = gs_model.render(
+            c2ws[idx],
+            Ks=Ks,
+            image_width=camera_params.resolution_hw[1],
+            image_height=camera_params.resolution_hw[0],
+        )
+        color = cv2.resize(
+            result.rgba,
+            (args.image_size, args.image_size),
+            interpolation=cv2.INTER_AREA,
+        )
+        color = cv2.cvtColor(color, cv2.COLOR_BGRA2RGBA)
+        images.append(Image.fromarray(color))
+
+    combine_images_to_grid(images, image_mode="RGBA")[0].save(args.output_path)
+
+    logger.info(f"Saved grid image to {args.output_path}")
+
+
+if __name__ == "__main__":
+    entrypoint()

embodied_gen/scripts/render_mv.py

@@ -0,0 +1,198 @@
+# 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.
+
+
+import logging
+import os
+import random
+from typing import List, Tuple
+
+import fire
+import numpy as np
+import torch
+from diffusers.utils import make_image_grid
+from kolors.pipelines.pipeline_controlnet_xl_kolors_img2img import (
+    StableDiffusionXLControlNetImg2ImgPipeline,
+)
+from PIL import Image, ImageEnhance, ImageFilter
+from torchvision import transforms
+from embodied_gen.data.datasets import Asset3dGenDataset
+from embodied_gen.models.texture_model import build_texture_gen_pipe
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def get_init_noise_image(image: Image.Image) -> Image.Image:
+    blurred_image = image.convert("L").filter(
+        ImageFilter.GaussianBlur(radius=3)
+    )
+
+    enhancer = ImageEnhance.Contrast(blurred_image)
+    image_decreased_contrast = enhancer.enhance(factor=0.5)
+
+    return image_decreased_contrast
+
+
+def infer_pipe(
+    index_file: str,
+    controlnet_ckpt: str = None,
+    uid: str = None,
+    prompt: str = None,
+    controlnet_cond_scale: float = 0.4,
+    control_guidance_end: float = 0.9,
+    strength: float = 1.0,
+    num_inference_steps: int = 50,
+    guidance_scale: float = 10,
+    ip_adapt_scale: float = 0,
+    ip_img_path: str = None,
+    sub_idxs: List[List[int]] = None,
+    num_images_per_prompt: int = 3,  # increase if want similar images.
+    device: str = "cuda",
+    save_dir: str = "infer_vis",
+    seed: int = None,
+    target_hw: tuple[int, int] = (512, 512),
+    pipeline: StableDiffusionXLControlNetImg2ImgPipeline = None,
+) -> str:
+    # sub_idxs = [[0, 1, 2], [3, 4, 5]] # None for single image.
+    if sub_idxs is None:
+        sub_idxs = [[random.randint(0, 5)]]  # 6 views.
+        target_hw = [2 * size for size in target_hw]
+
+    transform_list = [
+        transforms.Resize(
+            target_hw, interpolation=transforms.InterpolationMode.BILINEAR
+        ),
+        transforms.CenterCrop(target_hw),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5]),
+    ]
+    image_transform = transforms.Compose(transform_list)
+    control_transform = transforms.Compose(transform_list[:-1])
+
+    grid_hw = (target_hw[0] * len(sub_idxs), target_hw[1] * len(sub_idxs[0]))
+    dataset = Asset3dGenDataset(
+        index_file, target_hw=grid_hw, sub_idxs=sub_idxs
+    )
+
+    if uid is None:
+        uid = random.choice(list(dataset.meta_info.keys()))
+    if prompt is None:
+        prompt = dataset.meta_info[uid]["capture"]
+    if isinstance(prompt, List) or isinstance(prompt, Tuple):
+        prompt = ", ".join(map(str, prompt))
+    # prompt += "high quality, ultra-clear, high resolution, best quality, 4k"
+    # prompt += "高品质,清晰,细节"
+    prompt += ", high quality, high resolution, best quality"
+    # prompt += ", with diffuse lighting, showing no reflections."
+    logger.info(f"Inference with prompt: {prompt}")
+
+    negative_prompt = "nsfw,阴影,低分辨率,伪影、模糊,霓虹灯,高光,镜面反射"
+
+    control_image = dataset.fetch_sample_grid_images(
+        uid,
+        attrs=["image_view_normal", "image_position", "image_mask"],
+        sub_idxs=sub_idxs,
+        transform=control_transform,
+    )
+
+    color_image = dataset.fetch_sample_grid_images(
+        uid,
+        attrs=["image_color"],
+        sub_idxs=sub_idxs,
+        transform=image_transform,
+    )
+
+    normal_pil, position_pil, mask_pil, color_pil = dataset.visualize_item(
+        control_image,
+        color_image,
+        save_dir=save_dir,
+    )
+
+    if pipeline is None:
+        pipeline = build_texture_gen_pipe(
+            base_ckpt_dir="./weights",
+            controlnet_ckpt=controlnet_ckpt,
+            ip_adapt_scale=ip_adapt_scale,
+            device=device,
+        )
+
+    if ip_adapt_scale > 0 and ip_img_path is not None and len(ip_img_path) > 0:
+        ip_image = Image.open(ip_img_path).convert("RGB")
+        ip_image = ip_image.resize(target_hw[::-1])
+        ip_image = [ip_image]
+        pipeline.set_ip_adapter_scale([ip_adapt_scale])
+    else:
+        ip_image = None
+
+    generator = None
+    if seed is not None:
+        generator = torch.Generator(device).manual_seed(seed)
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+        random.seed(seed)
+
+    init_image = get_init_noise_image(normal_pil)
+    # init_image = get_init_noise_image(color_pil)
+
+    images = []
+    row_num, col_num = 2, 3
+    img_save_paths = []
+    while len(images) < col_num:
+        image = pipeline(
+            prompt=prompt,
+            image=init_image,
+            controlnet_conditioning_scale=controlnet_cond_scale,
+            control_guidance_end=control_guidance_end,
+            strength=strength,
+            control_image=control_image[None, ...],
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=num_images_per_prompt,
+            ip_adapter_image=ip_image,
+            generator=generator,
+        ).images
+        images.extend(image)
+
+    grid_image = [normal_pil, position_pil, color_pil] + images[:col_num]
+    # save_dir = os.path.join(save_dir, uid)
+    os.makedirs(save_dir, exist_ok=True)
+
+    for idx in range(col_num):
+        rgba_image = Image.merge("RGBA", (*images[idx].split(), mask_pil))
+        img_save_path = os.path.join(save_dir, f"color_sample{idx}.png")
+        rgba_image.save(img_save_path)
+        img_save_paths.append(img_save_path)
+
+    sub_idxs = "_".join(
+        [str(item) for sublist in sub_idxs for item in sublist]
+    )
+    save_path = os.path.join(
+        save_dir, f"sample_idx{str(sub_idxs)}_ip{ip_adapt_scale}.jpg"
+    )
+    make_image_grid(grid_image, row_num, col_num).save(save_path)
+    logger.info(f"Visualize in {save_path}")
+
+    return img_save_paths
+
+
+def entrypoint() -> None:
+    fire.Fire(infer_pipe)
+
+
+if __name__ == "__main__":
+    entrypoint()

embodied_gen/scripts/simulate_sapien.py

@@ -0,0 +1,196 @@
+# 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.
+
+
+import json
+import os
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import Literal
+
+import imageio
+import numpy as np
+import torch
+import tyro
+from tqdm import tqdm
+from embodied_gen.models.gs_model import GaussianOperator
+from embodied_gen.utils.enum import LayoutInfo, Scene3DItemEnum
+from embodied_gen.utils.geometry import quaternion_multiply
+from embodied_gen.utils.log import logger
+from embodied_gen.utils.process_media import alpha_blend_rgba
+from embodied_gen.utils.simulation import (
+    SIM_COORD_ALIGN,
+    FrankaPandaGrasper,
+    SapienSceneManager,
+    load_assets_from_layout_file,
+    load_mani_skill_robot,
+    render_images,
+)
+
+
+@dataclass
+class SapienSimConfig:
+    # Simulation settings.
+    layout_path: str
+    output_dir: str
+    sim_freq: int = 200
+    sim_step: int = 400
+    z_offset: float = 0.004
+    init_3dgs_quat: list[float] = field(
+        default_factory=lambda: [0.7071, 0, 0, 0.7071]
+    )  # xyzw
+    device: str = "cuda"
+    control_freq: int = 50
+    insert_robot: bool = False
+    # Camera settings.
+    render_interval: int = 10
+    num_cameras: int = 3
+    camera_radius: float = 0.9
+    camera_height: float = 1.1
+    image_hw: tuple[int, int] = (512, 512)
+    ray_tracing: bool = True
+    fovy_deg: float = 75.0
+    camera_target_pt: list[float] = field(
+        default_factory=lambda: [0.0, 0.0, 0.9]
+    )
+    render_keys: list[
+        Literal[
+            "Color", "Foreground", "Segmentation", "Normal", "Mask", "Depth"
+        ]
+    ] = field(default_factory=lambda: ["Foreground"])
+
+
+def entrypoint(**kwargs):
+    if kwargs is None or len(kwargs) == 0:
+        cfg = tyro.cli(SapienSimConfig)
+    else:
+        cfg = SapienSimConfig(**kwargs)
+
+    scene_manager = SapienSceneManager(
+        cfg.sim_freq, ray_tracing=cfg.ray_tracing
+    )
+    _ = scene_manager.initialize_circular_cameras(
+        num_cameras=cfg.num_cameras,
+        radius=cfg.camera_radius,
+        height=cfg.camera_height,
+        target_pt=cfg.camera_target_pt,
+        image_hw=cfg.image_hw,
+        fovy_deg=cfg.fovy_deg,
+    )
+    with open(cfg.layout_path, "r") as f:
+        layout_data: LayoutInfo = LayoutInfo.from_dict(json.load(f))
+
+    actors = load_assets_from_layout_file(
+        scene_manager.scene,
+        cfg.layout_path,
+        cfg.z_offset,
+    )
+    agent = load_mani_skill_robot(
+        scene_manager.scene, cfg.layout_path, cfg.control_freq
+    )
+
+    frames = defaultdict(list)
+    image_cnt = 0
+    for step in tqdm(range(cfg.sim_step), desc="Simulation"):
+        scene_manager.scene.step()
+        agent.reset(agent.init_qpos)
+        if step % cfg.render_interval != 0:
+            continue
+        scene_manager.scene.update_render()
+        image_cnt += 1
+        for camera in scene_manager.cameras:
+            camera.take_picture()
+            images = render_images(camera, cfg.render_keys)
+            frames[camera.name].append(images)
+
+    actions = dict()
+    if cfg.insert_robot:
+        grasper = FrankaPandaGrasper(
+            agent,
+            cfg.control_freq,
+        )
+        for node in layout_data.relation[
+            Scene3DItemEnum.MANIPULATED_OBJS.value
+        ]:
+            actions[node] = grasper.compute_grasp_action(
+                actor=actors[node], reach_target_only=True
+            )
+
+    if "Foreground" not in cfg.render_keys:
+        return
+
+    asset_root = os.path.dirname(cfg.layout_path)
+    bg_node = layout_data.relation[Scene3DItemEnum.BACKGROUND.value]
+    gs_path = f"{asset_root}/{layout_data.assets[bg_node]}/gs_model.ply"
+    gs_model: GaussianOperator = GaussianOperator.load_from_ply(gs_path)
+    x, y, z, qx, qy, qz, qw = layout_data.position[bg_node]
+    qx, qy, qz, qw = quaternion_multiply([qx, qy, qz, qw], cfg.init_3dgs_quat)
+    init_pose = torch.tensor([x, y, z, qx, qy, qz, qw])
+    gs_model = gs_model.get_gaussians(instance_pose=init_pose)
+
+    bg_images = dict()
+    for camera in scene_manager.cameras:
+        Ks = camera.get_intrinsic_matrix()
+        c2w = camera.get_model_matrix()
+        c2w = c2w @ SIM_COORD_ALIGN
+        result = gs_model.render(
+            torch.tensor(c2w, dtype=torch.float32).to(cfg.device),
+            torch.tensor(Ks, dtype=torch.float32).to(cfg.device),
+            image_width=cfg.image_hw[1],
+            image_height=cfg.image_hw[0],
+        )
+        bg_images[camera.name] = result.rgb[..., ::-1]
+
+    video_frames = []
+    for idx, camera in enumerate(scene_manager.cameras):
+        # Scene rendering
+        if idx == 0:
+            for step in range(image_cnt):
+                rgba = alpha_blend_rgba(
+                    frames[camera.name][step]["Foreground"],
+                    bg_images[camera.name],
+                )
+                video_frames.append(np.array(rgba))
+
+        # Grasp rendering
+        for node in actions:
+            if actions[node] is None:
+                continue
+            logger.info(f"Render SIM grasping in camera {idx} for {node}...")
+            for action in actions[node]:
+                grasp_frames = scene_manager.step_action(
+                    agent,
+                    torch.Tensor(action[None, ...]),
+                    scene_manager.cameras,
+                    cfg.render_keys,
+                    sim_steps_per_control=cfg.sim_freq // cfg.control_freq,
+                )
+                rgba = alpha_blend_rgba(
+                    grasp_frames[camera.name][0]["Foreground"],
+                    bg_images[camera.name],
+                )
+                video_frames.append(np.array(rgba))
+
+            agent.reset(agent.init_qpos)
+
+    os.makedirs(cfg.output_dir, exist_ok=True)
+    video_path = f"{cfg.output_dir}/Iscene.mp4"
+    imageio.mimsave(video_path, video_frames, fps=30)
+    logger.info(f"Interative 3D Scene Visualization saved in {video_path}")
+
+
+if __name__ == "__main__":
+    entrypoint()

embodied_gen/scripts/text2image.py

@@ -0,0 +1,162 @@
+# 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.
+
+
+import argparse
+import logging
+import os
+
+from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256 import (
+    StableDiffusionXLPipeline,
+)
+from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256_ipadapter import (  # noqa
+    StableDiffusionXLPipeline as StableDiffusionXLPipelineIP,
+)
+from tqdm import tqdm
+from embodied_gen.models.text_model import (
+    build_text2img_ip_pipeline,
+    build_text2img_pipeline,
+    text2img_gen,
+)
+from embodied_gen.utils.process_media import parse_text_prompts
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Text to Image.")
+    parser.add_argument(
+        "--prompts",
+        type=str,
+        nargs="+",
+        help="List of prompts (space-separated).",
+    )
+    parser.add_argument(
+        "--ref_image",
+        type=str,
+        nargs="+",
+        help="List of ref_image paths (space-separated).",
+    )
+    parser.add_argument(
+        "--output_root",
+        type=str,
+        help="Root directory for saving outputs.",
+    )
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=12.0,
+        help="Guidance scale for the diffusion model.",
+    )
+    parser.add_argument(
+        "--ref_scale",
+        type=float,
+        default=0.3,
+        help="Reference image scale for the IP adapter.",
+    )
+    parser.add_argument(
+        "--n_sample",
+        type=int,
+        default=1,
+    )
+    parser.add_argument(
+        "--resolution",
+        type=int,
+        default=1024,
+    )
+    parser.add_argument(
+        "--infer_step",
+        type=int,
+        default=50,
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+    )
+    args = parser.parse_args()
+
+    return args
+
+
+def entrypoint(
+    pipeline: StableDiffusionXLPipeline | StableDiffusionXLPipelineIP = None,
+    **kwargs,
+) -> list[str]:
+    args = parse_args()
+    for k, v in kwargs.items():
+        if hasattr(args, k) and v is not None:
+            setattr(args, k, v)
+
+    prompts = parse_text_prompts(args.prompts)
+    os.makedirs(args.output_root, exist_ok=True)
+
+    ip_img_paths = args.ref_image
+    if ip_img_paths is None or len(ip_img_paths) == 0:
+        args.ref_scale = 0
+        ip_img_paths = [None] * len(prompts)
+    elif isinstance(ip_img_paths, str):
+        ip_img_paths = [ip_img_paths] * len(prompts)
+    elif isinstance(ip_img_paths, list):
+        if len(ip_img_paths) == 1:
+            ip_img_paths = ip_img_paths * len(prompts)
+    else:
+        raise ValueError("Invalid ref_image paths.")
+    assert len(ip_img_paths) == len(
+        prompts
+    ), f"Number of ref images does not match prompts, {len(ip_img_paths)} != {len(prompts)}"  # noqa
+
+    if pipeline is None:
+        if args.ref_scale > 0:
+            pipeline = build_text2img_ip_pipeline(
+                "weights/Kolors",
+                ref_scale=args.ref_scale,
+            )
+        else:
+            pipeline = build_text2img_pipeline("weights/Kolors")
+
+    for idx, (prompt, ip_img_path) in tqdm(
+        enumerate(zip(prompts, ip_img_paths)),
+        desc="Generating images",
+        total=len(prompts),
+    ):
+        images = text2img_gen(
+            prompt=prompt,
+            n_sample=args.n_sample,
+            guidance_scale=args.guidance_scale,
+            pipeline=pipeline,
+            ip_image=ip_img_path,
+            image_wh=[args.resolution, args.resolution],
+            infer_step=args.infer_step,
+            seed=args.seed,
+        )
+
+        save_paths = []
+        for sub_idx, image in enumerate(images):
+            save_path = (
+                f"{args.output_root}/sample_{idx*args.n_sample+sub_idx}.png"
+            )
+            image.save(save_path)
+            save_paths.append(save_path)
+
+        logger.info(f"Images saved to {args.output_root}")
+
+    return save_paths
+
+
+if __name__ == "__main__":
+    entrypoint()

embodied_gen/scripts/textto3d.py

@@ -0,0 +1,282 @@
+# 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.
+
+import argparse
+import os
+import random
+from collections import defaultdict
+
+import numpy as np
+import torch
+from PIL import Image
+from embodied_gen.models.image_comm_model import build_hf_image_pipeline
+from embodied_gen.models.segment_model import RembgRemover
+from embodied_gen.models.text_model import PROMPT_APPEND
+from embodied_gen.scripts.imageto3d import entrypoint as imageto3d_api
+from embodied_gen.utils.gpt_clients import GPT_CLIENT
+from embodied_gen.utils.log import logger
+from embodied_gen.utils.process_media import (
+    check_object_edge_truncated,
+    render_asset3d,
+)
+from embodied_gen.validators.quality_checkers import (
+    ImageSegChecker,
+    SemanticConsistChecker,
+    TextGenAlignChecker,
+)
+
+# Avoid huggingface/tokenizers: The current process just got forked.
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+random.seed(0)
+
+logger.info("Loading TEXT2IMG_MODEL...")
+SEMANTIC_CHECKER = SemanticConsistChecker(GPT_CLIENT)
+SEG_CHECKER = ImageSegChecker(GPT_CLIENT)
+TXTGEN_CHECKER = TextGenAlignChecker(GPT_CLIENT)
+PIPE_IMG = build_hf_image_pipeline(os.environ.get("TEXT_MODEL", "sd35"))
+BG_REMOVER = RembgRemover()
+
+
+__all__ = [
+    "text_to_image",
+    "text_to_3d",
+]
+
+
+def text_to_image(
+    prompt: str,
+    save_path: str,
+    n_retry: int,
+    img_denoise_step: int,
+    text_guidance_scale: float,
+    n_img_sample: int,
+    image_hw: tuple[int, int] = (1024, 1024),
+    seed: int = None,
+) -> bool:
+    select_image = None
+    success_flag = False
+    assert save_path.endswith(".png"), "Image save path must end with `.png`."
+    for try_idx in range(n_retry):
+        if select_image is not None:
+            select_image[0].save(save_path.replace(".png", "_raw.png"))
+            select_image[1].save(save_path)
+            break
+
+        f_prompt = PROMPT_APPEND.format(object=prompt)
+        logger.info(
+            f"Image GEN for {os.path.basename(save_path)}\n"
+            f"Try: {try_idx + 1}/{n_retry}, Seed: {seed}, Prompt: {f_prompt}"
+        )
+        torch.cuda.empty_cache()
+        images = PIPE_IMG.run(
+            f_prompt,
+            num_inference_steps=img_denoise_step,
+            guidance_scale=text_guidance_scale,
+            num_images_per_prompt=n_img_sample,
+            height=image_hw[0],
+            width=image_hw[1],
+            generator=(
+                torch.Generator().manual_seed(seed)
+                if seed is not None
+                else None
+            ),
+        )
+
+        for idx in range(len(images)):
+            raw_image: Image.Image = images[idx]
+            image = BG_REMOVER(raw_image)
+            image.save(save_path)
+            semantic_flag, semantic_result = SEMANTIC_CHECKER(
+                prompt, [image.convert("RGB")]
+            )
+            seg_flag, seg_result = SEG_CHECKER(
+                [raw_image, image.convert("RGB")]
+            )
+            image_mask = np.array(image)[..., -1]
+            edge_flag = check_object_edge_truncated(image_mask)
+            logger.warning(
+                f"SEMANTIC: {semantic_result}. SEG: {seg_result}. EDGE: {edge_flag}"
+            )
+            if (
+                (edge_flag and semantic_flag and seg_flag)
+                or (edge_flag and semantic_flag is None)
+                or (edge_flag and seg_flag is None)
+            ):
+                select_image = [raw_image, image]
+                success_flag = True
+                break
+
+        seed = random.randint(0, 100000) if seed is not None else None
+
+    return success_flag
+
+
+def text_to_3d(**kwargs) -> dict:
+    args = parse_args()
+    for k, v in kwargs.items():
+        if hasattr(args, k) and v is not None:
+            setattr(args, k, v)
+
+    if args.asset_names is None or len(args.asset_names) == 0:
+        args.asset_names = [f"sample3d_{i}" for i in range(len(args.prompts))]
+    img_save_dir = os.path.join(args.output_root, "images")
+    asset_save_dir = os.path.join(args.output_root, "asset3d")
+    os.makedirs(img_save_dir, exist_ok=True)
+    os.makedirs(asset_save_dir, exist_ok=True)
+    results = defaultdict(dict)
+    for prompt, node in zip(args.prompts, args.asset_names):
+        success_flag = False
+        n_pipe_retry = args.n_pipe_retry
+        seed_img = args.seed_img
+        seed_3d = args.seed_3d
+        while success_flag is False and n_pipe_retry > 0:
+            logger.info(
+                f"GEN pipeline for node {node}\n"
+                f"Try round: {args.n_pipe_retry-n_pipe_retry+1}/{args.n_pipe_retry}, Prompt: {prompt}"
+            )
+            # Text-to-image GEN
+            save_node = node.replace(" ", "_")
+            gen_image_path = f"{img_save_dir}/{save_node}.png"
+            textgen_flag = text_to_image(
+                prompt,
+                gen_image_path,
+                args.n_image_retry,
+                args.img_denoise_step,
+                args.text_guidance_scale,
+                args.n_img_sample,
+                seed=seed_img,
+            )
+
+            # Asset 3D GEN
+            node_save_dir = f"{asset_save_dir}/{save_node}"
+            asset_type = node if "sample3d_" not in node else None
+            imageto3d_api(
+                image_path=[gen_image_path],
+                output_root=node_save_dir,
+                asset_type=[asset_type],
+                seed=random.randint(0, 100000) if seed_3d is None else seed_3d,
+                n_retry=args.n_asset_retry,
+                keep_intermediate=args.keep_intermediate,
+                disable_decompose_convex=args.disable_decompose_convex,
+            )
+            mesh_path = f"{node_save_dir}/result/mesh/{save_node}.obj"
+            image_path = render_asset3d(
+                mesh_path,
+                output_root=f"{node_save_dir}/result",
+                num_images=6,
+                elevation=(30, -30),
+                output_subdir="renders",
+                no_index_file=True,
+            )
+
+            check_text = asset_type if asset_type is not None else prompt
+            qa_flag, qa_result = TXTGEN_CHECKER(check_text, image_path)
+            logger.warning(
+                f"Node {node}, {TXTGEN_CHECKER.__class__.__name__}: {qa_result}"
+            )
+            results["assets"][node] = f"asset3d/{save_node}/result"
+            results["quality"][node] = qa_result
+
+            if qa_flag is None or qa_flag is True:
+                success_flag = True
+                break
+
+            n_pipe_retry -= 1
+            seed_img = (
+                random.randint(0, 100000) if seed_img is not None else None
+            )
+            seed_3d = (
+                random.randint(0, 100000) if seed_3d is not None else None
+            )
+
+        torch.cuda.empty_cache()
+
+    return results
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="3D Layout Generation Config")
+    parser.add_argument("--prompts", nargs="+", help="text descriptions")
+    parser.add_argument(
+        "--output_root",
+        type=str,
+        help="Directory to save outputs",
+    )
+    parser.add_argument(
+        "--asset_names",
+        type=str,
+        nargs="+",
+        default=None,
+        help="Asset names to generate",
+    )
+    parser.add_argument(
+        "--n_img_sample",
+        type=int,
+        default=3,
+        help="Number of image samples to generate",
+    )
+    parser.add_argument(
+        "--text_guidance_scale",
+        type=float,
+        default=7,
+        help="Text-to-image guidance scale",
+    )
+    parser.add_argument(
+        "--img_denoise_step",
+        type=int,
+        default=25,
+        help="Denoising steps for image generation",
+    )
+    parser.add_argument(
+        "--n_image_retry",
+        type=int,
+        default=2,
+        help="Max retry count for image generation",
+    )
+    parser.add_argument(
+        "--n_asset_retry",
+        type=int,
+        default=2,
+        help="Max retry count for 3D generation",
+    )
+    parser.add_argument(
+        "--n_pipe_retry",
+        type=int,
+        default=1,
+        help="Max retry count for 3D asset generation",
+    )
+    parser.add_argument(
+        "--seed_img",
+        type=int,
+        default=None,
+        help="Random seed for image generation",
+    )
+    parser.add_argument(
+        "--seed_3d",
+        type=int,
+        default=0,
+        help="Random seed for 3D generation",
+    )
+    parser.add_argument("--keep_intermediate", action="store_true")
+    parser.add_argument("--disable_decompose_convex", action="store_true")
+
+    args, unknown = parser.parse_known_args()
+
+    return args
+
+
+if __name__ == "__main__":
+    text_to_3d()

embodied_gen/scripts/textto3d.sh

@@ -0,0 +1,94 @@
+#!/bin/bash
+
+# Initialize variables
+prompts=()
+asset_types=()
+output_root=""
+seed=0
+
+# Parse arguments
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --prompts)
+            shift
+            while [[ $# -gt 0 && ! "$1" =~ ^-- ]]; do
+                prompts+=("$1")
+                shift
+            done
+            ;;
+        --asset_types)
+            shift
+            while [[ $# -gt 0 && ! "$1" =~ ^-- ]]; do
+                asset_types+=("$1")
+                shift
+            done
+            ;;
+        --output_root)
+            output_root="$2"
+            shift 2
+            ;;
+        --seed)
+            seed="$2"
+            shift 2
+            ;;
+        *)
+            echo "Unknown argument: $1"
+            exit 1
+            ;;
+    esac
+done
+
+# Validate required arguments
+if [[ ${#prompts[@]} -eq 0 || -z "$output_root" ]]; then
+    echo "Missing required arguments."
+    echo "Usage: bash run_text2asset3d.sh --prompts \"Prompt1\" \"Prompt2\" \
+    --asset_types \"type1\" \"type2\" --seed <seed_value> --output_root <path>"
+    exit 1
+fi
+
+# If no asset_types provided, default to ""
+if [[ ${#asset_types[@]} -eq 0 ]]; then
+    for (( i=0; i<${#prompts[@]}; i++ )); do
+        asset_types+=("")
+    done
+fi
+
+# Ensure the number of asset_types matches the number of prompts
+if [[ ${#prompts[@]} -ne ${#asset_types[@]} ]]; then
+    echo "The number of asset types must match the number of prompts."
+    exit 1
+fi
+
+# Print arguments (for debugging)
+echo "Prompts:"
+for p in "${prompts[@]}"; do
+    echo "   - $p"
+done
+# echo "Asset types:"
+# for at in "${asset_types[@]}"; do
+#     echo "   - $at"
+# done
+echo "Output root: ${output_root}"
+echo "Seed: ${seed}"
+
+# Concatenate prompts and asset types for Python command
+prompt_args=""
+asset_type_args=""
+for i in "${!prompts[@]}"; do
+    prompt_args+="\"${prompts[$i]}\" "
+    asset_type_args+="\"${asset_types[$i]}\" "
+done
+
+
+# Step 1: Text-to-Image
+echo ${prompt_args}
+eval python3 embodied_gen/scripts/text2image.py \
+    --prompts ${prompt_args} \
+    --output_root "${output_root}/images" \
+    --seed ${seed}
+
+# Step 2: Image-to-3D
+python3 embodied_gen/scripts/imageto3d.py \
+    --image_root "${output_root}/images" \
+    --output_root "${output_root}/asset3d" \
+    --asset_type ${asset_type_args}

embodied_gen/scripts/texture_gen.sh

@@ -0,0 +1,78 @@
+#!/bin/bash
+
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --mesh_path)
+            mesh_path="$2"
+            shift 2
+            ;;
+        --prompt)
+            prompt="$2"
+            shift 2
+            ;;
+        --output_root)
+            output_root="$2"
+            shift 2
+            ;;
+        *)
+            echo "unknown: $1"
+            exit 1
+            ;;
+    esac
+done
+
+
+if [[ -z "$mesh_path" || -z "$prompt" || -z "$output_root" ]]; then
+    echo "params missing"
+    echo "usage: bash run.sh --mesh_path <path> --prompt <text> --output_root <path>"
+    exit 1
+fi
+
+echo "Will be deprecated, recommended to use 'texture-cli' instead."
+uuid=$(basename "$output_root")
+# Step 1: drender-cli for condition rendering
+drender-cli --mesh_path ${mesh_path} \
+    --output_root ${output_root}/condition \
+    --uuid ${uuid}
+
+# Step 2: multi-view rendering
+python embodied_gen/scripts/render_mv.py \
+    --index_file "${output_root}/condition/index.json" \
+    --controlnet_cond_scale 0.7 \
+    --guidance_scale 9 \
+    --strength 0.9 \
+    --num_inference_steps 40 \
+    --ip_adapt_scale 0 \
+    --ip_img_path None \
+    --uid ${uuid} \
+    --prompt "${prompt}" \
+    --save_dir "${output_root}/multi_view" \
+    --sub_idxs "[[0,1,2],[3,4,5]]" \
+    --seed 0
+
+# Step 3: backprojection
+backproject-cli --mesh_path ${mesh_path} \
+    --color_path ${output_root}/multi_view/color_sample0.png \
+    --output_path "${output_root}/texture_mesh/${uuid}.obj" \
+    --save_glb_path "${output_root}/texture_mesh/${uuid}.glb" \
+    --skip_fix_mesh \
+    --delight \
+    --no_save_delight_img
+
+# Step 4: final rendering of textured mesh
+drender-cli --mesh_path "${output_root}/texture_mesh/${uuid}.obj" \
+    --output_root ${output_root}/texture_mesh \
+    --num_images 90 \
+    --elevation 20 \
+    --with_mtl \
+    --gen_color_mp4 \
+    --pbr_light_factor 1.2
+
+# Organize folders
+rm -rf ${output_root}/condition
+video_path="${output_root}/texture_mesh/${uuid}/color.mp4"
+if [ -f "${video_path}" ]; then
+    cp "${video_path}" "${output_root}/texture_mesh/color.mp4"
+    echo "Resave video to ${output_root}/texture_mesh/color.mp4"
+fi
+rm -rf ${output_root}/texture_mesh/${uuid}

embodied_gen/trainer/gsplat_trainer.py

@@ -0,0 +1,678 @@
+# 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.
+# Part of the code comes from https://github.com/nerfstudio-project/gsplat/blob/main/examples/simple_trainer.py
+# Both under the Apache License, Version 2.0.
+
+
+import json
+import os
+import time
+from collections import defaultdict
+from typing import Dict, Optional, Tuple
+
+import cv2
+import imageio
+import numpy as np
+import torch
+import torch.nn.functional as F
+import tqdm
+import tyro
+import yaml
+from fused_ssim import fused_ssim
+from gsplat.distributed import cli
+from gsplat.rendering import rasterization
+from gsplat.strategy import DefaultStrategy, MCMCStrategy
+from torch import Tensor
+from torch.utils.tensorboard import SummaryWriter
+from torchmetrics.image import (
+    PeakSignalNoiseRatio,
+    StructuralSimilarityIndexMeasure,
+)
+from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from typing_extensions import Literal, assert_never
+from embodied_gen.data.datasets import PanoGSplatDataset
+from embodied_gen.utils.config import GsplatTrainConfig
+from embodied_gen.utils.gaussian import (
+    create_splats_with_optimizers,
+    export_splats,
+    resize_pinhole_intrinsics,
+    set_random_seed,
+)
+
+
+class Runner:
+    """Engine for training and testing from gsplat example.
+
+    Code from https://github.com/nerfstudio-project/gsplat/blob/main/examples/simple_trainer.py
+    """
+
+    def __init__(
+        self,
+        local_rank: int,
+        world_rank,
+        world_size: int,
+        cfg: GsplatTrainConfig,
+    ) -> None:
+        set_random_seed(42 + local_rank)
+
+        self.cfg = cfg
+        self.world_rank = world_rank
+        self.local_rank = local_rank
+        self.world_size = world_size
+        self.device = f"cuda:{local_rank}"
+
+        # Where to dump results.
+        os.makedirs(cfg.result_dir, exist_ok=True)
+
+        # Setup output directories.
+        self.ckpt_dir = f"{cfg.result_dir}/ckpts"
+        os.makedirs(self.ckpt_dir, exist_ok=True)
+        self.stats_dir = f"{cfg.result_dir}/stats"
+        os.makedirs(self.stats_dir, exist_ok=True)
+        self.render_dir = f"{cfg.result_dir}/renders"
+        os.makedirs(self.render_dir, exist_ok=True)
+        self.ply_dir = f"{cfg.result_dir}/ply"
+        os.makedirs(self.ply_dir, exist_ok=True)
+
+        # Tensorboard
+        self.writer = SummaryWriter(log_dir=f"{cfg.result_dir}/tb")
+        self.trainset = PanoGSplatDataset(cfg.data_dir, split="train")
+        self.valset = PanoGSplatDataset(
+            cfg.data_dir, split="train", max_sample_num=6
+        )
+        self.testset = PanoGSplatDataset(cfg.data_dir, split="eval")
+        self.scene_scale = cfg.scene_scale
+
+        # Model
+        self.splats, self.optimizers = create_splats_with_optimizers(
+            self.trainset.points,
+            self.trainset.points_rgb,
+            init_num_pts=cfg.init_num_pts,
+            init_extent=cfg.init_extent,
+            init_opacity=cfg.init_opa,
+            init_scale=cfg.init_scale,
+            means_lr=cfg.means_lr,
+            scales_lr=cfg.scales_lr,
+            opacities_lr=cfg.opacities_lr,
+            quats_lr=cfg.quats_lr,
+            sh0_lr=cfg.sh0_lr,
+            shN_lr=cfg.shN_lr,
+            scene_scale=self.scene_scale,
+            sh_degree=cfg.sh_degree,
+            sparse_grad=cfg.sparse_grad,
+            visible_adam=cfg.visible_adam,
+            batch_size=cfg.batch_size,
+            feature_dim=None,
+            device=self.device,
+            world_rank=world_rank,
+            world_size=world_size,
+        )
+        print("Model initialized. Number of GS:", len(self.splats["means"]))
+
+        # Densification Strategy
+        self.cfg.strategy.check_sanity(self.splats, self.optimizers)
+
+        if isinstance(self.cfg.strategy, DefaultStrategy):
+            self.strategy_state = self.cfg.strategy.initialize_state(
+                scene_scale=self.scene_scale
+            )
+        elif isinstance(self.cfg.strategy, MCMCStrategy):
+            self.strategy_state = self.cfg.strategy.initialize_state()
+        else:
+            assert_never(self.cfg.strategy)
+
+        # Losses & Metrics.
+        self.ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to(
+            self.device
+        )
+        self.psnr = PeakSignalNoiseRatio(data_range=1.0).to(self.device)
+
+        if cfg.lpips_net == "alex":
+            self.lpips = LearnedPerceptualImagePatchSimilarity(
+                net_type="alex", normalize=True
+            ).to(self.device)
+        elif cfg.lpips_net == "vgg":
+            # The 3DGS official repo uses lpips vgg, which is equivalent with the following:
+            self.lpips = LearnedPerceptualImagePatchSimilarity(
+                net_type="vgg", normalize=False
+            ).to(self.device)
+        else:
+            raise ValueError(f"Unknown LPIPS network: {cfg.lpips_net}")
+
+    def rasterize_splats(
+        self,
+        camtoworlds: Tensor,
+        Ks: Tensor,
+        width: int,
+        height: int,
+        masks: Optional[Tensor] = None,
+        rasterize_mode: Optional[Literal["classic", "antialiased"]] = None,
+        camera_model: Optional[Literal["pinhole", "ortho", "fisheye"]] = None,
+        **kwargs,
+    ) -> Tuple[Tensor, Tensor, Dict]:
+        means = self.splats["means"]  # [N, 3]
+        # quats = F.normalize(self.splats["quats"], dim=-1)  # [N, 4]
+        # rasterization does normalization internally
+        quats = self.splats["quats"]  # [N, 4]
+        scales = torch.exp(self.splats["scales"])  # [N, 3]
+        opacities = torch.sigmoid(self.splats["opacities"])  # [N,]
+        image_ids = kwargs.pop("image_ids", None)
+
+        colors = torch.cat(
+            [self.splats["sh0"], self.splats["shN"]], 1
+        )  # [N, K, 3]
+
+        if rasterize_mode is None:
+            rasterize_mode = (
+                "antialiased" if self.cfg.antialiased else "classic"
+            )
+        if camera_model is None:
+            camera_model = self.cfg.camera_model
+
+        render_colors, render_alphas, info = rasterization(
+            means=means,
+            quats=quats,
+            scales=scales,
+            opacities=opacities,
+            colors=colors,
+            viewmats=torch.linalg.inv(camtoworlds),  # [C, 4, 4]
+            Ks=Ks,  # [C, 3, 3]
+            width=width,
+            height=height,
+            packed=self.cfg.packed,
+            absgrad=(
+                self.cfg.strategy.absgrad
+                if isinstance(self.cfg.strategy, DefaultStrategy)
+                else False
+            ),
+            sparse_grad=self.cfg.sparse_grad,
+            rasterize_mode=rasterize_mode,
+            distributed=self.world_size > 1,
+            camera_model=self.cfg.camera_model,
+            with_ut=self.cfg.with_ut,
+            with_eval3d=self.cfg.with_eval3d,
+            **kwargs,
+        )
+        if masks is not None:
+            render_colors[~masks] = 0
+        return render_colors, render_alphas, info
+
+    def train(self):
+        cfg = self.cfg
+        device = self.device
+        world_rank = self.world_rank
+
+        # Dump cfg.
+        if world_rank == 0:
+            with open(f"{cfg.result_dir}/cfg.yml", "w") as f:
+                yaml.dump(vars(cfg), f)
+
+        max_steps = cfg.max_steps
+        init_step = 0
+
+        schedulers = [
+            # means has a learning rate schedule, that end at 0.01 of the initial value
+            torch.optim.lr_scheduler.ExponentialLR(
+                self.optimizers["means"], gamma=0.01 ** (1.0 / max_steps)
+            ),
+        ]
+        trainloader = torch.utils.data.DataLoader(
+            self.trainset,
+            batch_size=cfg.batch_size,
+            shuffle=True,
+            num_workers=4,
+            persistent_workers=True,
+            pin_memory=True,
+        )
+        trainloader_iter = iter(trainloader)
+
+        # Training loop.
+        global_tic = time.time()
+        pbar = tqdm.tqdm(range(init_step, max_steps))
+        for step in pbar:
+            try:
+                data = next(trainloader_iter)
+            except StopIteration:
+                trainloader_iter = iter(trainloader)
+                data = next(trainloader_iter)
+
+            camtoworlds = data["camtoworld"].to(device)  # [1, 4, 4]
+            Ks = data["K"].to(device)  # [1, 3, 3]
+            pixels = data["image"].to(device) / 255.0  # [1, H, W, 3]
+            image_ids = data["image_id"].to(device)
+            masks = (
+                data["mask"].to(device) if "mask" in data else None
+            )  # [1, H, W]
+            if cfg.depth_loss:
+                points = data["points"].to(device)  # [1, M, 2]
+                depths_gt = data["depths"].to(device)  # [1, M]
+
+            height, width = pixels.shape[1:3]
+
+            # sh schedule
+            sh_degree_to_use = min(
+                step // cfg.sh_degree_interval, cfg.sh_degree
+            )
+
+            # forward
+            renders, alphas, info = self.rasterize_splats(
+                camtoworlds=camtoworlds,
+                Ks=Ks,
+                width=width,
+                height=height,
+                sh_degree=sh_degree_to_use,
+                near_plane=cfg.near_plane,
+                far_plane=cfg.far_plane,
+                image_ids=image_ids,
+                render_mode="RGB+ED" if cfg.depth_loss else "RGB",
+                masks=masks,
+            )
+            if renders.shape[-1] == 4:
+                colors, depths = renders[..., 0:3], renders[..., 3:4]
+            else:
+                colors, depths = renders, None
+
+            if cfg.random_bkgd:
+                bkgd = torch.rand(1, 3, device=device)
+                colors = colors + bkgd * (1.0 - alphas)
+
+            self.cfg.strategy.step_pre_backward(
+                params=self.splats,
+                optimizers=self.optimizers,
+                state=self.strategy_state,
+                step=step,
+                info=info,
+            )
+
+            # loss
+            l1loss = F.l1_loss(colors, pixels)
+            ssimloss = 1.0 - fused_ssim(
+                colors.permute(0, 3, 1, 2),
+                pixels.permute(0, 3, 1, 2),
+                padding="valid",
+            )
+            loss = (
+                l1loss * (1.0 - cfg.ssim_lambda) + ssimloss * cfg.ssim_lambda
+            )
+            if cfg.depth_loss:
+                # query depths from depth map
+                points = torch.stack(
+                    [
+                        points[:, :, 0] / (width - 1) * 2 - 1,
+                        points[:, :, 1] / (height - 1) * 2 - 1,
+                    ],
+                    dim=-1,
+                )  # normalize to [-1, 1]
+                grid = points.unsqueeze(2)  # [1, M, 1, 2]
+                depths = F.grid_sample(
+                    depths.permute(0, 3, 1, 2), grid, align_corners=True
+                )  # [1, 1, M, 1]
+                depths = depths.squeeze(3).squeeze(1)  # [1, M]
+                # calculate loss in disparity space
+                disp = torch.where(
+                    depths > 0.0, 1.0 / depths, torch.zeros_like(depths)
+                )
+                disp_gt = 1.0 / depths_gt  # [1, M]
+                depthloss = F.l1_loss(disp, disp_gt) * self.scene_scale
+                loss += depthloss * cfg.depth_lambda
+
+            # regularizations
+            if cfg.opacity_reg > 0.0:
+                loss += (
+                    cfg.opacity_reg
+                    * torch.sigmoid(self.splats["opacities"]).mean()
+                )
+            if cfg.scale_reg > 0.0:
+                loss += cfg.scale_reg * torch.exp(self.splats["scales"]).mean()
+
+            loss.backward()
+
+            desc = (
+                f"loss={loss.item():.3f}| " f"sh degree={sh_degree_to_use}| "
+            )
+            if cfg.depth_loss:
+                desc += f"depth loss={depthloss.item():.6f}| "
+            pbar.set_description(desc)
+
+            # write images (gt and render)
+            # if world_rank == 0 and step % 800 == 0:
+            #     canvas = torch.cat([pixels, colors], dim=2).detach().cpu().numpy()
+            #     canvas = canvas.reshape(-1, *canvas.shape[2:])
+            #     imageio.imwrite(
+            #         f"{self.render_dir}/train_rank{self.world_rank}.png",
+            #         (canvas * 255).astype(np.uint8),
+            #     )
+
+            if (
+                world_rank == 0
+                and cfg.tb_every > 0
+                and step % cfg.tb_every == 0
+            ):
+                mem = torch.cuda.max_memory_allocated() / 1024**3
+                self.writer.add_scalar("train/loss", loss.item(), step)
+                self.writer.add_scalar("train/l1loss", l1loss.item(), step)
+                self.writer.add_scalar("train/ssimloss", ssimloss.item(), step)
+                self.writer.add_scalar(
+                    "train/num_GS", len(self.splats["means"]), step
+                )
+                self.writer.add_scalar("train/mem", mem, step)
+                if cfg.depth_loss:
+                    self.writer.add_scalar(
+                        "train/depthloss", depthloss.item(), step
+                    )
+                if cfg.tb_save_image:
+                    canvas = (
+                        torch.cat([pixels, colors], dim=2)
+                        .detach()
+                        .cpu()
+                        .numpy()
+                    )
+                    canvas = canvas.reshape(-1, *canvas.shape[2:])
+                    self.writer.add_image("train/render", canvas, step)
+                self.writer.flush()
+
+            # save checkpoint before updating the model
+            if (
+                step in [i - 1 for i in cfg.save_steps]
+                or step == max_steps - 1
+            ):
+                mem = torch.cuda.max_memory_allocated() / 1024**3
+                stats = {
+                    "mem": mem,
+                    "ellipse_time": time.time() - global_tic,
+                    "num_GS": len(self.splats["means"]),
+                }
+                print("Step: ", step, stats)
+                with open(
+                    f"{self.stats_dir}/train_step{step:04d}_rank{self.world_rank}.json",
+                    "w",
+                ) as f:
+                    json.dump(stats, f)
+                data = {"step": step, "splats": self.splats.state_dict()}
+                torch.save(
+                    data,
+                    f"{self.ckpt_dir}/ckpt_{step}_rank{self.world_rank}.pt",
+                )
+            if (
+                step in [i - 1 for i in cfg.ply_steps] or step == max_steps - 1
+            ) and cfg.save_ply:
+                sh0 = self.splats["sh0"]
+                shN = self.splats["shN"]
+                means = self.splats["means"]
+                scales = self.splats["scales"]
+                quats = self.splats["quats"]
+                opacities = self.splats["opacities"]
+                export_splats(
+                    means=means,
+                    scales=scales,
+                    quats=quats,
+                    opacities=opacities,
+                    sh0=sh0,
+                    shN=shN,
+                    format="ply",
+                    save_to=f"{self.ply_dir}/point_cloud_{step}.ply",
+                )
+
+            # Turn Gradients into Sparse Tensor before running optimizer
+            if cfg.sparse_grad:
+                assert (
+                    cfg.packed
+                ), "Sparse gradients only work with packed mode."
+                gaussian_ids = info["gaussian_ids"]
+                for k in self.splats.keys():
+                    grad = self.splats[k].grad
+                    if grad is None or grad.is_sparse:
+                        continue
+                    self.splats[k].grad = torch.sparse_coo_tensor(
+                        indices=gaussian_ids[None],  # [1, nnz]
+                        values=grad[gaussian_ids],  # [nnz, ...]
+                        size=self.splats[k].size(),  # [N, ...]
+                        is_coalesced=len(Ks) == 1,
+                    )
+
+            if cfg.visible_adam:
+                gaussian_cnt = self.splats.means.shape[0]
+                if cfg.packed:
+                    visibility_mask = torch.zeros_like(
+                        self.splats["opacities"], dtype=bool
+                    )
+                    visibility_mask.scatter_(0, info["gaussian_ids"], 1)
+                else:
+                    visibility_mask = (info["radii"] > 0).all(-1).any(0)
+
+            # optimize
+            for optimizer in self.optimizers.values():
+                if cfg.visible_adam:
+                    optimizer.step(visibility_mask)
+                else:
+                    optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+            for scheduler in schedulers:
+                scheduler.step()
+
+            # Run post-backward steps after backward and optimizer
+            if isinstance(self.cfg.strategy, DefaultStrategy):
+                self.cfg.strategy.step_post_backward(
+                    params=self.splats,
+                    optimizers=self.optimizers,
+                    state=self.strategy_state,
+                    step=step,
+                    info=info,
+                    packed=cfg.packed,
+                )
+            elif isinstance(self.cfg.strategy, MCMCStrategy):
+                self.cfg.strategy.step_post_backward(
+                    params=self.splats,
+                    optimizers=self.optimizers,
+                    state=self.strategy_state,
+                    step=step,
+                    info=info,
+                    lr=schedulers[0].get_last_lr()[0],
+                )
+            else:
+                assert_never(self.cfg.strategy)
+
+            # eval the full set
+            if step in [i - 1 for i in cfg.eval_steps]:
+                self.eval(step)
+                self.render_video(step)
+
+    @torch.no_grad()
+    def eval(
+        self,
+        step: int,
+        stage: str = "val",
+        canvas_h: int = 512,
+        canvas_w: int = 1024,
+    ):
+        """Entry for evaluation."""
+        print("Running evaluation...")
+        cfg = self.cfg
+        device = self.device
+        world_rank = self.world_rank
+
+        valloader = torch.utils.data.DataLoader(
+            self.valset, batch_size=1, shuffle=False, num_workers=1
+        )
+        ellipse_time = 0
+        metrics = defaultdict(list)
+        for i, data in enumerate(valloader):
+            camtoworlds = data["camtoworld"].to(device)
+            Ks = data["K"].to(device)
+            pixels = data["image"].to(device) / 255.0
+            height, width = pixels.shape[1:3]
+            masks = data["mask"].to(device) if "mask" in data else None
+
+            pixels = pixels.permute(0, 3, 1, 2)  # NHWC -> NCHW
+            pixels = F.interpolate(pixels, size=(canvas_h, canvas_w // 2))
+
+            torch.cuda.synchronize()
+            tic = time.time()
+            colors, _, _ = self.rasterize_splats(
+                camtoworlds=camtoworlds,
+                Ks=Ks,
+                width=width,
+                height=height,
+                sh_degree=cfg.sh_degree,
+                near_plane=cfg.near_plane,
+                far_plane=cfg.far_plane,
+                masks=masks,
+            )  # [1, H, W, 3]
+            torch.cuda.synchronize()
+            ellipse_time += max(time.time() - tic, 1e-10)
+
+            colors = colors.permute(0, 3, 1, 2)  # NHWC -> NCHW
+            colors = F.interpolate(colors, size=(canvas_h, canvas_w // 2))
+            colors = torch.clamp(colors, 0.0, 1.0)
+            canvas_list = [pixels, colors]
+
+            if world_rank == 0:
+                canvas = torch.cat(canvas_list, dim=2).squeeze(0)
+                canvas = canvas.permute(1, 2, 0)  # CHW -> HWC
+                canvas = (canvas * 255).to(torch.uint8).cpu().numpy()
+                cv2.imwrite(
+                    f"{self.render_dir}/{stage}_step{step}_{i:04d}.png",
+                    canvas[..., ::-1],
+                )
+                metrics["psnr"].append(self.psnr(colors, pixels))
+                metrics["ssim"].append(self.ssim(colors, pixels))
+                metrics["lpips"].append(self.lpips(colors, pixels))
+
+        if world_rank == 0:
+            ellipse_time /= len(valloader)
+
+            stats = {
+                k: torch.stack(v).mean().item() for k, v in metrics.items()
+            }
+            stats.update(
+                {
+                    "ellipse_time": ellipse_time,
+                    "num_GS": len(self.splats["means"]),
+                }
+            )
+            print(
+                f"PSNR: {stats['psnr']:.3f}, SSIM: {stats['ssim']:.4f}, LPIPS: {stats['lpips']:.3f} "
+                f"Time: {stats['ellipse_time']:.3f}s/image "
+                f"Number of GS: {stats['num_GS']}"
+            )
+            # save stats as json
+            with open(
+                f"{self.stats_dir}/{stage}_step{step:04d}.json", "w"
+            ) as f:
+                json.dump(stats, f)
+            # save stats to tensorboard
+            for k, v in stats.items():
+                self.writer.add_scalar(f"{stage}/{k}", v, step)
+            self.writer.flush()
+
+    @torch.no_grad()
+    def render_video(
+        self, step: int, canvas_h: int = 512, canvas_w: int = 1024
+    ):
+        testloader = torch.utils.data.DataLoader(
+            self.testset, batch_size=1, shuffle=False, num_workers=1
+        )
+
+        images_cache = []
+        depth_global_min, depth_global_max = float("inf"), -float("inf")
+        for data in testloader:
+            camtoworlds = data["camtoworld"].to(self.device)
+            Ks = resize_pinhole_intrinsics(
+                data["K"].squeeze(),
+                raw_hw=(data["image_h"].item(), data["image_w"].item()),
+                new_hw=(canvas_h, canvas_w // 2),
+            ).to(self.device)
+            renders, _, _ = self.rasterize_splats(
+                camtoworlds=camtoworlds,
+                Ks=Ks[None, ...],
+                width=canvas_w // 2,
+                height=canvas_h,
+                sh_degree=self.cfg.sh_degree,
+                near_plane=self.cfg.near_plane,
+                far_plane=self.cfg.far_plane,
+                render_mode="RGB+ED",
+            )  # [1, H, W, 4]
+            colors = torch.clamp(renders[0, ..., 0:3], 0.0, 1.0)  # [H, W, 3]
+            colors = (colors * 255).to(torch.uint8).cpu().numpy()
+            depths = renders[0, ..., 3:4]  # [H, W, 1], tensor in device.
+            images_cache.append([colors, depths])
+            depth_global_min = min(depth_global_min, depths.min().item())
+            depth_global_max = max(depth_global_max, depths.max().item())
+
+        video_path = f"{self.render_dir}/video_step{step}.mp4"
+        writer = imageio.get_writer(video_path, fps=30)
+        for rgb, depth in images_cache:
+            depth_normalized = torch.clip(
+                (depth - depth_global_min)
+                / (depth_global_max - depth_global_min + 1e-8),
+                0,
+                1,
+            )
+            depth_normalized = (
+                (depth_normalized * 255).to(torch.uint8).cpu().numpy()
+            )
+            depth_map = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_JET)
+            image = np.concatenate([rgb, depth_map], axis=1)
+            writer.append_data(image)
+
+        writer.close()
+
+
+def entrypoint(
+    local_rank: int, world_rank, world_size: int, cfg: GsplatTrainConfig
+):
+    runner = Runner(local_rank, world_rank, world_size, cfg)
+
+    if cfg.ckpt is not None:
+        # run eval only
+        ckpts = [
+            torch.load(file, map_location=runner.device, weights_only=True)
+            for file in cfg.ckpt
+        ]
+        for k in runner.splats.keys():
+            runner.splats[k].data = torch.cat(
+                [ckpt["splats"][k] for ckpt in ckpts]
+            )
+        step = ckpts[0]["step"]
+        runner.eval(step=step)
+        runner.render_video(step=step)
+    else:
+        runner.train()
+        runner.render_video(step=cfg.max_steps - 1)
+
+
+if __name__ == "__main__":
+    configs = {
+        "default": (
+            "Gaussian splatting training using densification heuristics from the original paper.",
+            GsplatTrainConfig(
+                strategy=DefaultStrategy(verbose=True),
+            ),
+        ),
+        "mcmc": (
+            "Gaussian splatting training using densification from the paper '3D Gaussian Splatting as Markov Chain Monte Carlo'.",
+            GsplatTrainConfig(
+                init_scale=0.1,
+                opacity_reg=0.01,
+                scale_reg=0.01,
+                strategy=MCMCStrategy(verbose=True),
+            ),
+        ),
+    }
+    cfg = tyro.extras.overridable_config_cli(configs)
+    cfg.adjust_steps(cfg.steps_scaler)
+
+    cli(entrypoint, cfg, verbose=True)

embodied_gen/trainer/pono2mesh_trainer.py

@@ -0,0 +1,538 @@
+# 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_pano2room
+
+monkey_patch_pano2room()
+
+import os
+
+import cv2
+import numpy as np
+import torch
+import trimesh
+from equilib import cube2equi, equi2pers
+from kornia.morphology import dilation
+from PIL import Image
+from embodied_gen.models.sr_model import ImageRealESRGAN
+from embodied_gen.utils.config import Pano2MeshSRConfig
+from embodied_gen.utils.geometry import compute_pinhole_intrinsics
+from embodied_gen.utils.log import logger
+from thirdparty.pano2room.modules.geo_predictors import PanoJointPredictor
+from thirdparty.pano2room.modules.geo_predictors.PanoFusionDistancePredictor import (
+    PanoFusionDistancePredictor,
+)
+from thirdparty.pano2room.modules.inpainters import PanoPersFusionInpainter
+from thirdparty.pano2room.modules.mesh_fusion.render import (
+    features_to_world_space_mesh,
+    render_mesh,
+)
+from thirdparty.pano2room.modules.mesh_fusion.sup_info import SupInfoPool
+from thirdparty.pano2room.utils.camera_utils import gen_pano_rays
+from thirdparty.pano2room.utils.functions import (
+    depth_to_distance,
+    get_cubemap_views_world_to_cam,
+    resize_image_with_aspect_ratio,
+    rot_z_world_to_cam,
+    tensor_to_pil,
+)
+
+
+class Pano2MeshSRPipeline:
+    """Converting panoramic RGB image into 3D mesh representations, followed by inpainting and mesh refinement.
+
+    This class integrates several key components including:
+    - Depth estimation from RGB panorama
+    - Inpainting of missing regions under offsets
+    - RGB-D to mesh conversion
+    - Multi-view mesh repair
+    - 3D Gaussian Splatting (3DGS) dataset generation
+
+    Args:
+        config (Pano2MeshSRConfig): Configuration object containing model and pipeline parameters.
+
+    Example:
+        ```python
+        pipeline = Pano2MeshSRPipeline(config)
+        pipeline(pano_image='example.png', output_dir='./output')
+        ```
+    """
+
+    def __init__(self, config: Pano2MeshSRConfig) -> None:
+        self.cfg = config
+        self.device = config.device
+
+        # Init models.
+        self.inpainter = PanoPersFusionInpainter(save_path=None)
+        self.geo_predictor = PanoJointPredictor(save_path=None)
+        self.pano_fusion_distance_predictor = PanoFusionDistancePredictor()
+        self.super_model = ImageRealESRGAN(outscale=self.cfg.upscale_factor)
+
+        # Init poses.
+        cubemap_w2cs = get_cubemap_views_world_to_cam()
+        self.cubemap_w2cs = [p.to(self.device) for p in cubemap_w2cs]
+        self.camera_poses = self.load_camera_poses(self.cfg.trajectory_dir)
+
+        kernel = cv2.getStructuringElement(
+            cv2.MORPH_ELLIPSE, self.cfg.kernel_size
+        )
+        self.kernel = torch.from_numpy(kernel).float().to(self.device)
+
+    def init_mesh_params(self) -> None:
+        torch.set_default_device(self.device)
+        self.inpaint_mask = torch.ones(
+            (self.cfg.cubemap_h, self.cfg.cubemap_w), dtype=torch.bool
+        )
+        self.vertices = torch.empty((3, 0), requires_grad=False)
+        self.colors = torch.empty((3, 0), requires_grad=False)
+        self.faces = torch.empty((3, 0), dtype=torch.long, requires_grad=False)
+
+    @staticmethod
+    def read_camera_pose_file(filepath: str) -> np.ndarray:
+        with open(filepath, "r") as f:
+            values = [float(num) for line in f for num in line.split()]
+
+        return np.array(values).reshape(4, 4)
+
+    def load_camera_poses(
+        self, trajectory_dir: str
+    ) -> tuple[np.ndarray, list[torch.Tensor]]:
+        pose_filenames = sorted(
+            [
+                fname
+                for fname in os.listdir(trajectory_dir)
+                if fname.startswith("camera_pose")
+            ]
+        )
+
+        pano_pose_world = None
+        relative_poses = []
+        for idx, filename in enumerate(pose_filenames):
+            pose_path = os.path.join(trajectory_dir, filename)
+            pose_matrix = self.read_camera_pose_file(pose_path)
+
+            if pano_pose_world is None:
+                pano_pose_world = pose_matrix.copy()
+                pano_pose_world[0, 3] += self.cfg.pano_center_offset[0]
+                pano_pose_world[2, 3] += self.cfg.pano_center_offset[1]
+
+            # Use different reference for the first 6 cubemap views
+            reference_pose = pose_matrix if idx < 6 else pano_pose_world
+            relative_matrix = pose_matrix @ np.linalg.inv(reference_pose)
+            relative_matrix[0:2, :] *= -1  # flip_xy
+            relative_matrix = (
+                relative_matrix @ rot_z_world_to_cam(180).cpu().numpy()
+            )
+            relative_matrix[:3, 3] *= self.cfg.pose_scale
+            relative_matrix = torch.tensor(
+                relative_matrix, dtype=torch.float32
+            )
+            relative_poses.append(relative_matrix)
+
+        return relative_poses
+
+    def load_inpaint_poses(
+        self, poses: torch.Tensor
+    ) -> dict[int, torch.Tensor]:
+        inpaint_poses = dict()
+        sampled_views = poses[:: self.cfg.inpaint_frame_stride]
+        init_pose = torch.eye(4)
+        for idx, w2c_tensor in enumerate(sampled_views):
+            w2c = w2c_tensor.cpu().numpy().astype(np.float32)
+            c2w = np.linalg.inv(w2c)
+            pose_tensor = init_pose.clone()
+            pose_tensor[:3, 3] = torch.from_numpy(c2w[:3, 3])
+            pose_tensor[:3, 3] *= -1
+            inpaint_poses[idx] = pose_tensor.to(self.device)
+
+        return inpaint_poses
+
+    def project(self, world_to_cam: torch.Tensor):
+        (
+            project_image,
+            project_depth,
+            inpaint_mask,
+            _,
+            z_buf,
+            mesh,
+        ) = render_mesh(
+            vertices=self.vertices,
+            faces=self.faces,
+            vertex_features=self.colors,
+            H=self.cfg.cubemap_h,
+            W=self.cfg.cubemap_w,
+            fov_in_degrees=self.cfg.fov,
+            RT=world_to_cam,
+            blur_radius=self.cfg.blur_radius,
+            faces_per_pixel=self.cfg.faces_per_pixel,
+        )
+        project_image = project_image * ~inpaint_mask
+
+        return project_image[:3, ...], inpaint_mask, project_depth
+
+    def render_pano(self, pose: torch.Tensor):
+        cubemap_list = []
+        for cubemap_pose in self.cubemap_w2cs:
+            project_pose = cubemap_pose @ pose
+            rgb, inpaint_mask, depth = self.project(project_pose)
+            distance_map = depth_to_distance(depth[None, ...])
+            mask = inpaint_mask[None, ...]
+            cubemap_list.append(torch.cat([rgb, distance_map, mask], dim=0))
+
+        # Set default tensor type for CPU operation in cube2equi
+        with torch.device("cpu"):
+            pano_rgbd = cube2equi(
+                cubemap_list, "list", self.cfg.pano_h, self.cfg.pano_w
+            )
+
+        pano_rgb = pano_rgbd[:3, :, :]
+        pano_depth = pano_rgbd[3:4, :, :].squeeze(0)
+        pano_mask = pano_rgbd[4:, :, :].squeeze(0)
+
+        return pano_rgb, pano_depth, pano_mask
+
+    def rgbd_to_mesh(
+        self,
+        rgb: torch.Tensor,
+        depth: torch.Tensor,
+        inpaint_mask: torch.Tensor,
+        world_to_cam: torch.Tensor = None,
+        using_distance_map: bool = True,
+    ) -> None:
+        if world_to_cam is None:
+            world_to_cam = torch.eye(4, dtype=torch.float32).to(self.device)
+
+        if inpaint_mask.sum() == 0:
+            return
+
+        vertices, faces, colors = features_to_world_space_mesh(
+            colors=rgb.squeeze(0),
+            depth=depth,
+            fov_in_degrees=self.cfg.fov,
+            world_to_cam=world_to_cam,
+            mask=inpaint_mask,
+            faces=self.faces,
+            vertices=self.vertices,
+            using_distance_map=using_distance_map,
+            edge_threshold=0.05,
+        )
+
+        faces += self.vertices.shape[1]
+        self.vertices = torch.cat([self.vertices, vertices], dim=1)
+        self.colors = torch.cat([self.colors, colors], dim=1)
+        self.faces = torch.cat([self.faces, faces], dim=1)
+
+    def get_edge_image_by_depth(
+        self, depth: torch.Tensor, dilate_iter: int = 1
+    ) -> np.ndarray:
+        if isinstance(depth, torch.Tensor):
+            depth = depth.cpu().detach().numpy()
+
+        gray = (depth / depth.max() * 255).astype(np.uint8)
+        edges = cv2.Canny(gray, 60, 150)
+        if dilate_iter > 0:
+            kernel = np.ones((3, 3), np.uint8)
+            edges = cv2.dilate(edges, kernel, iterations=dilate_iter)
+
+        return edges
+
+    def mesh_repair_by_greedy_view_selection(
+        self, pose_dict: dict[str, torch.Tensor], output_dir: str
+    ) -> list:
+        inpainted_panos_w_pose = []
+        while len(pose_dict) > 0:
+            logger.info(f"Repairing mesh left rounds {len(pose_dict)}")
+            sampled_views = []
+            for key, pose in pose_dict.items():
+                pano_rgb, pano_distance, pano_mask = self.render_pano(pose)
+                completeness = torch.sum(1 - pano_mask) / (pano_mask.numel())
+                sampled_views.append((key, completeness.item(), pose))
+
+            if len(sampled_views) == 0:
+                break
+
+            # Find inpainting with least view completeness.
+            sampled_views = sorted(sampled_views, key=lambda x: x[1])
+            key, _, pose = sampled_views[len(sampled_views) * 2 // 3]
+            pose_dict.pop(key)
+
+            pano_rgb, pano_distance, pano_mask = self.render_pano(pose)
+
+            colors = pano_rgb.permute(1, 2, 0).clone()
+            distances = pano_distance.unsqueeze(-1).clone()
+            pano_inpaint_mask = pano_mask.clone()
+            init_pose = pose.clone()
+            normals = None
+            if pano_inpaint_mask.min().item() < 0.5:
+                colors, distances, normals = self.inpaint_panorama(
+                    idx=key,
+                    colors=colors,
+                    distances=distances,
+                    pano_mask=pano_inpaint_mask,
+                )
+
+                init_pose[0, 3], init_pose[1, 3], init_pose[2, 3] = (
+                    -pose[0, 3],
+                    pose[2, 3],
+                    0,
+                )
+                rays = gen_pano_rays(
+                    init_pose, self.cfg.pano_h, self.cfg.pano_w
+                )
+                conflict_mask = self.sup_pool.geo_check(
+                    rays, distances.unsqueeze(-1)
+                )  # 0 is conflict, 1 not conflict
+                pano_inpaint_mask *= conflict_mask
+
+            self.rgbd_to_mesh(
+                colors.permute(2, 0, 1),
+                distances,
+                pano_inpaint_mask,
+                world_to_cam=pose,
+            )
+
+            self.sup_pool.register_sup_info(
+                pose=init_pose,
+                mask=pano_inpaint_mask.clone(),
+                rgb=colors,
+                distance=distances.unsqueeze(-1),
+                normal=normals,
+            )
+
+            colors = colors.permute(2, 0, 1).unsqueeze(0)
+            inpainted_panos_w_pose.append([colors, pose])
+
+            if self.cfg.visualize:
+                from embodied_gen.data.utils import DiffrastRender
+
+                tensor_to_pil(pano_rgb.unsqueeze(0)).save(
+                    f"{output_dir}/rendered_pano_{key}.jpg"
+                )
+                tensor_to_pil(colors).save(
+                    f"{output_dir}/inpainted_pano_{key}.jpg"
+                )
+                norm_depth = DiffrastRender.normalize_map_by_mask(
+                    distances, torch.ones_like(distances)
+                )
+                heatmap = (norm_depth.cpu().numpy() * 255).astype(np.uint8)
+                heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+                Image.fromarray(heatmap).save(
+                    f"{output_dir}/inpainted_depth_{key}.png"
+                )
+
+        return inpainted_panos_w_pose
+
+    def inpaint_panorama(
+        self,
+        idx: int,
+        colors: torch.Tensor,
+        distances: torch.Tensor,
+        pano_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor]:
+        mask = (pano_mask[None, ..., None] > 0.5).float()
+        mask = mask.permute(0, 3, 1, 2)
+        mask = dilation(mask, kernel=self.kernel)
+        mask = mask[0, 0, ..., None]  # hwc
+        inpainted_img = self.inpainter.inpaint(idx, colors, mask)
+        inpainted_img = colors * (1 - mask) + inpainted_img * mask
+        inpainted_distances, inpainted_normals = self.geo_predictor(
+            idx,
+            inpainted_img,
+            distances[..., None],
+            mask=mask,
+            reg_loss_weight=0.0,
+            normal_loss_weight=5e-2,
+            normal_tv_loss_weight=5e-2,
+        )
+
+        return inpainted_img, inpainted_distances.squeeze(), inpainted_normals
+
+    def preprocess_pano(
+        self, image: Image.Image | str
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if isinstance(image, str):
+            image = Image.open(image)
+
+        image = image.convert("RGB")
+
+        if image.size[0] < image.size[1]:
+            image = image.transpose(Image.TRANSPOSE)
+
+        image = resize_image_with_aspect_ratio(image, self.cfg.pano_w)
+        image_rgb = torch.tensor(np.array(image)).permute(2, 0, 1) / 255
+        image_rgb = image_rgb.to(self.device)
+        image_depth = self.pano_fusion_distance_predictor.predict(
+            image_rgb.permute(1, 2, 0)
+        )
+        image_depth = (
+            image_depth / image_depth.max() * self.cfg.depth_scale_factor
+        )
+
+        return image_rgb, image_depth
+
+    def pano_to_perpective(
+        self, pano_image: torch.Tensor, pitch: float, yaw: float, fov: float
+    ) -> torch.Tensor:
+        rots = dict(
+            roll=0,
+            pitch=pitch,
+            yaw=yaw,
+        )
+        perspective = equi2pers(
+            equi=pano_image.squeeze(0),
+            rots=rots,
+            height=self.cfg.cubemap_h,
+            width=self.cfg.cubemap_w,
+            fov_x=fov,
+            mode="bilinear",
+        ).unsqueeze(0)
+
+        return perspective
+
+    def pano_to_cubemap(self, pano_rgb: torch.Tensor):
+        # Define six canonical cube directions in (pitch, yaw)
+        directions = [
+            (0, 0),
+            (0, 1.5 * np.pi),
+            (0, 1.0 * np.pi),
+            (0, 0.5 * np.pi),
+            (-0.5 * np.pi, 0),
+            (0.5 * np.pi, 0),
+        ]
+
+        cubemaps_rgb = []
+        for pitch, yaw in directions:
+            rgb_view = self.pano_to_perpective(
+                pano_rgb, pitch, yaw, fov=self.cfg.fov
+            )
+            cubemaps_rgb.append(rgb_view.cpu())
+
+        return cubemaps_rgb
+
+    def save_mesh(self, output_path: str) -> None:
+        vertices_np = self.vertices.T.cpu().numpy()
+        colors_np = self.colors.T.cpu().numpy()
+        faces_np = self.faces.T.cpu().numpy()
+        mesh = trimesh.Trimesh(
+            vertices=vertices_np, faces=faces_np, vertex_colors=colors_np
+        )
+
+        mesh.export(output_path)
+
+    def mesh_pose_to_gs_pose(self, mesh_pose: torch.Tensor) -> np.ndarray:
+        pose = mesh_pose.clone()
+        pose[0, :] *= -1
+        pose[1, :] *= -1
+
+        Rw2c = pose[:3, :3].cpu().numpy()
+        Tw2c = pose[:3, 3:].cpu().numpy()
+        yz_reverse = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
+
+        Rc2w = (yz_reverse @ Rw2c).T
+        Tc2w = -(Rc2w @ yz_reverse @ Tw2c)
+        c2w = np.concatenate((Rc2w, Tc2w), axis=1)
+        c2w = np.concatenate((c2w, np.array([[0, 0, 0, 1]])), axis=0)
+
+        return c2w
+
+    def __call__(self, pano_image: Image.Image | str, output_dir: str):
+        self.init_mesh_params()
+        pano_rgb, pano_depth = self.preprocess_pano(pano_image)
+        self.sup_pool = SupInfoPool()
+        self.sup_pool.register_sup_info(
+            pose=torch.eye(4).to(self.device),
+            mask=torch.ones([self.cfg.pano_h, self.cfg.pano_w]),
+            rgb=pano_rgb.permute(1, 2, 0),
+            distance=pano_depth[..., None],
+        )
+        self.sup_pool.gen_occ_grid(res=256)
+
+        logger.info("Init mesh from pano RGBD image...")
+        depth_edge = self.get_edge_image_by_depth(pano_depth)
+        inpaint_edge_mask = (
+            ~torch.from_numpy(depth_edge).to(self.device).bool()
+        )
+        self.rgbd_to_mesh(pano_rgb, pano_depth, inpaint_edge_mask)
+
+        repair_poses = self.load_inpaint_poses(self.camera_poses)
+        inpainted_panos_w_poses = self.mesh_repair_by_greedy_view_selection(
+            repair_poses, output_dir
+        )
+        torch.cuda.empty_cache()
+        torch.set_default_device("cpu")
+
+        if self.cfg.mesh_file is not None:
+            mesh_path = os.path.join(output_dir, self.cfg.mesh_file)
+            self.save_mesh(mesh_path)
+
+        if self.cfg.gs_data_file is None:
+            return
+
+        logger.info(f"Dump data for 3DGS training...")
+        points_rgb = (self.colors.clip(0, 1) * 255).to(torch.uint8)
+        data = {
+            "points": self.vertices.permute(1, 0).cpu().numpy(),  # (N, 3)
+            "points_rgb": points_rgb.permute(1, 0).cpu().numpy(),  # (N, 3)
+            "train": [],
+            "eval": [],
+        }
+        image_h = self.cfg.cubemap_h * self.cfg.upscale_factor
+        image_w = self.cfg.cubemap_w * self.cfg.upscale_factor
+        Ks = compute_pinhole_intrinsics(image_w, image_h, self.cfg.fov)
+        for idx, (pano_img, pano_pose) in enumerate(inpainted_panos_w_poses):
+            cubemaps = self.pano_to_cubemap(pano_img)
+            for i in range(len(cubemaps)):
+                cubemap = tensor_to_pil(cubemaps[i])
+                cubemap = self.super_model(cubemap)
+                mesh_pose = self.cubemap_w2cs[i] @ pano_pose
+                c2w = self.mesh_pose_to_gs_pose(mesh_pose)
+                data["train"].append(
+                    {
+                        "camtoworld": c2w.astype(np.float32),
+                        "K": Ks.astype(np.float32),
+                        "image": np.array(cubemap),
+                        "image_h": image_h,
+                        "image_w": image_w,
+                        "image_id": len(cubemaps) * idx + i,
+                    }
+                )
+
+        # Camera poses for evaluation.
+        for idx in range(len(self.camera_poses)):
+            c2w = self.mesh_pose_to_gs_pose(self.camera_poses[idx])
+            data["eval"].append(
+                {
+                    "camtoworld": c2w.astype(np.float32),
+                    "K": Ks.astype(np.float32),
+                    "image_h": image_h,
+                    "image_w": image_w,
+                    "image_id": idx,
+                }
+            )
+
+        data_path = os.path.join(output_dir, self.cfg.gs_data_file)
+        torch.save(data, data_path)
+
+        return
+
+
+if __name__ == "__main__":
+    output_dir = "outputs/bg_v2/test3"
+    input_pano = "apps/assets/example_scene/result_pano.png"
+    config = Pano2MeshSRConfig()
+    pipeline = Pano2MeshSRPipeline(config)
+    pipeline(input_pano, output_dir)

embodied_gen/utils/config.py

@@ -0,0 +1,202 @@
+# 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 dataclasses import dataclass, field
+from typing import List, Optional, Union
+
+from dataclasses_json import DataClassJsonMixin
+from gsplat.strategy import DefaultStrategy, MCMCStrategy
+from typing_extensions import Literal, assert_never
+
+__all__ = [
+    "GptParamsConfig",
+    "Pano2MeshSRConfig",
+    "GsplatTrainConfig",
+]
+
+
+@dataclass
+class GptParamsConfig(DataClassJsonMixin):
+    temperature: float = 0.1
+    top_p: float = 0.1
+    frequency_penalty: float = 0.0
+    presence_penalty: float = 0.0
+    stop: int | None = None
+    max_tokens: int = 500
+
+
+@dataclass
+class Pano2MeshSRConfig:
+    mesh_file: str = "mesh_model.ply"
+    gs_data_file: str = "gs_data.pt"
+    device: str = "cuda"
+    blur_radius: int = 0
+    faces_per_pixel: int = 8
+    fov: int = 90
+    pano_w: int = 2048
+    pano_h: int = 1024
+    cubemap_w: int = 512
+    cubemap_h: int = 512
+    pose_scale: float = 0.6
+    pano_center_offset: tuple = (-0.2, 0.3)
+    inpaint_frame_stride: int = 20
+    trajectory_dir: str = "apps/assets/example_scene/camera_trajectory"
+    visualize: bool = False
+    depth_scale_factor: float = 3.4092
+    kernel_size: tuple = (9, 9)
+    upscale_factor: int = 4
+
+
+@dataclass
+class GsplatTrainConfig:
+    # Path to the .pt files. If provide, it will skip training and run evaluation only.
+    ckpt: Optional[List[str]] = None
+    # Render trajectory path
+    render_traj_path: str = "interp"
+
+    # Path to the Mip-NeRF 360 dataset
+    data_dir: str = "outputs/bg"
+    # Downsample factor for the dataset
+    data_factor: int = 4
+    # Directory to save results
+    result_dir: str = "outputs/bg"
+    # Every N images there is a test image
+    test_every: int = 8
+    # Random crop size for training  (experimental)
+    patch_size: Optional[int] = None
+    # A global scaler that applies to the scene size related parameters
+    global_scale: float = 1.0
+    # Normalize the world space
+    normalize_world_space: bool = True
+    # Camera model
+    camera_model: Literal["pinhole", "ortho", "fisheye"] = "pinhole"
+
+    # Port for the viewer server
+    port: int = 8080
+
+    # Batch size for training. Learning rates are scaled automatically
+    batch_size: int = 1
+    # A global factor to scale the number of training steps
+    steps_scaler: float = 1.0
+
+    # Number of training steps
+    max_steps: int = 30_000
+    # Steps to evaluate the model
+    eval_steps: List[int] = field(default_factory=lambda: [7_000, 30_000])
+    # Steps to save the model
+    save_steps: List[int] = field(default_factory=lambda: [7_000, 30_000])
+    # Whether to save ply file (storage size can be large)
+    save_ply: bool = True
+    # Steps to save the model as ply
+    ply_steps: List[int] = field(default_factory=lambda: [7_000, 30_000])
+    # Whether to disable video generation during training and evaluation
+    disable_video: bool = False
+
+    # Initial number of GSs. Ignored if using sfm
+    init_num_pts: int = 100_000
+    # Initial extent of GSs as a multiple of the camera extent. Ignored if using sfm
+    init_extent: float = 3.0
+    # Degree of spherical harmonics
+    sh_degree: int = 1
+    # Turn on another SH degree every this steps
+    sh_degree_interval: int = 1000
+    # Initial opacity of GS
+    init_opa: float = 0.1
+    # Initial scale of GS
+    init_scale: float = 1.0
+    # Weight for SSIM loss
+    ssim_lambda: float = 0.2
+
+    # Near plane clipping distance
+    near_plane: float = 0.01
+    # Far plane clipping distance
+    far_plane: float = 1e10
+
+    # Strategy for GS densification
+    strategy: Union[DefaultStrategy, MCMCStrategy] = field(
+        default_factory=DefaultStrategy
+    )
+    # Use packed mode for rasterization, this leads to less memory usage but slightly slower.
+    packed: bool = False
+    # Use sparse gradients for optimization. (experimental)
+    sparse_grad: bool = False
+    # Use visible adam from Taming 3DGS. (experimental)
+    visible_adam: bool = False
+    # Anti-aliasing in rasterization. Might slightly hurt quantitative metrics.
+    antialiased: bool = False
+
+    # Use random background for training to discourage transparency
+    random_bkgd: bool = False
+
+    # LR for 3D point positions
+    means_lr: float = 1.6e-4
+    # LR for Gaussian scale factors
+    scales_lr: float = 5e-3
+    # LR for alpha blending weights
+    opacities_lr: float = 5e-2
+    # LR for orientation (quaternions)
+    quats_lr: float = 1e-3
+    # LR for SH band 0 (brightness)
+    sh0_lr: float = 2.5e-3
+    # LR for higher-order SH (detail)
+    shN_lr: float = 2.5e-3 / 20
+
+    # Opacity regularization
+    opacity_reg: float = 0.0
+    # Scale regularization
+    scale_reg: float = 0.0
+
+    # Enable depth loss. (experimental)
+    depth_loss: bool = False
+    # Weight for depth loss
+    depth_lambda: float = 1e-2
+
+    # Dump information to tensorboard every this steps
+    tb_every: int = 200
+    # Save training images to tensorboard
+    tb_save_image: bool = False
+
+    lpips_net: Literal["vgg", "alex"] = "alex"
+
+    # 3DGUT (uncented transform + eval 3D)
+    with_ut: bool = False
+    with_eval3d: bool = False
+
+    scene_scale: float = 1.0
+
+    def adjust_steps(self, factor: float):
+        self.eval_steps = [int(i * factor) for i in self.eval_steps]
+        self.save_steps = [int(i * factor) for i in self.save_steps]
+        self.ply_steps = [int(i * factor) for i in self.ply_steps]
+        self.max_steps = int(self.max_steps * factor)
+        self.sh_degree_interval = int(self.sh_degree_interval * factor)
+
+        strategy = self.strategy
+        if isinstance(strategy, DefaultStrategy):
+            strategy.refine_start_iter = int(
+                strategy.refine_start_iter * factor
+            )
+            strategy.refine_stop_iter = int(strategy.refine_stop_iter * factor)
+            strategy.reset_every = int(strategy.reset_every * factor)
+            strategy.refine_every = int(strategy.refine_every * factor)
+        elif isinstance(strategy, MCMCStrategy):
+            strategy.refine_start_iter = int(
+                strategy.refine_start_iter * factor
+            )
+            strategy.refine_stop_iter = int(strategy.refine_stop_iter * factor)
+            strategy.refine_every = int(strategy.refine_every * factor)
+        else:
+            assert_never(strategy)

embodied_gen/utils/enum.py

@@ -0,0 +1,108 @@
+# 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 dataclasses import dataclass, field
+from enum import Enum
+
+from dataclasses_json import DataClassJsonMixin
+
+__all__ = [
+    "RenderItems",
+    "Scene3DItemEnum",
+    "SpatialRelationEnum",
+    "RobotItemEnum",
+]
+
+
+@dataclass
+class RenderItems(str, Enum):
+    IMAGE = "image_color"
+    ALPHA = "image_mask"
+    VIEW_NORMAL = "image_view_normal"
+    GLOBAL_NORMAL = "image_global_normal"
+    POSITION_MAP = "image_position"
+    DEPTH = "image_depth"
+    ALBEDO = "image_albedo"
+    DIFFUSE = "image_diffuse"
+
+
+@dataclass
+class Scene3DItemEnum(str, Enum):
+    BACKGROUND = "background"
+    CONTEXT = "context"
+    ROBOT = "robot"
+    MANIPULATED_OBJS = "manipulated_objs"
+    DISTRACTOR_OBJS = "distractor_objs"
+    OTHERS = "others"
+
+    @classmethod
+    def object_list(cls, layout_relation: dict) -> list:
+        return (
+            [
+                layout_relation[cls.BACKGROUND.value],
+                layout_relation[cls.CONTEXT.value],
+            ]
+            + layout_relation[cls.MANIPULATED_OBJS.value]
+            + layout_relation[cls.DISTRACTOR_OBJS.value]
+        )
+
+    @classmethod
+    def object_mapping(cls, layout_relation):
+        relation_mapping = {
+            # layout_relation[cls.ROBOT.value]: cls.ROBOT.value,
+            layout_relation[cls.BACKGROUND.value]: cls.BACKGROUND.value,
+            layout_relation[cls.CONTEXT.value]: cls.CONTEXT.value,
+        }
+        relation_mapping.update(
+            {
+                item: cls.MANIPULATED_OBJS.value
+                for item in layout_relation[cls.MANIPULATED_OBJS.value]
+            }
+        )
+        relation_mapping.update(
+            {
+                item: cls.DISTRACTOR_OBJS.value
+                for item in layout_relation[cls.DISTRACTOR_OBJS.value]
+            }
+        )
+
+        return relation_mapping
+
+
+@dataclass
+class SpatialRelationEnum(str, Enum):
+    ON = "ON"  # objects on the table
+    IN = "IN"  # objects in the room
+    INSIDE = "INSIDE"  # objects inside the shelf/rack
+    FLOOR = "FLOOR"  # object floor room/bin
+
+
+@dataclass
+class RobotItemEnum(str, Enum):
+    FRANKA = "franka"
+    UR5 = "ur5"
+    PIPER = "piper"
+
+
+@dataclass
+class LayoutInfo(DataClassJsonMixin):
+    tree: dict[str, list]
+    relation: dict[str, str | list[str]]
+    objs_desc: dict[str, str] = field(default_factory=dict)
+    objs_mapping: dict[str, str] = field(default_factory=dict)
+    assets: dict[str, str] = field(default_factory=dict)
+    quality: dict[str, str] = field(default_factory=dict)
+    position: dict[str, list[float]] = field(default_factory=dict)

embodied_gen/utils/gaussian.py

@@ -0,0 +1,330 @@
+# 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.
+# Part of the code comes from https://github.com/nerfstudio-project/gsplat
+# Both under the Apache License, Version 2.0.
+
+
+import math
+import random
+from io import BytesIO
+from typing import Dict, Literal, Optional, Tuple
+
+import numpy as np
+import torch
+import trimesh
+from gsplat.optimizers import SelectiveAdam
+from scipy.spatial.transform import Rotation
+from sklearn.neighbors import NearestNeighbors
+from torch import Tensor
+from embodied_gen.models.gs_model import GaussianOperator
+
+__all__ = [
+    "set_random_seed",
+    "export_splats",
+    "create_splats_with_optimizers",
+    "resize_pinhole_intrinsics",
+    "restore_scene_scale_and_position",
+]
+
+
+def knn(x: Tensor, K: int = 4) -> Tensor:
+    x_np = x.cpu().numpy()
+    model = NearestNeighbors(n_neighbors=K, metric="euclidean").fit(x_np)
+    distances, _ = model.kneighbors(x_np)
+    return torch.from_numpy(distances).to(x)
+
+
+def rgb_to_sh(rgb: Tensor) -> Tensor:
+    C0 = 0.28209479177387814
+    return (rgb - 0.5) / C0
+
+
+def set_random_seed(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def splat2ply_bytes(
+    means: torch.Tensor,
+    scales: torch.Tensor,
+    quats: torch.Tensor,
+    opacities: torch.Tensor,
+    sh0: torch.Tensor,
+    shN: torch.Tensor,
+) -> bytes:
+    num_splats = means.shape[0]
+    buffer = BytesIO()
+
+    # Write PLY header
+    buffer.write(b"ply\n")
+    buffer.write(b"format binary_little_endian 1.0\n")
+    buffer.write(f"element vertex {num_splats}\n".encode())
+    buffer.write(b"property float x\n")
+    buffer.write(b"property float y\n")
+    buffer.write(b"property float z\n")
+    for i, data in enumerate([sh0, shN]):
+        prefix = "f_dc" if i == 0 else "f_rest"
+        for j in range(data.shape[1]):
+            buffer.write(f"property float {prefix}_{j}\n".encode())
+    buffer.write(b"property float opacity\n")
+    for i in range(scales.shape[1]):
+        buffer.write(f"property float scale_{i}\n".encode())
+    for i in range(quats.shape[1]):
+        buffer.write(f"property float rot_{i}\n".encode())
+    buffer.write(b"end_header\n")
+
+    # Concatenate all tensors in the correct order
+    splat_data = torch.cat(
+        [means, sh0, shN, opacities.unsqueeze(1), scales, quats], dim=1
+    )
+    # Ensure correct dtype
+    splat_data = splat_data.to(torch.float32)
+
+    # Write binary data
+    float_dtype = np.dtype(np.float32).newbyteorder("<")
+    buffer.write(
+        splat_data.detach().cpu().numpy().astype(float_dtype).tobytes()
+    )
+
+    return buffer.getvalue()
+
+
+def export_splats(
+    means: torch.Tensor,
+    scales: torch.Tensor,
+    quats: torch.Tensor,
+    opacities: torch.Tensor,
+    sh0: torch.Tensor,
+    shN: torch.Tensor,
+    format: Literal["ply"] = "ply",
+    save_to: Optional[str] = None,
+) -> bytes:
+    """Export a Gaussian Splats model to bytes in PLY file format."""
+    total_splats = means.shape[0]
+    assert means.shape == (total_splats, 3), "Means must be of shape (N, 3)"
+    assert scales.shape == (total_splats, 3), "Scales must be of shape (N, 3)"
+    assert quats.shape == (
+        total_splats,
+        4,
+    ), "Quaternions must be of shape (N, 4)"
+    assert opacities.shape == (
+        total_splats,
+    ), "Opacities must be of shape (N,)"
+    assert sh0.shape == (total_splats, 1, 3), "sh0 must be of shape (N, 1, 3)"
+    assert (
+        shN.ndim == 3 and shN.shape[0] == total_splats and shN.shape[2] == 3
+    ), f"shN must be of shape (N, K, 3), got {shN.shape}"
+
+    # Reshape spherical harmonics
+    sh0 = sh0.squeeze(1)  # Shape (N, 3)
+    shN = shN.permute(0, 2, 1).reshape(means.shape[0], -1)  # Shape (N, K * 3)
+
+    # Check for NaN or Inf values
+    invalid_mask = (
+        torch.isnan(means).any(dim=1)
+        | torch.isinf(means).any(dim=1)
+        | torch.isnan(scales).any(dim=1)
+        | torch.isinf(scales).any(dim=1)
+        | torch.isnan(quats).any(dim=1)
+        | torch.isinf(quats).any(dim=1)
+        | torch.isnan(opacities).any(dim=0)
+        | torch.isinf(opacities).any(dim=0)
+        | torch.isnan(sh0).any(dim=1)
+        | torch.isinf(sh0).any(dim=1)
+        | torch.isnan(shN).any(dim=1)
+        | torch.isinf(shN).any(dim=1)
+    )
+
+    # Filter out invalid entries
+    valid_mask = ~invalid_mask
+    means = means[valid_mask]
+    scales = scales[valid_mask]
+    quats = quats[valid_mask]
+    opacities = opacities[valid_mask]
+    sh0 = sh0[valid_mask]
+    shN = shN[valid_mask]
+
+    if format == "ply":
+        data = splat2ply_bytes(means, scales, quats, opacities, sh0, shN)
+    else:
+        raise ValueError(f"Unsupported format: {format}")
+
+    if save_to:
+        with open(save_to, "wb") as binary_file:
+            binary_file.write(data)
+
+    return data
+
+
+def create_splats_with_optimizers(
+    points: np.ndarray = None,
+    points_rgb: np.ndarray = None,
+    init_num_pts: int = 100_000,
+    init_extent: float = 3.0,
+    init_opacity: float = 0.1,
+    init_scale: float = 1.0,
+    means_lr: float = 1.6e-4,
+    scales_lr: float = 5e-3,
+    opacities_lr: float = 5e-2,
+    quats_lr: float = 1e-3,
+    sh0_lr: float = 2.5e-3,
+    shN_lr: float = 2.5e-3 / 20,
+    scene_scale: float = 1.0,
+    sh_degree: int = 3,
+    sparse_grad: bool = False,
+    visible_adam: bool = False,
+    batch_size: int = 1,
+    feature_dim: Optional[int] = None,
+    device: str = "cuda",
+    world_rank: int = 0,
+    world_size: int = 1,
+) -> Tuple[torch.nn.ParameterDict, Dict[str, torch.optim.Optimizer]]:
+    if points is not None and points_rgb is not None:
+        points = torch.from_numpy(points).float()
+        rgbs = torch.from_numpy(points_rgb / 255.0).float()
+    else:
+        points = (
+            init_extent * scene_scale * (torch.rand((init_num_pts, 3)) * 2 - 1)
+        )
+        rgbs = torch.rand((init_num_pts, 3))
+
+    # Initialize the GS size to be the average dist of the 3 nearest neighbors
+    dist2_avg = (knn(points, 4)[:, 1:] ** 2).mean(dim=-1)  # [N,]
+    dist_avg = torch.sqrt(dist2_avg)
+    scales = (
+        torch.log(dist_avg * init_scale).unsqueeze(-1).repeat(1, 3)
+    )  # [N, 3]
+
+    # Distribute the GSs to different ranks (also works for single rank)
+    points = points[world_rank::world_size]
+    rgbs = rgbs[world_rank::world_size]
+    scales = scales[world_rank::world_size]
+
+    N = points.shape[0]
+    quats = torch.rand((N, 4))  # [N, 4]
+    opacities = torch.logit(torch.full((N,), init_opacity))  # [N,]
+
+    params = [
+        # name, value, lr
+        ("means", torch.nn.Parameter(points), means_lr * scene_scale),
+        ("scales", torch.nn.Parameter(scales), scales_lr),
+        ("quats", torch.nn.Parameter(quats), quats_lr),
+        ("opacities", torch.nn.Parameter(opacities), opacities_lr),
+    ]
+
+    if feature_dim is None:
+        # color is SH coefficients.
+        colors = torch.zeros((N, (sh_degree + 1) ** 2, 3))  # [N, K, 3]
+        colors[:, 0, :] = rgb_to_sh(rgbs)
+        params.append(("sh0", torch.nn.Parameter(colors[:, :1, :]), sh0_lr))
+        params.append(("shN", torch.nn.Parameter(colors[:, 1:, :]), shN_lr))
+    else:
+        # features will be used for appearance and view-dependent shading
+        features = torch.rand(N, feature_dim)  # [N, feature_dim]
+        params.append(("features", torch.nn.Parameter(features), sh0_lr))
+        colors = torch.logit(rgbs)  # [N, 3]
+        params.append(("colors", torch.nn.Parameter(colors), sh0_lr))
+
+    splats = torch.nn.ParameterDict({n: v for n, v, _ in params}).to(device)
+    # Scale learning rate based on batch size, reference:
+    # https://www.cs.princeton.edu/~smalladi/blog/2024/01/22/SDEs-ScalingRules/
+    # Note that this would not make the training exactly equivalent, see
+    # https://arxiv.org/pdf/2402.18824v1
+    BS = batch_size * world_size
+    optimizer_class = None
+    if sparse_grad:
+        optimizer_class = torch.optim.SparseAdam
+    elif visible_adam:
+        optimizer_class = SelectiveAdam
+    else:
+        optimizer_class = torch.optim.Adam
+    optimizers = {
+        name: optimizer_class(
+            [{"params": splats[name], "lr": lr * math.sqrt(BS), "name": name}],
+            eps=1e-15 / math.sqrt(BS),
+            # TODO: check betas logic when BS is larger than 10 betas[0] will be zero.
+            betas=(1 - BS * (1 - 0.9), 1 - BS * (1 - 0.999)),
+        )
+        for name, _, lr in params
+    }
+    return splats, optimizers
+
+
+def compute_intrinsics_from_fovy(
+    image_w: int, image_h: int, fovy_deg: float
+) -> np.ndarray:
+    fovy_rad = np.deg2rad(fovy_deg)
+    fy = image_h / (2 * np.tan(fovy_rad / 2))
+    fx = fy * (image_w / image_h)
+    cx = image_w / 2
+    cy = image_h / 2
+    K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+
+    return K
+
+
+def resize_pinhole_intrinsics(
+    raw_K: np.ndarray | torch.Tensor,
+    raw_hw: tuple[int, int],
+    new_hw: tuple[int, int],
+) -> np.ndarray:
+    raw_h, raw_w = raw_hw
+    new_h, new_w = new_hw
+
+    scale_x = new_w / raw_w
+    scale_y = new_h / raw_h
+
+    new_K = raw_K.copy() if isinstance(raw_K, np.ndarray) else raw_K.clone()
+    new_K[0, 0] *= scale_x  # fx
+    new_K[0, 2] *= scale_x  # cx
+    new_K[1, 1] *= scale_y  # fy
+    new_K[1, 2] *= scale_y  # cy
+
+    return new_K
+
+
+def restore_scene_scale_and_position(
+    real_height: float, mesh_path: str, gs_path: str
+) -> None:
+    """Scales a mesh and corresponding GS model to match a given real-world height.
+
+    Uses the 1st and 99th percentile of mesh Z-axis to estimate height,
+    applies scaling and vertical alignment, and updates both the mesh and GS model.
+
+    Args:
+        real_height (float): Target real-world height among Z axis.
+        mesh_path (str): Path to the input mesh file.
+        gs_path (str): Path to the Gaussian Splatting model file.
+    """
+    mesh = trimesh.load(mesh_path)
+    z_min = np.percentile(mesh.vertices[:, 1], 1)
+    z_max = np.percentile(mesh.vertices[:, 1], 99)
+    height = z_max - z_min
+    scale = real_height / height
+
+    rot = Rotation.from_quat([0, 1, 0, 0])
+    mesh.vertices = rot.apply(mesh.vertices)
+    mesh.vertices[:, 1] -= z_min
+    mesh.vertices *= scale
+    mesh.export(mesh_path)
+
+    gs_model: GaussianOperator = GaussianOperator.load_from_ply(gs_path)
+    gs_model = gs_model.get_gaussians(
+        instance_pose=torch.tensor([0.0, -z_min, 0, 0, 1, 0, 0])
+    )
+    gs_model.rescale(scale)
+    gs_model.save_to_ply(gs_path)

embodied_gen/utils/geometry.py

@@ -0,0 +1,515 @@
+# 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.
+
+import json
+import os
+import random
+from collections import defaultdict, deque
+from functools import wraps
+from typing import Literal
+
+import numpy as np
+import torch
+import trimesh
+from matplotlib.path import Path
+from pyquaternion import Quaternion
+from scipy.spatial import ConvexHull
+from scipy.spatial.transform import Rotation as R
+from shapely.geometry import Polygon
+from embodied_gen.utils.enum import LayoutInfo, Scene3DItemEnum
+from embodied_gen.utils.log import logger
+
+__all__ = [
+    "with_seed",
+    "matrix_to_pose",
+    "pose_to_matrix",
+    "quaternion_multiply",
+    "check_reachable",
+    "bfs_placement",
+    "compose_mesh_scene",
+    "compute_pinhole_intrinsics",
+]
+
+
+def matrix_to_pose(matrix: np.ndarray) -> list[float]:
+    """Convert a 4x4 transformation matrix to a pose (x, y, z, qx, qy, qz, qw).
+
+    Args:
+        matrix (np.ndarray): 4x4 transformation matrix.
+
+    Returns:
+        List[float]: Pose as [x, y, z, qx, qy, qz, qw].
+    """
+    x, y, z = matrix[:3, 3]
+    rot_mat = matrix[:3, :3]
+    quat = R.from_matrix(rot_mat).as_quat()
+    qx, qy, qz, qw = quat
+
+    return [x, y, z, qx, qy, qz, qw]
+
+
+def pose_to_matrix(pose: list[float]) -> np.ndarray:
+    """Convert pose (x, y, z, qx, qy, qz, qw) to a 4x4 transformation matrix.
+
+    Args:
+        List[float]: Pose as [x, y, z, qx, qy, qz, qw].
+
+    Returns:
+        matrix (np.ndarray): 4x4 transformation matrix.
+    """
+    x, y, z, qx, qy, qz, qw = pose
+    r = R.from_quat([qx, qy, qz, qw])
+    matrix = np.eye(4)
+    matrix[:3, :3] = r.as_matrix()
+    matrix[:3, 3] = [x, y, z]
+
+    return matrix
+
+
+def compute_xy_bbox(
+    vertices: np.ndarray, col_x: int = 0, col_y: int = 1
+) -> list[float]:
+    x_vals = vertices[:, col_x]
+    y_vals = vertices[:, col_y]
+    return x_vals.min(), x_vals.max(), y_vals.min(), y_vals.max()
+
+
+def has_iou_conflict(
+    new_box: list[float],
+    placed_boxes: list[list[float]],
+    iou_threshold: float = 0.0,
+) -> bool:
+    new_min_x, new_max_x, new_min_y, new_max_y = new_box
+    for min_x, max_x, min_y, max_y in placed_boxes:
+        ix1 = max(new_min_x, min_x)
+        iy1 = max(new_min_y, min_y)
+        ix2 = min(new_max_x, max_x)
+        iy2 = min(new_max_y, max_y)
+        inter_area = max(0, ix2 - ix1) * max(0, iy2 - iy1)
+        if inter_area > iou_threshold:
+            return True
+    return False
+
+
+def with_seed(seed_attr_name: str = "seed"):
+    """A parameterized decorator that temporarily sets the random seed."""
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            seed = kwargs.get(seed_attr_name, None)
+            if seed is not None:
+                py_state = random.getstate()
+                np_state = np.random.get_state()
+                torch_state = torch.get_rng_state()
+
+                random.seed(seed)
+                np.random.seed(seed)
+                torch.manual_seed(seed)
+                try:
+                    result = func(*args, **kwargs)
+                finally:
+                    random.setstate(py_state)
+                    np.random.set_state(np_state)
+                    torch.set_rng_state(torch_state)
+                return result
+            else:
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def compute_convex_hull_path(
+    vertices: np.ndarray,
+    z_threshold: float = 0.05,
+    interp_per_edge: int = 10,
+    margin: float = -0.02,
+    x_axis: int = 0,
+    y_axis: int = 1,
+    z_axis: int = 2,
+) -> Path:
+    top_vertices = vertices[
+        vertices[:, z_axis] > vertices[:, z_axis].max() - z_threshold
+    ]
+    top_xy = top_vertices[:, [x_axis, y_axis]]
+
+    if len(top_xy) < 3:
+        raise ValueError("Not enough points to form a convex hull")
+
+    hull = ConvexHull(top_xy)
+    hull_points = top_xy[hull.vertices]
+
+    polygon = Polygon(hull_points)
+    polygon = polygon.buffer(margin)
+    hull_points = np.array(polygon.exterior.coords)
+
+    dense_points = []
+    for i in range(len(hull_points)):
+        p1 = hull_points[i]
+        p2 = hull_points[(i + 1) % len(hull_points)]
+        for t in np.linspace(0, 1, interp_per_edge, endpoint=False):
+            pt = (1 - t) * p1 + t * p2
+            dense_points.append(pt)
+
+    return Path(np.array(dense_points), closed=True)
+
+
+def find_parent_node(node: str, tree: dict) -> str | None:
+    for parent, children in tree.items():
+        if any(child[0] == node for child in children):
+            return parent
+    return None
+
+
+def all_corners_inside(hull: Path, box: list, threshold: int = 3) -> bool:
+    x1, x2, y1, y2 = box
+    corners = [[x1, y1], [x2, y1], [x1, y2], [x2, y2]]
+
+    num_inside = sum(hull.contains_point(c) for c in corners)
+    return num_inside >= threshold
+
+
+def compute_axis_rotation_quat(
+    axis: Literal["x", "y", "z"], angle_rad: float
+) -> list[float]:
+    if axis.lower() == "x":
+        q = Quaternion(axis=[1, 0, 0], angle=angle_rad)
+    elif axis.lower() == "y":
+        q = Quaternion(axis=[0, 1, 0], angle=angle_rad)
+    elif axis.lower() == "z":
+        q = Quaternion(axis=[0, 0, 1], angle=angle_rad)
+    else:
+        raise ValueError(f"Unsupported axis '{axis}', must be one of x, y, z")
+
+    return [q.x, q.y, q.z, q.w]
+
+
+def quaternion_multiply(
+    init_quat: list[float], rotate_quat: list[float]
+) -> list[float]:
+    qx, qy, qz, qw = init_quat
+    q1 = Quaternion(w=qw, x=qx, y=qy, z=qz)
+    qx, qy, qz, qw = rotate_quat
+    q2 = Quaternion(w=qw, x=qx, y=qy, z=qz)
+    quat = q2 * q1
+
+    return [quat.x, quat.y, quat.z, quat.w]
+
+
+def check_reachable(
+    base_xyz: np.ndarray,
+    reach_xyz: np.ndarray,
+    min_reach: float = 0.25,
+    max_reach: float = 0.85,
+) -> bool:
+    """Check if the target point is within the reachable range."""
+    distance = np.linalg.norm(reach_xyz - base_xyz)
+
+    return min_reach < distance < max_reach
+
+
+@with_seed("seed")
+def bfs_placement(
+    layout_file: str,
+    floor_margin: float = 0,
+    beside_margin: float = 0.1,
+    max_attempts: int = 3000,
+    init_rpy: tuple = (1.5708, 0.0, 0.0),
+    rotate_objs: bool = True,
+    rotate_bg: bool = True,
+    rotate_context: bool = True,
+    limit_reach_range: tuple[float, float] | None = (0.20, 0.85),
+    max_orient_diff: float | None = 60,
+    robot_dim: float = 0.12,
+    seed: int = None,
+) -> LayoutInfo:
+    """Place objects in the layout using BFS traversal.
+
+    Args:
+        layout_file: Path to the JSON file defining the layout structure and assets.
+        floor_margin: Z-offset for the background object, typically for objects placed on the floor.
+        beside_margin: Minimum margin for objects placed 'beside' their parent, used when 'on' placement fails.
+        max_attempts: Maximum number of attempts to find a non-overlapping position for an object.
+        init_rpy: Initial Roll-Pitch-Yaw rotation rad applied to all object meshes to align the mesh's
+            coordinate system with the world's (e.g., Z-up).
+        rotate_objs: If True, apply a random rotation around the Z-axis for manipulated and distractor objects.
+        rotate_bg: If True, apply a random rotation around the Y-axis for the background object.
+        rotate_context: If True, apply a random rotation around the Z-axis for the context object.
+        limit_reach_range: If set, enforce a check that manipulated objects are within the robot's reach range, in meter.
+        max_orient_diff: If set, enforce a check that manipulated objects are within the robot's orientation range, in degree.
+        robot_dim: The approximate dimension (e.g., diameter) of the robot for box representation.
+        seed: Random seed for reproducible placement.
+
+    Returns:
+        A :class:`LayoutInfo` object containing the objects and their final computed 7D poses
+        ([x, y, z, qx, qy, qz, qw]).
+    """
+    layout_info = LayoutInfo.from_dict(json.load(open(layout_file, "r")))
+    asset_dir = os.path.dirname(layout_file)
+    object_mapping = layout_info.objs_mapping
+    position = {}  # node: [x, y, z, qx, qy, qz, qw]
+    parent_bbox_xy = {}
+    placed_boxes_map = defaultdict(list)
+    mesh_info = defaultdict(dict)
+    robot_node = layout_info.relation[Scene3DItemEnum.ROBOT.value]
+    for node in object_mapping:
+        if object_mapping[node] == Scene3DItemEnum.BACKGROUND.value:
+            bg_quat = (
+                compute_axis_rotation_quat(
+                    axis="y",
+                    angle_rad=np.random.uniform(0, 2 * np.pi),
+                )
+                if rotate_bg
+                else [0, 0, 0, 1]
+            )
+            bg_quat = [round(q, 4) for q in bg_quat]
+            continue
+
+        mesh_path = (
+            f"{layout_info.assets[node]}/mesh/{node.replace(' ', '_')}.obj"
+        )
+        mesh_path = os.path.join(asset_dir, mesh_path)
+        mesh_info[node]["path"] = mesh_path
+        mesh = trimesh.load(mesh_path)
+        rotation = R.from_euler("xyz", init_rpy, degrees=False)
+        vertices = mesh.vertices @ rotation.as_matrix().T
+        z1 = np.percentile(vertices[:, 2], 1)
+        z2 = np.percentile(vertices[:, 2], 99)
+
+        if object_mapping[node] == Scene3DItemEnum.CONTEXT.value:
+            object_quat = [0, 0, 0, 1]
+            if rotate_context:
+                angle_rad = np.random.uniform(0, 2 * np.pi)
+                object_quat = compute_axis_rotation_quat(
+                    axis="z", angle_rad=angle_rad
+                )
+                rotation = R.from_quat(object_quat).as_matrix()
+                vertices = vertices @ rotation.T
+
+            mesh_info[node]["surface"] = compute_convex_hull_path(vertices)
+
+            # Put robot in the CONTEXT edge.
+            x, y = random.choice(mesh_info[node]["surface"].vertices)
+            theta = np.arctan2(y, x)
+            quat_initial = Quaternion(axis=[0, 0, 1], angle=theta)
+            quat_extra = Quaternion(axis=[0, 0, 1], angle=np.pi)
+            quat = quat_extra * quat_initial
+            _pose = [x, y, z2 - z1, quat.x, quat.y, quat.z, quat.w]
+            position[robot_node] = [round(v, 4) for v in _pose]
+            node_box = [
+                x - robot_dim / 2,
+                x + robot_dim / 2,
+                y - robot_dim / 2,
+                y + robot_dim / 2,
+            ]
+            placed_boxes_map[node].append(node_box)
+        elif rotate_objs:
+            # For manipulated and distractor objects, apply random rotation
+            angle_rad = np.random.uniform(0, 2 * np.pi)
+            object_quat = compute_axis_rotation_quat(
+                axis="z", angle_rad=angle_rad
+            )
+            rotation = R.from_quat(object_quat).as_matrix()
+            vertices = vertices @ rotation.T
+
+        x1, x2, y1, y2 = compute_xy_bbox(vertices)
+        mesh_info[node]["pose"] = [x1, x2, y1, y2, z1, z2, *object_quat]
+        mesh_info[node]["area"] = max(1e-5, (x2 - x1) * (y2 - y1))
+
+    root = list(layout_info.tree.keys())[0]
+    queue = deque([((root, None), layout_info.tree.get(root, []))])
+    while queue:
+        (node, relation), children = queue.popleft()
+        if node not in object_mapping:
+            continue
+
+        if object_mapping[node] == Scene3DItemEnum.BACKGROUND.value:
+            position[node] = [0, 0, floor_margin, *bg_quat]
+        else:
+            x1, x2, y1, y2, z1, z2, qx, qy, qz, qw = mesh_info[node]["pose"]
+            if object_mapping[node] == Scene3DItemEnum.CONTEXT.value:
+                position[node] = [0, 0, -round(z1, 4), qx, qy, qz, qw]
+                parent_bbox_xy[node] = [x1, x2, y1, y2, z1, z2]
+            elif object_mapping[node] in [
+                Scene3DItemEnum.MANIPULATED_OBJS.value,
+                Scene3DItemEnum.DISTRACTOR_OBJS.value,
+            ]:
+                parent_node = find_parent_node(node, layout_info.tree)
+                parent_pos = position[parent_node]
+                (
+                    p_x1,
+                    p_x2,
+                    p_y1,
+                    p_y2,
+                    p_z1,
+                    p_z2,
+                ) = parent_bbox_xy[parent_node]
+
+                obj_dx = x2 - x1
+                obj_dy = y2 - y1
+                hull_path = mesh_info[parent_node].get("surface")
+                for _ in range(max_attempts):
+                    node_x1 = random.uniform(p_x1, p_x2 - obj_dx)
+                    node_y1 = random.uniform(p_y1, p_y2 - obj_dy)
+                    node_box = [
+                        node_x1,
+                        node_x1 + obj_dx,
+                        node_y1,
+                        node_y1 + obj_dy,
+                    ]
+                    if hull_path and not all_corners_inside(
+                        hull_path, node_box
+                    ):
+                        continue
+                    # Make sure the manipulated object is reachable by robot.
+                    if (
+                        limit_reach_range is not None
+                        and object_mapping[node]
+                        == Scene3DItemEnum.MANIPULATED_OBJS.value
+                    ):
+                        cx = parent_pos[0] + node_box[0] + obj_dx / 2
+                        cy = parent_pos[1] + node_box[2] + obj_dy / 2
+                        cz = parent_pos[2] + p_z2 - z1
+                        robot_pos = position[robot_node][:3]
+                        if not check_reachable(
+                            base_xyz=np.array(robot_pos),
+                            reach_xyz=np.array([cx, cy, cz]),
+                            min_reach=limit_reach_range[0],
+                            max_reach=limit_reach_range[1],
+                        ):
+                            continue
+
+                    # Make sure the manipulated object is inside the robot's orientation.
+                    if (
+                        max_orient_diff is not None
+                        and object_mapping[node]
+                        == Scene3DItemEnum.MANIPULATED_OBJS.value
+                    ):
+                        cx = parent_pos[0] + node_box[0] + obj_dx / 2
+                        cy = parent_pos[1] + node_box[2] + obj_dy / 2
+                        cx2, cy2 = position[robot_node][:2]
+                        v1 = np.array([-cx2, -cy2])
+                        v2 = np.array([cx - cx2, cy - cy2])
+                        dot = np.dot(v1, v2)
+                        norms = np.linalg.norm(v1) * np.linalg.norm(v2)
+                        theta = np.arccos(np.clip(dot / norms, -1.0, 1.0))
+                        theta = np.rad2deg(theta)
+                        if theta > max_orient_diff:
+                            continue
+
+                    if not has_iou_conflict(
+                        node_box, placed_boxes_map[parent_node]
+                    ):
+                        z_offset = 0
+                        break
+                else:
+                    logger.warning(
+                        f"Cannot place {node} on {parent_node} without overlap"
+                        f" after {max_attempts} attempts, place beside {parent_node}."
+                    )
+                    for _ in range(max_attempts):
+                        node_x1 = random.choice(
+                            [
+                                random.uniform(
+                                    p_x1 - obj_dx - beside_margin,
+                                    p_x1 - obj_dx,
+                                ),
+                                random.uniform(p_x2, p_x2 + beside_margin),
+                            ]
+                        )
+                        node_y1 = random.choice(
+                            [
+                                random.uniform(
+                                    p_y1 - obj_dy - beside_margin,
+                                    p_y1 - obj_dy,
+                                ),
+                                random.uniform(p_y2, p_y2 + beside_margin),
+                            ]
+                        )
+                        node_box = [
+                            node_x1,
+                            node_x1 + obj_dx,
+                            node_y1,
+                            node_y1 + obj_dy,
+                        ]
+                        z_offset = -(parent_pos[2] + p_z2)
+                        if not has_iou_conflict(
+                            node_box, placed_boxes_map[parent_node]
+                        ):
+                            break
+
+                placed_boxes_map[parent_node].append(node_box)
+
+                abs_cx = parent_pos[0] + node_box[0] + obj_dx / 2
+                abs_cy = parent_pos[1] + node_box[2] + obj_dy / 2
+                abs_cz = parent_pos[2] + p_z2 - z1 + z_offset
+                position[node] = [
+                    round(v, 4)
+                    for v in [abs_cx, abs_cy, abs_cz, qx, qy, qz, qw]
+                ]
+                parent_bbox_xy[node] = [x1, x2, y1, y2, z1, z2]
+
+        sorted_children = sorted(
+            children, key=lambda x: -mesh_info[x[0]].get("area", 0)
+        )
+        for child, rel in sorted_children:
+            queue.append(((child, rel), layout_info.tree.get(child, [])))
+
+    layout_info.position = position
+
+    return layout_info
+
+
+def compose_mesh_scene(
+    layout_info: LayoutInfo, out_scene_path: str, with_bg: bool = False
+) -> None:
+    object_mapping = Scene3DItemEnum.object_mapping(layout_info.relation)
+    scene = trimesh.Scene()
+    for node in layout_info.assets:
+        if object_mapping[node] == Scene3DItemEnum.BACKGROUND.value:
+            mesh_path = f"{layout_info.assets[node]}/mesh_model.ply"
+            if not with_bg:
+                continue
+        else:
+            mesh_path = (
+                f"{layout_info.assets[node]}/mesh/{node.replace(' ', '_')}.obj"
+            )
+
+        mesh = trimesh.load(mesh_path)
+        offset = np.array(layout_info.position[node])[[0, 2, 1]]
+        mesh.vertices += offset
+        scene.add_geometry(mesh, node_name=node)
+
+    os.makedirs(os.path.dirname(out_scene_path), exist_ok=True)
+    scene.export(out_scene_path)
+    logger.info(f"Composed interactive 3D layout saved in {out_scene_path}")
+
+    return
+
+
+def compute_pinhole_intrinsics(
+    image_w: int, image_h: int, fov_deg: float
+) -> np.ndarray:
+    fov_rad = np.deg2rad(fov_deg)
+    fx = image_w / (2 * np.tan(fov_rad / 2))
+    fy = fx  # assuming square pixels
+    cx = image_w / 2
+    cy = image_h / 2
+    K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+
+    return K

embodied_gen/utils/gpt_clients.py

@@ -0,0 +1,218 @@
+# 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.
+
+
+import base64
+import logging
+import os
+from io import BytesIO
+from typing import Optional
+
+import yaml
+from openai import AzureOpenAI, OpenAI  # pip install openai
+from PIL import Image
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    stop_after_delay,
+    wait_random_exponential,
+)
+from embodied_gen.utils.process_media import combine_images_to_grid
+
+logging.getLogger("httpx").setLevel(logging.WARNING)
+logging.basicConfig(level=logging.WARNING)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "GPTclient",
+]
+
+CONFIG_FILE = "embodied_gen/utils/gpt_config.yaml"
+
+
+class GPTclient:
+    """A client to interact with the GPT model via OpenAI or Azure API."""
+
+    def __init__(
+        self,
+        endpoint: str,
+        api_key: str,
+        model_name: str = "yfb-gpt-4o",
+        api_version: str = None,
+        check_connection: bool = True,
+        verbose: bool = False,
+    ):
+        if api_version is not None:
+            self.client = AzureOpenAI(
+                azure_endpoint=endpoint,
+                api_key=api_key,
+                api_version=api_version,
+            )
+        else:
+            self.client = OpenAI(
+                base_url=endpoint,
+                api_key=api_key,
+            )
+
+        self.endpoint = endpoint
+        self.model_name = model_name
+        self.image_formats = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif"}
+        self.verbose = verbose
+        if check_connection:
+            self.check_connection()
+
+        logger.info(f"Using GPT model: {self.model_name}.")
+
+    @retry(
+        wait=wait_random_exponential(min=1, max=20),
+        stop=(stop_after_attempt(10) | stop_after_delay(30)),
+    )
+    def completion_with_backoff(self, **kwargs):
+        return self.client.chat.completions.create(**kwargs)
+
+    def query(
+        self,
+        text_prompt: str,
+        image_base64: Optional[list[str | Image.Image]] = None,
+        system_role: Optional[str] = None,
+        params: Optional[dict] = None,
+    ) -> Optional[str]:
+        """Queries the GPT model with a text and optional image prompts.
+
+        Args:
+            text_prompt (str): The main text input that the model responds to.
+            image_base64 (Optional[List[str]]): A list of image base64 strings
+                or local image paths or PIL.Image to accompany the text prompt.
+            system_role (Optional[str]): Optional system-level instructions
+                that specify the behavior of the assistant.
+            params (Optional[dict]): Additional parameters for GPT setting.
+
+        Returns:
+            Optional[str]: The response content generated by the model based on
+                the prompt. Returns `None` if an error occurs.
+        """
+        if system_role is None:
+            system_role = "You are a highly knowledgeable assistant specializing in physics, engineering, and object properties."  # noqa
+
+        content_user = [
+            {
+                "type": "text",
+                "text": text_prompt,
+            },
+        ]
+
+        # Process images if provided
+        if image_base64 is not None:
+            if not isinstance(image_base64, list):
+                image_base64 = [image_base64]
+            # Hardcode tmp because of the openrouter can't input multi images.
+            if "openrouter" in self.endpoint:
+                image_base64 = combine_images_to_grid(image_base64)
+            for img in image_base64:
+                if isinstance(img, Image.Image):
+                    buffer = BytesIO()
+                    img.save(buffer, format=img.format or "PNG")
+                    buffer.seek(0)
+                    image_binary = buffer.read()
+                    img = base64.b64encode(image_binary).decode("utf-8")
+                elif (
+                    len(os.path.splitext(img)) > 1
+                    and os.path.splitext(img)[-1].lower() in self.image_formats
+                ):
+                    if not os.path.exists(img):
+                        raise FileNotFoundError(f"Image file not found: {img}")
+                    with open(img, "rb") as f:
+                        img = base64.b64encode(f.read()).decode("utf-8")
+
+                content_user.append(
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/png;base64,{img}"},
+                    }
+                )
+
+        payload = {
+            "messages": [
+                {"role": "system", "content": system_role},
+                {"role": "user", "content": content_user},
+            ],
+            "temperature": 0.1,
+            "max_tokens": 500,
+            "top_p": 0.1,
+            "frequency_penalty": 0,
+            "presence_penalty": 0,
+            "stop": None,
+            "model": self.model_name,
+        }
+
+        if params:
+            payload.update(params)
+
+        response = None
+        try:
+            response = self.completion_with_backoff(**payload)
+            response = response.choices[0].message.content
+        except Exception as e:
+            logger.error(f"Error GPTclint {self.endpoint} API call: {e}")
+            response = None
+
+        if self.verbose:
+            logger.info(f"Prompt: {text_prompt}")
+            logger.info(f"Response: {response}")
+
+        return response
+
+    def check_connection(self) -> None:
+        """Check whether the GPT API connection is working."""
+        try:
+            response = self.completion_with_backoff(
+                messages=[
+                    {"role": "system", "content": "You are a test system."},
+                    {"role": "user", "content": "Hello"},
+                ],
+                model=self.model_name,
+                temperature=0,
+                max_tokens=100,
+            )
+            content = response.choices[0].message.content
+            logger.info(f"Connection check success.")
+        except Exception as e:
+            raise ConnectionError(
+                f"Failed to connect to GPT API at {self.endpoint}, "
+                f"please check setting in `{CONFIG_FILE}` and `README`."
+            )
+
+
+with open(CONFIG_FILE, "r") as f:
+    config = yaml.safe_load(f)
+
+agent_type = config["agent_type"]
+agent_config = config.get(agent_type, {})
+
+# Prefer environment variables, fallback to YAML config
+endpoint = os.environ.get("ENDPOINT", agent_config.get("endpoint"))
+api_key = os.environ.get("API_KEY", agent_config.get("api_key"))
+api_version = os.environ.get("API_VERSION", agent_config.get("api_version"))
+model_name = os.environ.get("MODEL_NAME", agent_config.get("model_name"))
+
+GPT_CLIENT = GPTclient(
+    endpoint=endpoint,
+    api_key=api_key,
+    api_version=api_version,
+    model_name=model_name,
+    check_connection=False,
+)

embodied_gen/utils/gpt_config.yaml

@@ -0,0 +1,14 @@
+# config.yaml
+agent_type: "qwen2.5-vl" # gpt-4o or qwen2.5-vl
+
+gpt-4o:
+  endpoint: https://xxx.openai.azure.com
+  api_key: xxx
+  api_version: 2025-xx-xx
+  model_name: yfb-gpt-4o
+
+qwen2.5-vl:
+  endpoint: https://openrouter.ai/api/v1
+  api_key: sk-or-v1-xxx
+  api_version: null
+  model_name: qwen/qwen2.5-vl-72b-instruct:free

embodied_gen/utils/log.py

@@ -0,0 +1,48 @@
+# 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.
+
+import logging
+
+from colorlog import ColoredFormatter
+
+__all__ = [
+    "logger",
+]
+
+LOG_FORMAT = (
+    "%(log_color)s[%(asctime)s] %(levelname)-8s | %(message)s%(reset)s"
+)
+DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+formatter = ColoredFormatter(
+    LOG_FORMAT,
+    datefmt=DATE_FORMAT,
+    log_colors={
+        "DEBUG": "cyan",
+        "INFO": "green",
+        "WARNING": "yellow",
+        "ERROR": "red",
+        "CRITICAL": "bold_red",
+    },
+)
+
+handler = logging.StreamHandler()
+handler.setFormatter(formatter)
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logger.addHandler(handler)
+logger.propagate = False

embodied_gen/utils/monkey_patches.py

@@ -0,0 +1,218 @@
+# 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.
+
+import os
+import sys
+import zipfile
+
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download
+from omegaconf import OmegaConf
+from PIL import Image
+from torchvision import transforms
+
+
+def monkey_patch_pano2room():
+    current_file_path = os.path.abspath(__file__)
+    current_dir = os.path.dirname(current_file_path)
+    sys.path.append(os.path.join(current_dir, "../.."))
+    sys.path.append(os.path.join(current_dir, "../../thirdparty/pano2room"))
+    from thirdparty.pano2room.modules.geo_predictors.omnidata.omnidata_normal_predictor import (
+        OmnidataNormalPredictor,
+    )
+    from thirdparty.pano2room.modules.geo_predictors.omnidata.omnidata_predictor import (
+        OmnidataPredictor,
+    )
+
+    def patched_omni_depth_init(self):
+        self.img_size = 384
+        self.model = torch.hub.load(
+            'alexsax/omnidata_models', 'depth_dpt_hybrid_384'
+        )
+        self.model.eval()
+        self.trans_totensor = transforms.Compose(
+            [
+                transforms.Resize(self.img_size, interpolation=Image.BILINEAR),
+                transforms.CenterCrop(self.img_size),
+                transforms.Normalize(mean=0.5, std=0.5),
+            ]
+        )
+
+    OmnidataPredictor.__init__ = patched_omni_depth_init
+
+    def patched_omni_normal_init(self):
+        self.img_size = 384
+        self.model = torch.hub.load(
+            'alexsax/omnidata_models', 'surface_normal_dpt_hybrid_384'
+        )
+        self.model.eval()
+        self.trans_totensor = transforms.Compose(
+            [
+                transforms.Resize(self.img_size, interpolation=Image.BILINEAR),
+                transforms.CenterCrop(self.img_size),
+                transforms.Normalize(mean=0.5, std=0.5),
+            ]
+        )
+
+    OmnidataNormalPredictor.__init__ = patched_omni_normal_init
+
+    def patched_panojoint_init(self, save_path=None):
+        self.depth_predictor = OmnidataPredictor()
+        self.normal_predictor = OmnidataNormalPredictor()
+        self.save_path = save_path
+
+    from modules.geo_predictors import PanoJointPredictor
+
+    PanoJointPredictor.__init__ = patched_panojoint_init
+
+    # NOTE: We use gsplat instead.
+    # import depth_diff_gaussian_rasterization_min as ddgr
+    # from dataclasses import dataclass
+    # @dataclass
+    # class PatchedGaussianRasterizationSettings:
+    #     image_height: int
+    #     image_width: int
+    #     tanfovx: float
+    #     tanfovy: float
+    #     bg: torch.Tensor
+    #     scale_modifier: float
+    #     viewmatrix: torch.Tensor
+    #     projmatrix: torch.Tensor
+    #     sh_degree: int
+    #     campos: torch.Tensor
+    #     prefiltered: bool
+    #     debug: bool = False
+    # ddgr.GaussianRasterizationSettings = PatchedGaussianRasterizationSettings
+
+    # disable get_has_ddp_rank print in `BaseInpaintingTrainingModule`
+    os.environ["NODE_RANK"] = "0"
+
+    from thirdparty.pano2room.modules.inpainters.lama.saicinpainting.training.trainers import (
+        load_checkpoint,
+    )
+    from thirdparty.pano2room.modules.inpainters.lama_inpainter import (
+        LamaInpainter,
+    )
+
+    def patched_lama_inpaint_init(self):
+        zip_path = hf_hub_download(
+            repo_id="smartywu/big-lama",
+            filename="big-lama.zip",
+            repo_type="model",
+        )
+        extract_dir = os.path.splitext(zip_path)[0]
+
+        if not os.path.exists(extract_dir):
+            os.makedirs(extract_dir, exist_ok=True)
+            with zipfile.ZipFile(zip_path, "r") as zip_ref:
+                zip_ref.extractall(extract_dir)
+
+        config_path = os.path.join(extract_dir, 'big-lama', 'config.yaml')
+        checkpoint_path = os.path.join(
+            extract_dir, 'big-lama/models/best.ckpt'
+        )
+        train_config = OmegaConf.load(config_path)
+        train_config.training_model.predict_only = True
+        train_config.visualizer.kind = 'noop'
+
+        self.model = load_checkpoint(
+            train_config, checkpoint_path, strict=False, map_location='cpu'
+        )
+        self.model.freeze()
+
+    LamaInpainter.__init__ = patched_lama_inpaint_init
+
+    from diffusers import StableDiffusionInpaintPipeline
+    from thirdparty.pano2room.modules.inpainters.SDFT_inpainter import (
+        SDFTInpainter,
+    )
+
+    def patched_sd_inpaint_init(self, subset_name=None):
+        super(SDFTInpainter, self).__init__()
+        pipe = StableDiffusionInpaintPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-2-inpainting",
+            torch_dtype=torch.float16,
+        ).to("cuda")
+        pipe.enable_model_cpu_offload()
+        self.inpaint_pipe = pipe
+
+    SDFTInpainter.__init__ = patched_sd_inpaint_init
+
+
+def monkey_patch_maniskill():
+    from mani_skill.envs.scene import ManiSkillScene
+
+    def get_sensor_images(
+        self, obs: dict[str, any]
+    ) -> dict[str, dict[str, torch.Tensor]]:
+        sensor_data = dict()
+        for name, sensor in self.sensors.items():
+            sensor_data[name] = sensor.get_images(obs[name])
+        return sensor_data
+
+    def get_human_render_camera_images(
+        self, camera_name: str = None, return_alpha: bool = False
+    ) -> dict[str, torch.Tensor]:
+        def get_rgba_tensor(camera, return_alpha):
+            color = camera.get_obs(
+                rgb=True, depth=False, segmentation=False, position=False
+            )["rgb"]
+            if return_alpha:
+                seg_labels = camera.get_obs(
+                    rgb=False, depth=False, segmentation=True, position=False
+                )["segmentation"]
+                masks = np.where((seg_labels.cpu() > 1), 255, 0).astype(
+                    np.uint8
+                )
+                masks = torch.tensor(masks).to(color.device)
+                color = torch.concat([color, masks], dim=-1)
+
+            return color
+
+        image_data = dict()
+        if self.gpu_sim_enabled:
+            if self.parallel_in_single_scene:
+                for name, camera in self.human_render_cameras.items():
+                    camera.camera._render_cameras[0].take_picture()
+                    rgba = get_rgba_tensor(camera, return_alpha)
+                    image_data[name] = rgba
+            else:
+                for name, camera in self.human_render_cameras.items():
+                    if camera_name is not None and name != camera_name:
+                        continue
+                    assert camera.config.shader_config.shader_pack not in [
+                        "rt",
+                        "rt-fast",
+                        "rt-med",
+                    ], "ray tracing shaders do not work with parallel rendering"
+                    camera.capture()
+                    rgba = get_rgba_tensor(camera, return_alpha)
+                    image_data[name] = rgba
+        else:
+            for name, camera in self.human_render_cameras.items():
+                if camera_name is not None and name != camera_name:
+                    continue
+                camera.capture()
+                rgba = get_rgba_tensor(camera, return_alpha)
+                image_data[name] = rgba
+
+        return image_data
+
+    ManiSkillScene.get_sensor_images = get_sensor_images
+    ManiSkillScene.get_human_render_camera_images = (
+        get_human_render_camera_images
+    )

embodied_gen/utils/process_media.py

@@ -0,0 +1,467 @@
+# 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.
+
+
+import logging
+import math
+import mimetypes
+import os
+import textwrap
+from glob import glob
+from typing import Union
+
+import cv2
+import imageio
+import matplotlib.pyplot as plt
+import networkx as nx
+import numpy as np
+import spaces
+from matplotlib.patches import Patch
+from moviepy.editor import VideoFileClip, clips_array
+from PIL import Image
+from embodied_gen.data.differentiable_render import entrypoint as render_api
+from embodied_gen.utils.enum import LayoutInfo, Scene3DItemEnum
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "render_asset3d",
+    "merge_images_video",
+    "filter_small_connected_components",
+    "filter_image_small_connected_components",
+    "combine_images_to_grid",
+    "SceneTreeVisualizer",
+    "is_image_file",
+    "parse_text_prompts",
+    "check_object_edge_truncated",
+    "vcat_pil_images",
+]
+
+
+@spaces.GPU
+def render_asset3d(
+    mesh_path: str,
+    output_root: str,
+    distance: float = 5.0,
+    num_images: int = 1,
+    elevation: list[float] = (0.0,),
+    pbr_light_factor: float = 1.2,
+    return_key: str = "image_color/*",
+    output_subdir: str = "renders",
+    gen_color_mp4: bool = False,
+    gen_viewnormal_mp4: bool = False,
+    gen_glonormal_mp4: bool = False,
+    no_index_file: bool = False,
+    with_mtl: bool = True,
+) -> list[str]:
+    input_args = dict(
+        mesh_path=mesh_path,
+        output_root=output_root,
+        uuid=output_subdir,
+        distance=distance,
+        num_images=num_images,
+        elevation=elevation,
+        pbr_light_factor=pbr_light_factor,
+        with_mtl=with_mtl,
+        gen_color_mp4=gen_color_mp4,
+        gen_viewnormal_mp4=gen_viewnormal_mp4,
+        gen_glonormal_mp4=gen_glonormal_mp4,
+        no_index_file=no_index_file,
+    )
+
+    try:
+        _ = render_api(**input_args)
+    except Exception as e:
+        logger.error(f"Error occurred during rendering: {e}.")
+
+    dst_paths = glob(os.path.join(output_root, output_subdir, return_key))
+
+    return dst_paths
+
+
+def merge_images_video(color_images, normal_images, output_path) -> None:
+    width = color_images[0].shape[1]
+    combined_video = [
+        np.hstack([rgb_img[:, : width // 2], normal_img[:, width // 2 :]])
+        for rgb_img, normal_img in zip(color_images, normal_images)
+    ]
+    imageio.mimsave(output_path, combined_video, fps=50)
+
+    return
+
+
+def merge_video_video(
+    video_path1: str, video_path2: str, output_path: str
+) -> None:
+    """Merge two videos by the left half and the right half of the videos."""
+    clip1 = VideoFileClip(video_path1)
+    clip2 = VideoFileClip(video_path2)
+
+    if clip1.size != clip2.size:
+        raise ValueError("The resolutions of the two videos do not match.")
+
+    width, height = clip1.size
+    clip1_half = clip1.crop(x1=0, y1=0, x2=width // 2, y2=height)
+    clip2_half = clip2.crop(x1=width // 2, y1=0, x2=width, y2=height)
+    final_clip = clips_array([[clip1_half, clip2_half]])
+    final_clip.write_videofile(output_path, codec="libx264")
+
+
+def filter_small_connected_components(
+    mask: Union[Image.Image, np.ndarray],
+    area_ratio: float,
+    connectivity: int = 8,
+) -> np.ndarray:
+    if isinstance(mask, Image.Image):
+        mask = np.array(mask)
+    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
+        mask,
+        connectivity=connectivity,
+    )
+
+    small_components = np.zeros_like(mask, dtype=np.uint8)
+    mask_area = (mask != 0).sum()
+    min_area = mask_area // area_ratio
+    for label in range(1, num_labels):
+        area = stats[label, cv2.CC_STAT_AREA]
+        if area < min_area:
+            small_components[labels == label] = 255
+
+    mask = cv2.bitwise_and(mask, cv2.bitwise_not(small_components))
+
+    return mask
+
+
+def filter_image_small_connected_components(
+    image: Union[Image.Image, np.ndarray],
+    area_ratio: float = 10,
+    connectivity: int = 8,
+) -> np.ndarray:
+    if isinstance(image, Image.Image):
+        image = image.convert("RGBA")
+        image = np.array(image)
+
+    mask = image[..., 3]
+    mask = filter_small_connected_components(mask, area_ratio, connectivity)
+    image[..., 3] = mask
+
+    return image
+
+
+def combine_images_to_grid(
+    images: list[str | Image.Image],
+    cat_row_col: tuple[int, int] = None,
+    target_wh: tuple[int, int] = (512, 512),
+    image_mode: str = "RGB",
+) -> list[Image.Image]:
+    n_images = len(images)
+    if n_images == 1:
+        return images
+
+    if cat_row_col is None:
+        n_col = math.ceil(math.sqrt(n_images))
+        n_row = math.ceil(n_images / n_col)
+    else:
+        n_row, n_col = cat_row_col
+
+    images = [
+        Image.open(p).convert(image_mode) if isinstance(p, str) else p
+        for p in images
+    ]
+    images = [img.resize(target_wh) for img in images]
+
+    grid_w, grid_h = n_col * target_wh[0], n_row * target_wh[1]
+    grid = Image.new(image_mode, (grid_w, grid_h), (0, 0, 0))
+
+    for idx, img in enumerate(images):
+        row, col = divmod(idx, n_col)
+        grid.paste(img, (col * target_wh[0], row * target_wh[1]))
+
+    return [grid]
+
+
+class SceneTreeVisualizer:
+    def __init__(self, layout_info: LayoutInfo) -> None:
+        self.tree = layout_info.tree
+        self.relation = layout_info.relation
+        self.objs_desc = layout_info.objs_desc
+        self.G = nx.DiGraph()
+        self.root = self._find_root()
+        self._build_graph()
+
+        self.role_colors = {
+            Scene3DItemEnum.BACKGROUND.value: "plum",
+            Scene3DItemEnum.CONTEXT.value: "lightblue",
+            Scene3DItemEnum.ROBOT.value: "lightcoral",
+            Scene3DItemEnum.MANIPULATED_OBJS.value: "lightgreen",
+            Scene3DItemEnum.DISTRACTOR_OBJS.value: "lightgray",
+            Scene3DItemEnum.OTHERS.value: "orange",
+        }
+
+    def _find_root(self) -> str:
+        children = {c for cs in self.tree.values() for c, _ in cs}
+        parents = set(self.tree.keys())
+        roots = parents - children
+        if not roots:
+            raise ValueError("No root node found.")
+        return next(iter(roots))
+
+    def _build_graph(self):
+        for parent, children in self.tree.items():
+            for child, relation in children:
+                self.G.add_edge(parent, child, relation=relation)
+
+    def _get_node_role(self, node: str) -> str:
+        if node == self.relation.get(Scene3DItemEnum.BACKGROUND.value):
+            return Scene3DItemEnum.BACKGROUND.value
+        if node == self.relation.get(Scene3DItemEnum.CONTEXT.value):
+            return Scene3DItemEnum.CONTEXT.value
+        if node == self.relation.get(Scene3DItemEnum.ROBOT.value):
+            return Scene3DItemEnum.ROBOT.value
+        if node in self.relation.get(
+            Scene3DItemEnum.MANIPULATED_OBJS.value, []
+        ):
+            return Scene3DItemEnum.MANIPULATED_OBJS.value
+        if node in self.relation.get(
+            Scene3DItemEnum.DISTRACTOR_OBJS.value, []
+        ):
+            return Scene3DItemEnum.DISTRACTOR_OBJS.value
+        return Scene3DItemEnum.OTHERS.value
+
+    def _get_positions(
+        self, root, width=1.0, vert_gap=0.1, vert_loc=1, xcenter=0.5, pos=None
+    ):
+        if pos is None:
+            pos = {root: (xcenter, vert_loc)}
+        else:
+            pos[root] = (xcenter, vert_loc)
+
+        children = list(self.G.successors(root))
+        if children:
+            dx = width / len(children)
+            next_x = xcenter - width / 2 - dx / 2
+            for child in children:
+                next_x += dx
+                pos = self._get_positions(
+                    child,
+                    width=dx,
+                    vert_gap=vert_gap,
+                    vert_loc=vert_loc - vert_gap,
+                    xcenter=next_x,
+                    pos=pos,
+                )
+        return pos
+
+    def render(
+        self,
+        save_path: str,
+        figsize=(8, 6),
+        dpi=300,
+        title: str = "Scene 3D Hierarchy Tree",
+    ):
+        node_colors = [
+            self.role_colors[self._get_node_role(n)] for n in self.G.nodes
+        ]
+        pos = self._get_positions(self.root)
+
+        plt.figure(figsize=figsize)
+        nx.draw(
+            self.G,
+            pos,
+            with_labels=True,
+            arrows=False,
+            node_size=2000,
+            node_color=node_colors,
+            font_size=10,
+            font_weight="bold",
+        )
+
+        # Draw edge labels
+        edge_labels = nx.get_edge_attributes(self.G, "relation")
+        nx.draw_networkx_edge_labels(
+            self.G,
+            pos,
+            edge_labels=edge_labels,
+            font_size=9,
+            font_color="black",
+        )
+
+        # Draw small description text under each node (if available)
+        for node, (x, y) in pos.items():
+            desc = self.objs_desc.get(node)
+            if desc:
+                wrapped = "\n".join(textwrap.wrap(desc, width=30))
+                plt.text(
+                    x,
+                    y - 0.006,
+                    wrapped,
+                    fontsize=6,
+                    ha="center",
+                    va="top",
+                    wrap=True,
+                    color="black",
+                    bbox=dict(
+                        facecolor="dimgray",
+                        edgecolor="darkgray",
+                        alpha=0.1,
+                        boxstyle="round,pad=0.2",
+                    ),
+                )
+
+        plt.title(title, fontsize=12)
+        task_desc = self.relation.get("task_desc", "")
+        if task_desc:
+            plt.suptitle(
+                f"Task Description: {task_desc}", fontsize=10, y=0.999
+            )
+
+        plt.axis("off")
+
+        legend_handles = [
+            Patch(facecolor=color, edgecolor='black', label=role)
+            for role, color in self.role_colors.items()
+        ]
+        plt.legend(
+            handles=legend_handles,
+            loc="lower center",
+            ncol=3,
+            bbox_to_anchor=(0.5, -0.1),
+            fontsize=9,
+        )
+
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        plt.savefig(save_path, dpi=dpi, bbox_inches="tight")
+        plt.close()
+
+
+def load_scene_dict(file_path: str) -> dict:
+    scene_dict = {}
+    with open(file_path, "r", encoding='utf-8') as f:
+        for line in f:
+            line = line.strip()
+            if not line or ":" not in line:
+                continue
+            scene_id, desc = line.split(":", 1)
+            scene_dict[scene_id.strip()] = desc.strip()
+
+    return scene_dict
+
+
+def is_image_file(filename: str) -> bool:
+    mime_type, _ = mimetypes.guess_type(filename)
+
+    return mime_type is not None and mime_type.startswith('image')
+
+
+def parse_text_prompts(prompts: list[str]) -> list[str]:
+    if len(prompts) == 1 and prompts[0].endswith(".txt"):
+        with open(prompts[0], "r") as f:
+            prompts = [
+                line.strip()
+                for line in f
+                if line.strip() and not line.strip().startswith("#")
+            ]
+    return prompts
+
+
+def alpha_blend_rgba(
+    fg_image: Union[str, Image.Image, np.ndarray],
+    bg_image: Union[str, Image.Image, np.ndarray],
+) -> Image.Image:
+    """Alpha blends a foreground RGBA image over a background RGBA image.
+
+    Args:
+        fg_image: Foreground image. Can be a file path (str), a PIL Image,
+            or a NumPy ndarray.
+        bg_image: Background image. Can be a file path (str), a PIL Image,
+            or a NumPy ndarray.
+
+    Returns:
+        A PIL Image representing the alpha-blended result in RGBA mode.
+    """
+    if isinstance(fg_image, str):
+        fg_image = Image.open(fg_image)
+    elif isinstance(fg_image, np.ndarray):
+        fg_image = Image.fromarray(fg_image)
+
+    if isinstance(bg_image, str):
+        bg_image = Image.open(bg_image)
+    elif isinstance(bg_image, np.ndarray):
+        bg_image = Image.fromarray(bg_image)
+
+    if fg_image.size != bg_image.size:
+        raise ValueError(
+            f"Image sizes not match {fg_image.size} v.s. {bg_image.size}."
+        )
+
+    fg = fg_image.convert("RGBA")
+    bg = bg_image.convert("RGBA")
+
+    return Image.alpha_composite(bg, fg)
+
+
+def check_object_edge_truncated(
+    mask: np.ndarray, edge_threshold: int = 5
+) -> bool:
+    """Checks if a binary object mask is truncated at the image edges.
+
+    Args:
+        mask: A 2D binary NumPy array where nonzero values indicate the object region.
+        edge_threshold: Number of pixels from each image edge to consider for truncation.
+            Defaults to 5.
+
+    Returns:
+        True if the object is fully enclosed (not truncated).
+        False if the object touches or crosses any image boundary.
+    """
+    top = mask[:edge_threshold, :].any()
+    bottom = mask[-edge_threshold:, :].any()
+    left = mask[:, :edge_threshold].any()
+    right = mask[:, -edge_threshold:].any()
+
+    return not (top or bottom or left or right)
+
+
+def vcat_pil_images(
+    images: list[Image.Image], image_mode: str = "RGB"
+) -> Image.Image:
+    widths, heights = zip(*(img.size for img in images))
+    total_height = sum(heights)
+    max_width = max(widths)
+    new_image = Image.new(image_mode, (max_width, total_height))
+    y_offset = 0
+    for image in images:
+        new_image.paste(image, (0, y_offset))
+        y_offset += image.size[1]
+
+    return new_image
+
+
+if __name__ == "__main__":
+    image_paths = [
+        "outputs/layouts_sim/task_0000/images/pen.png",
+        "outputs/layouts_sim/task_0000/images/notebook.png",
+        "outputs/layouts_sim/task_0000/images/mug.png",
+        "outputs/layouts_sim/task_0000/images/lamp.png",
+        "outputs/layouts_sim2/task_0014/images/cloth.png",  # TODO
+    ]
+    for image_path in image_paths:
+        image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
+        mask = image[..., -1]
+        flag = check_object_edge_truncated(mask)
+        print(flag, image_path)

embodied_gen/utils/simulation.py

@@ -0,0 +1,667 @@
+# 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.
+
+import json
+import os
+import xml.etree.ElementTree as ET
+from collections import defaultdict
+from typing import Literal
+
+import mplib
+import numpy as np
+import sapien.core as sapien
+import sapien.physx as physx
+import torch
+from mani_skill.agents.base_agent import BaseAgent
+from mani_skill.envs.scene import ManiSkillScene
+from mani_skill.examples.motionplanning.panda.utils import (
+    compute_grasp_info_by_obb,
+)
+from mani_skill.utils.geometry.trimesh_utils import get_component_mesh
+from PIL import Image, ImageColor
+from scipy.spatial.transform import Rotation as R
+from embodied_gen.data.utils import DiffrastRender
+from embodied_gen.utils.enum import LayoutInfo, Scene3DItemEnum
+from embodied_gen.utils.geometry import quaternion_multiply
+from embodied_gen.utils.log import logger
+
+COLORMAP = list(set(ImageColor.colormap.values()))
+COLOR_PALETTE = np.array(
+    [ImageColor.getrgb(c) for c in COLORMAP], dtype=np.uint8
+)
+SIM_COORD_ALIGN = np.array(
+    [
+        [1.0, 0.0, 0.0, 0.0],
+        [0.0, -1.0, 0.0, 0.0],
+        [0.0, 0.0, -1.0, 0.0],
+        [0.0, 0.0, 0.0, 1.0],
+    ]
+)  # Used to align SAPIEN, MuJoCo coordinate system with the world coordinate system
+
+__all__ = [
+    "SIM_COORD_ALIGN",
+    "FrankaPandaGrasper",
+    "load_assets_from_layout_file",
+    "load_mani_skill_robot",
+    "render_images",
+]
+
+
+def load_actor_from_urdf(
+    scene: sapien.Scene | ManiSkillScene,
+    file_path: str,
+    pose: sapien.Pose | None = None,
+    env_idx: int = None,
+    use_static: bool = False,
+    update_mass: bool = False,
+    scale: float | np.ndarray = 1.0,
+) -> sapien.pysapien.Entity:
+    def _get_local_pose(origin_tag: ET.Element | None) -> sapien.Pose:
+        local_pose = sapien.Pose(p=[0, 0, 0], q=[1, 0, 0, 0])
+        if origin_tag is not None:
+            xyz = list(map(float, origin_tag.get("xyz", "0 0 0").split()))
+            rpy = list(map(float, origin_tag.get("rpy", "0 0 0").split()))
+            qx, qy, qz, qw = R.from_euler("xyz", rpy, degrees=False).as_quat()
+            local_pose = sapien.Pose(p=xyz, q=[qw, qx, qy, qz])
+
+        return local_pose
+
+    tree = ET.parse(file_path)
+    root = tree.getroot()
+    node_name = root.get("name")
+    file_dir = os.path.dirname(file_path)
+
+    visual_mesh = root.find(".//visual/geometry/mesh")
+    visual_file = visual_mesh.get("filename")
+    visual_scale = visual_mesh.get("scale", "1.0 1.0 1.0")
+    visual_scale = np.array([float(x) for x in visual_scale.split()])
+    visual_scale *= np.array(scale)
+
+    collision_mesh = root.find(".//collision/geometry/mesh")
+    collision_file = collision_mesh.get("filename")
+    collision_scale = collision_mesh.get("scale", "1.0 1.0 1.0")
+    collision_scale = np.array([float(x) for x in collision_scale.split()])
+    collision_scale *= np.array(scale)
+
+    visual_pose = _get_local_pose(root.find(".//visual/origin"))
+    collision_pose = _get_local_pose(root.find(".//collision/origin"))
+
+    visual_file = os.path.join(file_dir, visual_file)
+    collision_file = os.path.join(file_dir, collision_file)
+    static_fric = root.find(".//collision/gazebo/mu1").text
+    dynamic_fric = root.find(".//collision/gazebo/mu2").text
+
+    material = physx.PhysxMaterial(
+        static_friction=np.clip(float(static_fric), 0.1, 0.7),
+        dynamic_friction=np.clip(float(dynamic_fric), 0.1, 0.6),
+        restitution=0.05,
+    )
+    builder = scene.create_actor_builder()
+
+    body_type = "static" if use_static else "dynamic"
+    builder.set_physx_body_type(body_type)
+    builder.add_multiple_convex_collisions_from_file(
+        collision_file,
+        material=material,
+        scale=collision_scale,
+        # decomposition="coacd",
+        # decomposition_params=dict(
+        #     threshold=0.05, max_convex_hull=64, verbose=False
+        # ),
+        pose=collision_pose,
+    )
+
+    builder.add_visual_from_file(
+        visual_file,
+        scale=visual_scale,
+        pose=visual_pose,
+    )
+    if pose is None:
+        pose = sapien.Pose(p=[0, 0, 0], q=[1, 0, 0, 0])
+
+    builder.set_initial_pose(pose)
+    if isinstance(scene, ManiSkillScene) and env_idx is not None:
+        builder.set_scene_idxs([env_idx])
+
+    actor = builder.build(
+        name=node_name if env_idx is None else f"{node_name}-{env_idx}"
+    )
+
+    if update_mass and hasattr(actor.components[1], "mass"):
+        node_mass = float(root.find(".//inertial/mass").get("value"))
+        actor.components[1].set_mass(node_mass)
+
+    return actor
+
+
+def load_assets_from_layout_file(
+    scene: ManiSkillScene | sapien.Scene,
+    layout: str,
+    z_offset: float = 0.0,
+    init_quat: list[float] = [0, 0, 0, 1],
+    env_idx: int = None,
+) -> dict[str, sapien.pysapien.Entity]:
+    """Load assets from `EmbodiedGen` layout-gen output and create actors in the scene.
+
+    Args:
+        scene (sapien.Scene | ManiSkillScene): The SAPIEN or ManiSkill scene to load assets into.
+        layout (str): The layout file path.
+        z_offset (float): Offset to apply to the Z-coordinate of non-context objects.
+        init_quat (List[float]): Initial quaternion (x, y, z, w) for orientation adjustment.
+        env_idx (int): Environment index for multi-environment setup.
+    """
+    asset_root = os.path.dirname(layout)
+    layout = LayoutInfo.from_dict(json.load(open(layout, "r")))
+    actors = dict()
+    for node in layout.assets:
+        file_dir = layout.assets[node]
+        file_name = f"{node.replace(' ', '_')}.urdf"
+        urdf_file = os.path.join(asset_root, file_dir, file_name)
+
+        if layout.objs_mapping[node] == Scene3DItemEnum.BACKGROUND.value:
+            continue
+
+        position = layout.position[node].copy()
+        if layout.objs_mapping[node] != Scene3DItemEnum.CONTEXT.value:
+            position[2] += z_offset
+
+        use_static = (
+            layout.relation.get(Scene3DItemEnum.CONTEXT.value, None) == node
+        )
+
+        # Combine initial quaternion with object quaternion
+        x, y, z, qx, qy, qz, qw = position
+        qx, qy, qz, qw = quaternion_multiply([qx, qy, qz, qw], init_quat)
+        actor = load_actor_from_urdf(
+            scene,
+            urdf_file,
+            sapien.Pose(p=[x, y, z], q=[qw, qx, qy, qz]),
+            env_idx,
+            use_static=use_static,
+            update_mass=False,
+        )
+        actors[node] = actor
+
+    return actors
+
+
+def load_mani_skill_robot(
+    scene: sapien.Scene | ManiSkillScene,
+    layout: LayoutInfo | str,
+    control_freq: int = 20,
+    robot_init_qpos_noise: float = 0.0,
+    control_mode: str = "pd_joint_pos",
+    backend_str: tuple[str, str] = ("cpu", "gpu"),
+) -> BaseAgent:
+    from mani_skill.agents import REGISTERED_AGENTS
+    from mani_skill.envs.scene import ManiSkillScene
+    from mani_skill.envs.utils.system.backend import (
+        parse_sim_and_render_backend,
+    )
+
+    if isinstance(layout, str) and layout.endswith(".json"):
+        layout = LayoutInfo.from_dict(json.load(open(layout, "r")))
+
+    robot_name = layout.relation[Scene3DItemEnum.ROBOT.value]
+    x, y, z, qx, qy, qz, qw = layout.position[robot_name]
+    delta_z = 0.002  # Add small offset to avoid collision.
+    pose = sapien.Pose([x, y, z + delta_z], [qw, qx, qy, qz])
+
+    if robot_name not in REGISTERED_AGENTS:
+        logger.warning(
+            f"Robot `{robot_name}` not registered, chosen from {REGISTERED_AGENTS.keys()}, use `panda` instead."
+        )
+        robot_name = "panda"
+
+    ROBOT_CLS = REGISTERED_AGENTS[robot_name].agent_cls
+    backend = parse_sim_and_render_backend(*backend_str)
+    if isinstance(scene, sapien.Scene):
+        scene = ManiSkillScene([scene], device=backend_str[0], backend=backend)
+    robot = ROBOT_CLS(
+        scene=scene,
+        control_freq=control_freq,
+        control_mode=control_mode,
+        initial_pose=pose,
+    )
+
+    # Set robot init joint rad agree(joint0 to joint6 w 2 finger).
+    qpos = np.array(
+        [
+            0.0,
+            np.pi / 8,
+            0,
+            -np.pi * 3 / 8,
+            0,
+            np.pi * 3 / 4,
+            np.pi / 4,
+            0.04,
+            0.04,
+        ]
+    )
+    qpos = (
+        np.random.normal(
+            0, robot_init_qpos_noise, (len(scene.sub_scenes), len(qpos))
+        )
+        + qpos
+    )
+    qpos[:, -2:] = 0.04
+    robot.reset(qpos)
+    robot.init_qpos = robot.robot.qpos
+    robot.controller.controllers["gripper"].reset()
+
+    return robot
+
+
+def render_images(
+    camera: sapien.render.RenderCameraComponent,
+    render_keys: list[
+        Literal[
+            "Color",
+            "Segmentation",
+            "Normal",
+            "Mask",
+            "Depth",
+            "Foreground",
+        ]
+    ] = None,
+) -> dict[str, Image.Image]:
+    """Render images from a given sapien camera.
+
+    Args:
+        camera (sapien.render.RenderCameraComponent): The camera to render from.
+        render_keys (List[str]): Types of images to render (e.g., Color, Segmentation).
+
+    Returns:
+        Dict[str, Image.Image]: Dictionary of rendered images.
+    """
+    if render_keys is None:
+        render_keys = [
+            "Color",
+            "Segmentation",
+            "Normal",
+            "Mask",
+            "Depth",
+            "Foreground",
+        ]
+
+    results: dict[str, Image.Image] = {}
+    if "Color" in render_keys:
+        color = camera.get_picture("Color")
+        color_rgb = (np.clip(color[..., :3], 0, 1) * 255).astype(np.uint8)
+        results["Color"] = Image.fromarray(color_rgb)
+
+    if "Mask" in render_keys:
+        alpha = (np.clip(color[..., 3], 0, 1) * 255).astype(np.uint8)
+        results["Mask"] = Image.fromarray(alpha)
+
+    if "Segmentation" in render_keys:
+        seg_labels = camera.get_picture("Segmentation")
+        label0 = seg_labels[..., 0].astype(np.uint8)
+        seg_color = COLOR_PALETTE[label0]
+        results["Segmentation"] = Image.fromarray(seg_color)
+
+    if "Foreground" in render_keys:
+        seg_labels = camera.get_picture("Segmentation")
+        label0 = seg_labels[..., 0]
+        mask = np.where((label0 > 1), 255, 0).astype(np.uint8)
+        color = camera.get_picture("Color")
+        color_rgb = (np.clip(color[..., :3], 0, 1) * 255).astype(np.uint8)
+        foreground = np.concatenate([color_rgb, mask[..., None]], axis=-1)
+        results["Foreground"] = Image.fromarray(foreground)
+
+    if "Normal" in render_keys:
+        normal = camera.get_picture("Normal")[..., :3]
+        normal_img = (((normal + 1) / 2) * 255).astype(np.uint8)
+        results["Normal"] = Image.fromarray(normal_img)
+
+    if "Depth" in render_keys:
+        position_map = camera.get_picture("Position")
+        depth = -position_map[..., 2]
+        alpha = torch.tensor(color[..., 3], dtype=torch.float32)
+        norm_depth = DiffrastRender.normalize_map_by_mask(
+            torch.tensor(depth), alpha
+        )
+        depth_img = (norm_depth * 255).to(torch.uint8).numpy()
+        results["Depth"] = Image.fromarray(depth_img)
+
+    return results
+
+
+class SapienSceneManager:
+    """A class to manage SAPIEN simulator."""
+
+    def __init__(
+        self, sim_freq: int, ray_tracing: bool, device: str = "cuda"
+    ) -> None:
+        self.sim_freq = sim_freq
+        self.ray_tracing = ray_tracing
+        self.device = device
+        self.renderer = sapien.SapienRenderer()
+        self.scene = self._setup_scene()
+        self.cameras: list[sapien.render.RenderCameraComponent] = []
+        self.actors: dict[str, sapien.pysapien.Entity] = {}
+
+    def _setup_scene(self) -> sapien.Scene:
+        """Set up the SAPIEN scene with lighting and ground."""
+        # Ray tracing settings
+        if self.ray_tracing:
+            sapien.render.set_camera_shader_dir("rt")
+            sapien.render.set_ray_tracing_samples_per_pixel(64)
+            sapien.render.set_ray_tracing_path_depth(10)
+            sapien.render.set_ray_tracing_denoiser("oidn")
+
+        scene = sapien.Scene()
+        scene.set_timestep(1 / self.sim_freq)
+
+        # Add lighting
+        scene.set_ambient_light([0.2, 0.2, 0.2])
+        scene.add_directional_light(
+            direction=[0, 1, -1],
+            color=[1.5, 1.45, 1.4],
+            shadow=True,
+            shadow_map_size=2048,
+        )
+        scene.add_directional_light(
+            direction=[0, -0.5, 1], color=[0.8, 0.8, 0.85], shadow=False
+        )
+        scene.add_directional_light(
+            direction=[0, -1, 1], color=[1.0, 1.0, 1.0], shadow=False
+        )
+
+        ground_material = self.renderer.create_material()
+        ground_material.base_color = [0.5, 0.5, 0.5, 1]  # rgba, gray
+        ground_material.roughness = 0.7
+        ground_material.metallic = 0.0
+        scene.add_ground(0, render_material=ground_material)
+
+        return scene
+
+    def step_action(
+        self,
+        agent: BaseAgent,
+        action: torch.Tensor,
+        cameras: list[sapien.render.RenderCameraComponent],
+        render_keys: list[str],
+        sim_steps_per_control: int = 1,
+    ) -> dict:
+        agent.set_action(action)
+        frames = defaultdict(list)
+        for _ in range(sim_steps_per_control):
+            self.scene.step()
+
+        self.scene.update_render()
+        for camera in cameras:
+            camera.take_picture()
+            images = render_images(camera, render_keys=render_keys)
+            frames[camera.name].append(images)
+
+        return frames
+
+    def create_camera(
+        self,
+        cam_name: str,
+        pose: sapien.Pose,
+        image_hw: tuple[int, int],
+        fovy_deg: float,
+    ) -> sapien.render.RenderCameraComponent:
+        """Create a single camera in the scene.
+
+        Args:
+            cam_name (str): Name of the camera.
+            pose (sapien.Pose): Camera pose p=(x, y, z), q=(w, x, y, z)
+            image_hw (Tuple[int, int]): Image resolution (height, width) for cameras.
+            fovy_deg (float): Field of view in degrees for cameras.
+
+        Returns:
+            sapien.render.RenderCameraComponent: The created camera.
+        """
+        cam_actor = self.scene.create_actor_builder().build_kinematic()
+        cam_actor.set_pose(pose)
+        camera = self.scene.add_mounted_camera(
+            name=cam_name,
+            mount=cam_actor,
+            pose=sapien.Pose(p=[0, 0, 0], q=[1, 0, 0, 0]),
+            width=image_hw[1],
+            height=image_hw[0],
+            fovy=np.deg2rad(fovy_deg),
+            near=0.01,
+            far=100,
+        )
+        self.cameras.append(camera)
+
+        return camera
+
+    def initialize_circular_cameras(
+        self,
+        num_cameras: int,
+        radius: float,
+        height: float,
+        target_pt: list[float],
+        image_hw: tuple[int, int],
+        fovy_deg: float,
+    ) -> list[sapien.render.RenderCameraComponent]:
+        """Initialize multiple cameras arranged in a circle.
+
+        Args:
+            num_cameras (int): Number of cameras to create.
+            radius (float): Radius of the camera circle.
+            height (float): Fixed Z-coordinate of the cameras.
+            target_pt (list[float]): 3D point (x, y, z) that cameras look at.
+            image_hw (Tuple[int, int]): Image resolution (height, width) for cameras.
+            fovy_deg (float): Field of view in degrees for cameras.
+
+        Returns:
+            List[sapien.render.RenderCameraComponent]: List of created cameras.
+        """
+        angle_step = 2 * np.pi / num_cameras
+        world_up_vec = np.array([0.0, 0.0, 1.0])
+        target_pt = np.array(target_pt)
+
+        for i in range(num_cameras):
+            angle = i * angle_step
+            cam_x = radius * np.cos(angle)
+            cam_y = radius * np.sin(angle)
+            cam_z = height
+            eye_pos = [cam_x, cam_y, cam_z]
+
+            forward_vec = target_pt - eye_pos
+            forward_vec = forward_vec / np.linalg.norm(forward_vec)
+            temp_right_vec = np.cross(forward_vec, world_up_vec)
+
+            if np.linalg.norm(temp_right_vec) < 1e-6:
+                temp_right_vec = np.array([1.0, 0.0, 0.0])
+                if np.abs(np.dot(temp_right_vec, forward_vec)) > 0.99:
+                    temp_right_vec = np.array([0.0, 1.0, 0.0])
+
+            right_vec = temp_right_vec / np.linalg.norm(temp_right_vec)
+            up_vec = np.cross(right_vec, forward_vec)
+            rotation_matrix = np.array([forward_vec, -right_vec, up_vec]).T
+
+            rot = R.from_matrix(rotation_matrix)
+            scipy_quat = rot.as_quat()  # (x, y, z, w)
+            quat = [
+                scipy_quat[3],
+                scipy_quat[0],
+                scipy_quat[1],
+                scipy_quat[2],
+            ]  # (w, x, y, z)
+
+            self.create_camera(
+                f"camera_{i}",
+                sapien.Pose(p=eye_pos, q=quat),
+                image_hw,
+                fovy_deg,
+            )
+
+        return self.cameras
+
+
+class FrankaPandaGrasper(object):
+    def __init__(
+        self,
+        agent: BaseAgent,
+        control_freq: float,
+        joint_vel_limits: float = 2.0,
+        joint_acc_limits: float = 1.0,
+        finger_length: float = 0.025,
+    ) -> None:
+        self.agent = agent
+        self.robot = agent.robot
+        self.control_freq = control_freq
+        self.control_timestep = 1 / control_freq
+        self.joint_vel_limits = joint_vel_limits
+        self.joint_acc_limits = joint_acc_limits
+        self.finger_length = finger_length
+        self.planners = self._setup_planner()
+
+    def _setup_planner(self) -> mplib.Planner:
+        planners = []
+        for pose in self.robot.pose:
+            link_names = [link.get_name() for link in self.robot.get_links()]
+            joint_names = [
+                joint.get_name() for joint in self.robot.get_active_joints()
+            ]
+            planner = mplib.Planner(
+                urdf=self.agent.urdf_path,
+                srdf=self.agent.urdf_path.replace(".urdf", ".srdf"),
+                user_link_names=link_names,
+                user_joint_names=joint_names,
+                move_group="panda_hand_tcp",
+                joint_vel_limits=np.ones(7) * self.joint_vel_limits,
+                joint_acc_limits=np.ones(7) * self.joint_acc_limits,
+            )
+            planner.set_base_pose(pose.raw_pose[0].tolist())
+            planners.append(planner)
+
+        return planners
+
+    def control_gripper(
+        self,
+        gripper_state: Literal[-1, 1],
+        n_step: int = 10,
+    ) -> np.ndarray:
+        qpos = self.robot.get_qpos()[0, :-2].cpu().numpy()
+        actions = []
+        for _ in range(n_step):
+            action = np.hstack([qpos, gripper_state])[None, ...]
+            actions.append(action)
+
+        return np.concatenate(actions, axis=0)
+
+    def move_to_pose(
+        self,
+        pose: sapien.Pose,
+        control_timestep: float,
+        gripper_state: Literal[-1, 1],
+        use_point_cloud: bool = False,
+        n_max_step: int = 100,
+        action_key: str = "position",
+        env_idx: int = 0,
+    ) -> np.ndarray:
+        result = self.planners[env_idx].plan_qpos_to_pose(
+            np.concatenate([pose.p, pose.q]),
+            self.robot.get_qpos().cpu().numpy()[0],
+            time_step=control_timestep,
+            use_point_cloud=use_point_cloud,
+        )
+
+        if result["status"] != "Success":
+            result = self.planners[env_idx].plan_screw(
+                np.concatenate([pose.p, pose.q]),
+                self.robot.get_qpos().cpu().numpy()[0],
+                time_step=control_timestep,
+                use_point_cloud=use_point_cloud,
+            )
+
+        if result["status"] != "Success":
+            return
+
+        sample_ratio = (len(result[action_key]) // n_max_step) + 1
+        result[action_key] = result[action_key][::sample_ratio]
+
+        n_step = len(result[action_key])
+        actions = []
+        for i in range(n_step):
+            qpos = result[action_key][i]
+            action = np.hstack([qpos, gripper_state])[None, ...]
+            actions.append(action)
+
+        return np.concatenate(actions, axis=0)
+
+    def compute_grasp_action(
+        self,
+        actor: sapien.pysapien.Entity,
+        reach_target_only: bool = True,
+        offset: tuple[float, float, float] = [0, 0, -0.05],
+        env_idx: int = 0,
+    ) -> np.ndarray:
+        physx_rigid = actor.components[1]
+        mesh = get_component_mesh(physx_rigid, to_world_frame=True)
+        obb = mesh.bounding_box_oriented
+        approaching = np.array([0, 0, -1])
+        tcp_pose = self.agent.tcp.pose[env_idx]
+        target_closing = (
+            tcp_pose.to_transformation_matrix()[0, :3, 1].cpu().numpy()
+        )
+        grasp_info = compute_grasp_info_by_obb(
+            obb,
+            approaching=approaching,
+            target_closing=target_closing,
+            depth=self.finger_length,
+        )
+
+        closing, center = grasp_info["closing"], grasp_info["center"]
+        raw_tcp_pose = tcp_pose.sp
+        grasp_pose = self.agent.build_grasp_pose(approaching, closing, center)
+        reach_pose = grasp_pose * sapien.Pose(p=offset)
+        grasp_pose = grasp_pose * sapien.Pose(p=[0, 0, 0.01])
+        actions = []
+        reach_actions = self.move_to_pose(
+            reach_pose,
+            self.control_timestep,
+            gripper_state=1,
+            env_idx=env_idx,
+        )
+        actions.append(reach_actions)
+
+        if reach_actions is None:
+            logger.warning(
+                f"Failed to reach the grasp pose for node `{actor.name}`, skipping grasping."
+            )
+            return None
+
+        if not reach_target_only:
+            grasp_actions = self.move_to_pose(
+                grasp_pose,
+                self.control_timestep,
+                gripper_state=1,
+                env_idx=env_idx,
+            )
+            actions.append(grasp_actions)
+            close_actions = self.control_gripper(
+                gripper_state=-1,
+                env_idx=env_idx,
+            )
+            actions.append(close_actions)
+            back_actions = self.move_to_pose(
+                raw_tcp_pose,
+                self.control_timestep,
+                gripper_state=-1,
+                env_idx=env_idx,
+            )
+            actions.append(back_actions)
+
+        return np.concatenate(actions, axis=0)

embodied_gen/utils/tags.py

@@ -0,0 +1 @@
+VERSION = "v0.1.5"

embodied_gen/utils/trender.py

@@ -0,0 +1,90 @@
+# 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.
+
+import os
+import sys
+
+import numpy as np
+import spaces
+import torch
+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.utils.render_utils import (
+    render_frames,
+    yaw_pitch_r_fov_to_extrinsics_intrinsics,
+)
+
+__all__ = [
+    "render_video",
+]
+
+
+@spaces.GPU
+def render_mesh(sample, extrinsics, intrinsics, options={}, **kwargs):
+    renderer = MeshRenderer()
+    renderer.rendering_options.resolution = options.get("resolution", 512)
+    renderer.rendering_options.near = options.get("near", 1)
+    renderer.rendering_options.far = options.get("far", 100)
+    renderer.rendering_options.ssaa = options.get("ssaa", 4)
+    rets = {}
+    for extr, intr in tqdm(zip(extrinsics, intrinsics), desc="Rendering"):
+        res = renderer.render(sample, extr, intr)
+        if "normal" not in rets:
+            rets["normal"] = []
+        normal = torch.lerp(
+            torch.zeros_like(res["normal"]), res["normal"], res["mask"]
+        )
+        normal = np.clip(
+            normal.detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255
+        ).astype(np.uint8)
+        rets["normal"].append(normal)
+
+    return rets
+
+
+@spaces.GPU
+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)
+    yaws = yaws.tolist()
+    pitch = [0.5] * num_frames
+    extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(
+        yaws, pitch, r, fov
+    )
+    render_fn = (
+        render_mesh if isinstance(sample, MeshExtractResult) else render_frames
+    )
+    result = render_fn(
+        sample,
+        extrinsics,
+        intrinsics,
+        {"resolution": resolution, "bg_color": bg_color},
+        **kwargs,
+    )
+
+    return result

embodied_gen/validators/aesthetic_predictor.py

@@ -0,0 +1,137 @@
+# 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.
+
+
+import os
+
+import clip
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+from huggingface_hub import snapshot_download
+from PIL import Image
+
+
+class AestheticPredictor:
+    """Aesthetic Score Predictor.
+
+    Checkpoints from https://github.com/christophschuhmann/improved-aesthetic-predictor/tree/main
+
+    Args:
+        clip_model_dir (str): Path to the directory of the CLIP model.
+        sac_model_path (str): Path to the pre-trained SAC model.
+        device (str): Device to use for computation ("cuda" or "cpu").
+    """
+
+    def __init__(self, clip_model_dir=None, sac_model_path=None, device="cpu"):
+
+        self.device = device
+
+        if clip_model_dir is None:
+            model_path = snapshot_download(
+                repo_id="xinjjj/RoboAssetGen", allow_patterns="aesthetic/*"
+            )
+            suffix = "aesthetic"
+            model_path = snapshot_download(
+                repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
+            )
+            clip_model_dir = os.path.join(model_path, suffix)
+
+        if sac_model_path is None:
+            model_path = snapshot_download(
+                repo_id="xinjjj/RoboAssetGen", allow_patterns="aesthetic/*"
+            )
+            suffix = "aesthetic"
+            model_path = snapshot_download(
+                repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
+            )
+            sac_model_path = os.path.join(
+                model_path, suffix, "sac+logos+ava1-l14-linearMSE.pth"
+            )
+
+        self.clip_model, self.preprocess = self._load_clip_model(
+            clip_model_dir
+        )
+        self.sac_model = self._load_sac_model(sac_model_path, input_size=768)
+
+    class MLP(pl.LightningModule):  # noqa
+        def __init__(self, input_size):
+            super().__init__()
+            self.layers = nn.Sequential(
+                nn.Linear(input_size, 1024),
+                nn.Dropout(0.2),
+                nn.Linear(1024, 128),
+                nn.Dropout(0.2),
+                nn.Linear(128, 64),
+                nn.Dropout(0.1),
+                nn.Linear(64, 16),
+                nn.Linear(16, 1),
+            )
+
+        def forward(self, x):
+            return self.layers(x)
+
+    @staticmethod
+    def normalized(a, axis=-1, order=2):
+        """Normalize the array to unit norm."""
+        l2 = np.atleast_1d(np.linalg.norm(a, order, axis))
+        l2[l2 == 0] = 1
+        return a / np.expand_dims(l2, axis)
+
+    def _load_clip_model(self, model_dir: str, model_name: str = "ViT-L/14"):
+        """Load the CLIP model."""
+        model, preprocess = clip.load(
+            model_name, download_root=model_dir, device=self.device
+        )
+        return model, preprocess
+
+    def _load_sac_model(self, model_path, input_size):
+        """Load the SAC model."""
+        model = self.MLP(input_size)
+        ckpt = torch.load(model_path, weights_only=True)
+        model.load_state_dict(ckpt)
+        model.to(self.device)
+        model.eval()
+        return model
+
+    def predict(self, image_path):
+        """Predict the aesthetic score for a given image.
+
+        Args:
+            image_path (str): Path to the image file.
+
+        Returns:
+            float: Predicted aesthetic score.
+        """
+        pil_image = Image.open(image_path)
+        image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            # Extract CLIP features
+            image_features = self.clip_model.encode_image(image)
+            # Normalize features
+            normalized_features = self.normalized(
+                image_features.cpu().detach().numpy()
+            )
+            # Predict score
+            prediction = self.sac_model(
+                torch.from_numpy(normalized_features)
+                .type(torch.FloatTensor)
+                .to(self.device)
+            )
+
+        return prediction.item()