Initial clean commit

.gitattributes

@@ -0,0 +1,48 @@
+*.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
+assets/example_image/sample_01.jpg filter=lfs diff=lfs merge=lfs -text
+assets/example_image/sample_02.jpg filter=lfs diff=lfs merge=lfs -text
+assets/example_image/sample_03.jpg filter=lfs diff=lfs merge=lfs -text
+assets/example_image/sample_04.jpg filter=lfs diff=lfs merge=lfs -text
+assets/example_image/sample_05.jpg filter=lfs diff=lfs merge=lfs -text
+assets/example_image/sample_06.jpg filter=lfs diff=lfs merge=lfs -text
+assets/example_image/sample_07.jpg filter=lfs diff=lfs merge=lfs -text
+assets/example_image/sample_08.jpg filter=lfs diff=lfs merge=lfs -text
+assets/example_image/sample_11.jpg filter=lfs diff=lfs merge=lfs -text
+assets/example_image/sample_14.jpg filter=lfs diff=lfs merge=lfs -text
+assets/example_image/sample_20.jpg filter=lfs diff=lfs merge=lfs -text
+assets/logo.png filter=lfs diff=lfs merge=lfs -text
+assets/overall.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,15 @@
+---
+title: EmbodiedGen Image To 3D
+emoji: 🖼️
+colorFrom: blue
+colorTo: red
+sdk: gradio
+sdk_version: 5.33.1
+app_file: app.py
+pinned: false
+license: apache-2.0
+short_description: Generate physically plausible 3D model from single image.
+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,498 @@
+# 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"] = "imageto3d"
+from glob import glob
+
+import gradio as gr
+from app_style import custom_theme, image_css, lighting_css
+from common import (
+    MAX_SEED,
+    VERSION,
+    active_btn_by_content,
+    end_session,
+    extract_3d_representations_v2,
+    extract_urdf,
+    get_seed,
+    image_to_3d,
+    preprocess_image_fn,
+    preprocess_sam_image_fn,
+    select_point,
+    start_session,
+)
+
+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***: Image-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>
+
+        🖼️ Generate physically plausible 3D asset from single input image.
+        """.format(
+            VERSION=VERSION
+        ),
+        elem_classes=["header"],
+    )
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            with gr.Tabs() as input_tabs:
+                with gr.Tab(
+                    label="Image(auto seg)", id=0
+                ) as single_image_input_tab:
+                    raw_image_cache = gr.Image(
+                        format="png",
+                        image_mode="RGB",
+                        type="pil",
+                        visible=False,
+                    )
+                    image_prompt = gr.Image(
+                        label="Input Image",
+                        format="png",
+                        image_mode="RGBA",
+                        type="pil",
+                        height=400,
+                        elem_classes=["image_fit"],
+                    )
+                    gr.Markdown(
+                        """
+                        If you are not satisfied with the auto segmentation
+                        result, please switch to the `Image(SAM seg)` tab."""
+                    )
+                with gr.Tab(
+                    label="Image(SAM seg)", id=1
+                ) as samimage_input_tab:
+                    with gr.Row():
+                        with gr.Column(scale=1):
+                            image_prompt_sam = gr.Image(
+                                label="Input Image",
+                                type="numpy",
+                                height=400,
+                                elem_classes=["image_fit"],
+                            )
+                            image_seg_sam = gr.Image(
+                                label="SAM Seg Image",
+                                image_mode="RGBA",
+                                type="pil",
+                                height=400,
+                                visible=False,
+                            )
+                        with gr.Column(scale=1):
+                            image_mask_sam = gr.AnnotatedImage(
+                                elem_classes=["image_fit"]
+                            )
+
+                    fg_bg_radio = gr.Radio(
+                        ["foreground_point", "background_point"],
+                        label="Select foreground(green) or background(red) points, by default foreground",  # noqa
+                        value="foreground_point",
+                    )
+                    gr.Markdown(
+                        """ Click the `Input Image` to select SAM points,
+                        after get the satisified segmentation, click `Generate`
+                         button to generate the 3D asset. \n
+                        Note: If the segmented foreground is too small relative
+                         to the entire image area, the generation will fail.
+                    """
+                    )
+
+            with gr.Accordion(label="Generation Settings", open=False):
+                with gr.Row():
+                    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,
+                    )
+                    rmbg_tag = gr.Radio(
+                        choices=["rembg", "rmbg14"],
+                        value="rembg",
+                        label="Background Removal Model",
+                    )
+                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(
+                "🚀 1. Generate(~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(
+                    "🔍 2. Extract 3D Representation(~2 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(
+                    "🧩 3. Extract URDF with physics(~1 mins)",
+                    variant="primary",
+                    interactive=False,
+                )
+            with gr.Row():
+                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
+                )
+            with gr.Row():
+                download_urdf = gr.DownloadButton(
+                    label="⬇️ 4. Download URDF",
+                    variant="primary",
+                    interactive=False,
+                )
+
+            gr.Markdown(
+                """ NOTE: If `Asset Attributes` are provided, it will guide
+                GPT to perform physical attributes restoration. \n
+                The `Download URDF` file is restored to the real scale and
+                has quality inspection, open with an editor to view details.
+            """
+            )
+
+            with gr.Row() as single_image_example:
+                examples = gr.Examples(
+                    label="Image Gallery",
+                    examples=[
+                        [image_path]
+                        for image_path in sorted(
+                            glob("assets/example_image/*")
+                        )
+                    ],
+                    inputs=[image_prompt, rmbg_tag],
+                    fn=preprocess_image_fn,
+                    outputs=[image_prompt, raw_image_cache],
+                    run_on_click=True,
+                    examples_per_page=10,
+                )
+
+            with gr.Row(visible=False) as single_sam_image_example:
+                examples = gr.Examples(
+                    label="Image Gallery",
+                    examples=[
+                        [image_path]
+                        for image_path in sorted(
+                            glob("assets/example_image/*")
+                        )
+                    ],
+                    inputs=[image_prompt_sam],
+                    fn=preprocess_sam_image_fn,
+                    outputs=[image_prompt_sam, raw_image_cache],
+                    run_on_click=True,
+                    examples_per_page=10,
+                )
+        with gr.Column(scale=1):
+            gr.Markdown("<br>")
+            video_output = gr.Video(
+                label="Generated 3D Asset",
+                autoplay=True,
+                loop=True,
+                height=300,
+            )
+            model_output_gs = gr.Model3D(
+                label="Gaussian Representation", height=300, interactive=False
+            )
+            aligned_gs = gr.Textbox(visible=False)
+            gr.Markdown(
+                """ The rendering of `Gaussian Representation` takes additional 10s. """  # noqa
+            )
+            with gr.Row():
+                model_output_mesh = gr.Model3D(
+                    label="Mesh Representation",
+                    height=300,
+                    interactive=False,
+                    clear_color=[0.8, 0.8, 0.8, 1],
+                    elem_id="lighter_mesh",
+                )
+
+    is_samimage = gr.State(False)
+    output_buf = gr.State()
+    selected_points = gr.State(value=[])
+
+    demo.load(start_session)
+    demo.unload(end_session)
+
+    single_image_input_tab.select(
+        lambda: tuple(
+            [False, gr.Row.update(visible=True), gr.Row.update(visible=False)]
+        ),
+        outputs=[is_samimage, single_image_example, single_sam_image_example],
+    )
+    samimage_input_tab.select(
+        lambda: tuple(
+            [True, gr.Row.update(visible=True), gr.Row.update(visible=False)]
+        ),
+        outputs=[is_samimage, single_sam_image_example, single_image_example],
+    )
+
+    image_prompt.upload(
+        preprocess_image_fn,
+        inputs=[image_prompt, rmbg_tag],
+        outputs=[image_prompt, raw_image_cache],
+    )
+    image_prompt.change(
+        lambda: tuple(
+            [
+                gr.Button(interactive=False),
+                gr.Button(interactive=False),
+                gr.Button(interactive=False),
+                None,
+                "",
+                None,
+                None,
+                "",
+                "",
+                "",
+                "",
+                "",
+                "",
+                "",
+                "",
+            ]
+        ),
+        outputs=[
+            extract_rep3d_btn,
+            extract_urdf_btn,
+            download_urdf,
+            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_prompt.change(
+        active_btn_by_content,
+        inputs=image_prompt,
+        outputs=generate_btn,
+    )
+
+    image_prompt_sam.upload(
+        preprocess_sam_image_fn,
+        inputs=[image_prompt_sam],
+        outputs=[image_prompt_sam, raw_image_cache],
+    )
+    image_prompt_sam.change(
+        lambda: tuple(
+            [
+                gr.Button(interactive=False),
+                gr.Button(interactive=False),
+                gr.Button(interactive=False),
+                None,
+                None,
+                None,
+                "",
+                "",
+                "",
+                "",
+                "",
+                "",
+                "",
+                "",
+                None,
+                [],
+            ]
+        ),
+        outputs=[
+            extract_rep3d_btn,
+            extract_urdf_btn,
+            download_urdf,
+            model_output_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_mask_sam,
+            selected_points,
+        ],
+    )
+
+    image_prompt_sam.select(
+        select_point,
+        [
+            image_prompt_sam,
+            selected_points,
+            fg_bg_radio,
+        ],
+        [image_mask_sam, image_seg_sam],
+    )
+    image_seg_sam.change(
+        active_btn_by_content,
+        inputs=image_seg_sam,
+        outputs=generate_btn,
+    )
+
+    generate_btn.click(
+        get_seed,
+        inputs=[randomize_seed, seed],
+        outputs=[seed],
+    ).success(
+        image_to_3d,
+        inputs=[
+            image_prompt,
+            seed,
+            ss_guidance_strength,
+            ss_sampling_steps,
+            slat_guidance_strength,
+            slat_sampling_steps,
+            raw_image_cache,
+            image_seg_sam,
+            is_samimage,
+        ],
+        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()