test

app.py

@@ -1,21 +1,496 @@
-import torch
+import os
+import shutil
+from functools import partial
+
 import gradio as gr
-import spaces
-import nvdiffrast.torch as dr
+from common import (
+    MAX_SEED,
+    VERSION,
+    TrellisImageTo3DPipeline,
+    active_btn_by_content,
+    extract_3d_representations_v2,
+    extract_urdf,
+    get_seed,
+    image_to_3d,
+    preprocess_image_fn,
+    preprocess_sam_image_fn,
+    select_point,
+)
+from gradio.themes import Default
+from gradio.themes.utils.colors import slate
+from gradio_litmodel3d import LitModel3D
+from asset3d_gen.models.delight import DelightingModel
+from asset3d_gen.models.segment import RembgRemover, SAMPredictor
+from asset3d_gen.models.super_resolution import ImageRealESRGAN
+from asset3d_gen.utils.gpt_clients import GPT_CLIENT
+from asset3d_gen.validators.quality_checkers import (
+    ImageAestheticChecker,
+    ImageSegChecker,
+    MeshGeoChecker,
+)
+from asset3d_gen.validators.urdf_convertor import URDFGenerator
+
+TMP_DIR = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)), "sessions/imageto3d"
+)
+os.makedirs(TMP_DIR, exist_ok=True)
+
+RBG_REMOVER = RembgRemover()
+SAM_PREDICTOR = SAMPredictor(model_type="vit_h")
+DELIGHT = DelightingModel()
+IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
+PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
+    "JeffreyXiang/TRELLIS-image-large"
+)
+# PIPELINE.cuda()
+
+IMAGE_BUFFER = {}
+SEG_CHECKER = ImageSegChecker(GPT_CLIENT)
+GEO_CHECKER = MeshGeoChecker(GPT_CLIENT)
+AESTHETIC_CHECKER = ImageAestheticChecker()
+CHECKERS = [GEO_CHECKER, SEG_CHECKER, AESTHETIC_CHECKER]
+URDF_CONVERTOR = URDFGenerator(GPT_CLIENT, render_view_num=4)
+
+
+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)
+
+
+with gr.Blocks(
+    delete_cache=(43200, 43200), theme=Default(primary_hue=slate)
+) as demo:
+    gr.Markdown(
+        f"""
+        ## Image to 3D Asset Pipeline \n
+        version: {VERSION} \n
+        The service is temporarily deployed on `dev015-10.34.8.82: CUDA 4`.
+    """
+    )
+    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:
+                    image_prompt = gr.Image(
+                        label="Input Image",
+                        format="png",
+                        image_mode="RGBA",
+                        type="pil",
+                        height=300,
+                    )
+                    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
+                            )
+                            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()
+
+                    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
+                    )
+                with gr.Row():
+                    randomize_seed = gr.Checkbox(
+                        label="Randomize Seed", value=False
+                    )
+                    project_delight = gr.Checkbox(
+                        label="Backproject delighting",
+                        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(
+                "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(
+                    "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(
+                    "Extract URDF(~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="Download URDF", variant="primary", interactive=False
+                )
+
+            gr.Markdown(
+                """ NOTE: If `Asset Attributes` are provided, the provided
+                properties will be used; otherwise, the GPT-preset properties
+                will be applied. \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=[
+                        [f"scripts/apps/assets/example_image/{image}"]
+                        for image in os.listdir(
+                            "scripts/apps/assets/example_image"
+                        )
+                    ],
+                    inputs=[image_prompt],
+                    fn=partial(
+                        preprocess_image_fn,
+                        model=RBG_REMOVER,
+                        buffer=IMAGE_BUFFER,
+                    ),
+                    outputs=[image_prompt],
+                    run_on_click=True,
+                    examples_per_page=32,
+                )
+
+            with gr.Row(visible=False) as single_sam_image_example:
+                examples = gr.Examples(
+                    label="Image Gallery",
+                    examples=[
+                        f"scripts/apps/assets/example_image/{image}"
+                        for image in os.listdir(
+                            "scripts/apps/assets/example_image"
+                        )
+                    ],
+                    inputs=[image_prompt_sam],
+                    fn=partial(
+                        preprocess_sam_image_fn,
+                        buffer=IMAGE_BUFFER,
+                        model=SAM_PREDICTOR,
+                    ),
+                    outputs=[image_prompt_sam],
+                    run_on_click=True,
+                    examples_per_page=32,
+                )
+        with gr.Column(scale=1):
+            video_output = gr.Video(
+                label="Generated 3D Asset",
+                autoplay=True,
+                loop=True,
+                height=300,
+            )
+            model_output_gs = LitModel3D(
+                label="Gaussian Representation", height=300, interactive=False
+            )
+            aligned_gs = gr.Textbox(visible=False)
+            with gr.Row():
+                model_output_mesh = LitModel3D(
+                    label="Mesh Representation",
+                    exposure=10.0,
+                    height=300,
+                    interactive=False,
+                )
+            gr.Markdown(
+                """ The rendering of `Gaussian Representation` takes additional 10s. """  # noqa
+            )
+
+    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(
+        partial(preprocess_image_fn, model=RBG_REMOVER, buffer=IMAGE_BUFFER),
+        inputs=[image_prompt],
+        outputs=[image_prompt],
+    )
+    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(
+        partial(
+            preprocess_sam_image_fn, buffer=IMAGE_BUFFER, model=SAM_PREDICTOR
+        ),
+        inputs=[image_prompt_sam],
+        outputs=[image_prompt_sam],
+    )
+    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,
+            gr.State(lambda: SAM_PREDICTOR),
+        ],
+        [image_mask_sam, image_seg_sam],
+    )
+    image_seg_sam.change(
+        active_btn_by_content,
+        inputs=image_seg_sam,
+        outputs=generate_btn,
+    )
 
-@spaces.GPU
-def pinfo():
-    print("CUDA Version:", torch.version.cuda)
-    print(torch.version.cuda)
-    print(torch.cuda.is_available())
-    zero = torch.Tensor([0]).cuda()
-    print(zero.device)
+    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,
+            gr.State(lambda: IMAGE_BUFFER),
+            gr.State(lambda: PIPELINE),
+            gr.State(lambda: TMP_DIR),
+            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,
+            gr.State(lambda: TMP_DIR),
+            gr.State(lambda: DELIGHT),
+            gr.State(lambda: IMAGESR_MODEL),
+        ],
+        outputs=[
+            model_output_mesh,
+            model_output_gs,
+            model_output_obj,
+            aligned_gs,
+        ],
+    ).success(
+        lambda: gr.Button(interactive=True),
+        outputs=[extract_urdf_btn],
+    )
 
-def my_function(input_text):
-    pinfo()
-    return f"Received: {input_text}"
+    extract_urdf_btn.click(
+        extract_urdf,
+        inputs=[
+            aligned_gs,
+            model_output_obj,
+            asset_cat_text,
+            height_range_text,
+            mass_range_text,
+            asset_version_text,
+            gr.State(lambda: TMP_DIR),
+            gr.State(lambda: URDF_CONVERTOR),
+            gr.State(lambda: IMAGE_BUFFER),
+            gr.State(lambda: CHECKERS),
+        ],
+        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],
+    )
 
 
-iface = gr.Interface(fn=my_function, inputs="text", outputs="text")
-iface.launch()
+if __name__ == "__main__":
+    demo.queue()
+    demo.launch(server_name="10.34.8.82", server_port=8084)

asset3d_gen/data/backproject.py

@@ -0,0 +1,503 @@
+import argparse
+import logging
+import math
+import os
+from typing import List, Literal, Tuple, 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 asset3d_gen.data.mesh_operator import MeshFixer
+from asset3d_gen.data.utils import (
+    CameraSetting,
+    get_images_from_grid,
+    init_kal_camera,
+    normalize_vertices_array,
+    post_process_texture,
+    save_mesh_with_mtl,
+)
+from asset3d_gen.models.delight 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.input_image, 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(
+        "--input_image",
+        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.input_image
+    ):
+        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 = torch.Tensor([[1, 0, 0], [0, 0, 1], [0, -1, 0]])
+            z_rot = torch.Tensor([[0, 1, 0], [-1, 0, 0], [0, 0, 1]])
+            vertices = vertices @ x_rot
+            vertices = vertices @ z_rot
+
+        faces = mesh.faces.cpu().numpy().astype(np.int32)
+        vertices = vertices.cpu().numpy().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(
+                model_path="/horizon-bucket/robot_lab/users/xinjie.wang/weights/hunyuan3d-delight-v2-0"  # noqa
+            )
+            delight_images = [delight_model(img) for img in images]
+            images = [np.array(img) for img in delight_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()

asset3d_gen/data/backproject_v2.py

@@ -0,0 +1,613 @@
+import argparse
+import logging
+import math
+import os
+
+import cv2
+import numpy as np
+import nvdiffrast.torch as dr
+import torch
+import torch.nn.functional as F
+import trimesh
+import xatlas
+from PIL import Image
+from asset3d_gen.data.mesh_operator import MeshFixer
+from asset3d_gen.data.utils import (
+    CameraSetting,
+    DiffrastRender,
+    get_images_from_grid,
+    init_kal_camera,
+    normalize_vertices_array,
+    post_process_texture,
+    save_mesh_with_mtl,
+)
+from asset3d_gen.models.delight import DelightingModel
+from asset3d_gen.models.super_resolution import ImageRealESRGAN
+
+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."""
+
+    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,
+    ):
+        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]
+        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=camera_params.device,
+            antialias_mask=True,
+        )
+        self.camera = camera
+        self.renderer = renderer
+        self.view_weights = view_weights
+        self.device = camera_params.device
+        self.render_wh = render_wh
+        self.texture_wh = texture_wh
+
+        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 load_mesh(self, mesh: trimesh.Trimesh) -> None:
+        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
+
+        self.vertices = torch.from_numpy(mesh.vertices).to(self.device).float()
+        self.faces = torch.from_numpy(mesh.faces).to(self.device).to(torch.int)
+        self.uv_map = torch.from_numpy(mesh.visual.uv).to(self.device).float()
+
+    def get_mesh_np_attrs(
+        self,
+        scale: float = None,
+        center: np.ndarray = None,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        vertices = self.vertices.cpu().numpy()
+        faces = self.faces.cpu().numpy()
+        uv_map = self.uv_map.cpu().numpy()
+        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, texture: torch.Tensor, mask: torch.Tensor
+    ) -> np.ndarray:
+        texture_np = texture.cpu().numpy()
+        mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)
+        vertices, faces, uv_map = self.get_mesh_np_attrs()
+
+        texture_np, mask_np = _texture_inpaint_smooth(
+            texture_np, mask_np, vertices, faces, uv_map
+        )
+        texture_np = texture_np.clip(0, 1)
+        texture_np = cv2.inpaint(
+            (texture_np * 255).astype(np.uint8),
+            255 - mask_np,
+            3,
+            cv2.INPAINT_NS,
+        )
+
+        return texture_np
+
+    def __call__(
+        self,
+        colors: list[Image.Image],
+        mesh: trimesh.Trimesh,
+        output_path: str,
+    ) -> trimesh.Trimesh:
+        self.load_mesh(mesh)
+        rendered_depth, masks = self.renderer.render_depth(
+            self.vertices, self.faces
+        )
+        norm_deps = self.renderer.normalize_map_by_mask(rendered_depth, masks)
+        render_uvs, _ = self.renderer.render_uv(
+            self.vertices, self.faces, self.uv_map
+        )
+        view_normals = self.compute_enhanced_viewnormal(
+            self.renderer.mv_mtx, self.vertices, self.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 = self.uv_inpaint(texture, mask)
+        texture_np = post_process_texture(texture_np)
+        vertices, faces, uv_map = self.get_mesh_np_attrs(
+            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",
+        type=str,
+        help="Multiview color image in 6x512x512 file path",
+    )
+    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=2,
+        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."
+    )
+    args = parser.parse_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,
+    )
+    view_weights = [1, 0.1, 0.02, 0.1, 1, 0.02]
+
+    color_grid = Image.open(args.color_path)
+    if args.delight:
+        if delight_model is None:
+            delight_model = DelightingModel(
+                model_path="/horizon-bucket/robot_lab/users/xinjie.wang/weights/hunyuan3d-delight-v2-0"  # noqa
+            )
+        save_dir = os.path.dirname(args.output_path)
+        os.makedirs(save_dir, exist_ok=True)
+        color_grid.save(f"{save_dir}/color_grid.png")
+        color_grid = delight_model(color_grid)
+        color_grid.save(f"{save_dir}/color_grid_delight.png")
+
+    multiviews = get_images_from_grid(color_grid, img_size=512)
+
+    # 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,
+        )
+        # 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=camera_params.resolution_hw,
+        texture_wh=args.texture_wh,
+    )
+
+    textured_mesh = texture_backer(multiviews, mesh, args.output_path)
+
+    return textured_mesh
+
+
+if __name__ == "__main__":
+    entrypoint()

asset3d_gen/data/backup/backproject_v2 copy.py

@@ -0,0 +1,650 @@
+import argparse
+import logging
+import math
+import os
+
+import cv2
+import numpy as np
+import nvdiffrast.torch as dr
+import torch
+import torch.nn.functional as F
+from torchvision.transforms import functional as tF
+
+import trimesh
+import xatlas
+from PIL import Image
+from asset3d_gen.data.mesh_operator import MeshFixer
+from asset3d_gen.data.utils import (
+    CameraSetting,
+    DiffrastRender,
+    get_images_from_grid,
+    init_kal_camera,
+    normalize_vertices_array,
+    post_process_texture,
+    save_mesh_with_mtl,
+)
+from asset3d_gen.models.delight import DelightingModel
+from asset3d_gen.models.super_resolution import ImageRealESRGAN
+
+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
+
+
+def interp_tensers(tensors: list[torch.Tensor], target_wh: tuple[int, int]) -> list[torch.Tensor]:
+    for idx in range(len(tensors)):
+        tensor = tensors[idx].permute(2, 0, 1)
+        tensor = tF.resize(tensor, target_wh[::-1], antialias=True)
+        tensors[idx] = tensor.permute(1, 2, 0)
+
+    return tensors
+
+
+class TextureBacker:
+    """Texture baking pipeline for multi-view projection and fusion."""
+
+    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,
+    ):
+        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=camera_params.device,
+            antialias_mask=True,
+        )
+        self.camera = camera
+        self.view_weights = view_weights
+        self.device = camera_params.device
+        self.render_wh = render_wh
+        self.texture_wh = texture_wh
+
+        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 load_mesh(self, mesh: trimesh.Trimesh) -> None:
+        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
+
+        self.vertices = torch.from_numpy(mesh.vertices).to(self.device).float()
+        self.faces = torch.from_numpy(mesh.faces).to(self.device).to(torch.int)
+        self.uv_map = torch.from_numpy(mesh.visual.uv).to(self.device).float()
+
+    def get_mesh_np_attrs(
+        self,
+        scale: float = None,
+        center: np.ndarray = None,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        vertices = self.vertices.cpu().numpy()
+        faces = self.faces.cpu().numpy()
+        uv_map = self.uv_map.cpu().numpy()
+        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, texture: torch.Tensor, mask: torch.Tensor
+    ) -> np.ndarray:
+        texture_np = texture.cpu().numpy()
+        mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)
+        vertices, faces, uv_map = self.get_mesh_np_attrs()
+
+        texture_np, mask_np = _texture_inpaint_smooth(
+            texture_np, mask_np, vertices, faces, uv_map
+        )
+        texture_np = texture_np.clip(0, 1)
+        texture_np = cv2.inpaint(
+            (texture_np * 255).astype(np.uint8),
+            255 - mask_np,
+            3,
+            cv2.INPAINT_NS,
+        )
+
+        return texture_np
+
+    def __call__(
+        self,
+        colors: list[Image.Image],
+        mesh: trimesh.Trimesh,
+        output_path: str,
+    ) -> trimesh.Trimesh:
+        import time
+        start = time.time()
+        self.load_mesh(mesh)
+        print("load_mesh", time.time() - start)
+        
+        start = time.time()
+        rendered_depth, masks = self.renderer.render_depth(
+            self.vertices, self.faces
+        )
+        norm_deps = self.renderer.normalize_map_by_mask(rendered_depth, masks)
+        render_uvs, _ = self.renderer.render_uv(
+            self.vertices, self.faces, self.uv_map
+        )
+        view_normals = self.compute_enhanced_viewnormal(
+            self.renderer.mv_mtx, self.vertices, self.faces
+        )
+        print("0", time.time() - start)
+
+        textures, weighted_cos_maps = [], []
+        
+        start = time.time()
+        for color, mask, dep, normal, uv, weight in zip(
+            colors,
+            masks,
+            norm_deps,
+            view_normals,
+            render_uvs,
+            self.view_weights,
+        ):
+            mask, dep, normal, uv = interp_tensers([mask, dep, normal, uv], self.render_wh)
+            texture, cos_map = self.back_project(color, mask, dep, normal, uv)
+            textures.append(texture)
+            weighted_cos_maps.append(weight * (cos_map**4))
+        print("1", time.time() - start)
+        start = time.time()
+        texture, mask = self.fast_bake_texture(textures, weighted_cos_maps)
+        print("2", time.time() - start)
+        start = time.time()
+        texture_np = self.uv_inpaint(texture, mask)
+        print("3", time.time() - start)
+        start = time.time()
+        texture_np = post_process_texture(texture_np)
+        vertices, faces, uv_map = self.get_mesh_np_attrs(
+            self.scale, self.center
+        )
+
+        textured_mesh = save_mesh_with_mtl(
+            vertices, faces, uv_map, texture_np, output_path
+        )
+        print("4", time.time() - start)
+
+        return textured_mesh
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Backproject texture")
+    parser.add_argument(
+        "--color_path",
+        type=str,
+        help="Multiview color image in 6x512x512 file path",
+    )
+    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=2,
+        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 mesh rendering",
+    )
+    parser.add_argument(
+        "--target_hw",
+        type=int,
+        nargs=2,
+        default=(2048, 2048),
+        help="Target rendering images resolution",
+    )
+    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."
+    )
+    args = parser.parse_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,
+    )
+    view_weights = [1, 0.1, 0.02, 0.1, 1, 0.02]
+
+    color_grid = Image.open(args.color_path)
+    if args.delight:
+        if delight_model is None:
+            delight_model = DelightingModel(
+                model_path="/horizon-bucket/robot_lab/users/xinjie.wang/weights/hunyuan3d-delight-v2-0"  # noqa
+            )
+        save_dir = os.path.dirname(args.output_path)
+        os.makedirs(save_dir, exist_ok=True)
+        color_grid.save(f"{save_dir}/color_grid.png")
+        color_grid = delight_model(color_grid)
+        color_grid.save(f"{save_dir}/color_grid_delight.png")
+
+    multiviews = get_images_from_grid(color_grid, img_size=512)
+
+    # Use RealESRGAN_x4plus for x4 (512->2048) image super resolution.
+    if imagesr_model is None:
+        imagesr_model = ImageRealESRGAN(outscale=4)
+    multiviews = [imagesr_model(img.convert("RGB")) for img in multiviews]
+    multiviews = [img.resize(args.target_hw[::-1]) 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,
+        )
+        # Restore scale.
+        mesh.vertices = mesh.vertices / scale
+        mesh.vertices = mesh.vertices + center
+
+    # Baking texture to mesh.
+    import time
+    start = time.time()
+    texture_backer = TextureBacker(
+        camera_params=camera_params,
+        view_weights=view_weights,
+        render_wh=args.target_hw,
+        texture_wh=args.texture_wh,
+    )
+    print(time.time()-start)
+    start = time.time()
+    textured_mesh = texture_backer(multiviews, mesh, args.output_path)
+    print(f"Texture backproject time: {time.time() - start:.2f}s")
+
+    return textured_mesh
+
+
+if __name__ == "__main__":
+    entrypoint()

asset3d_gen/data/backup/backproject_v2.py

@@ -0,0 +1,700 @@
+import logging
+import math
+from typing import Union
+
+import custom_rasterizer as cr
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+import trimesh
+import xatlas
+from PIL import Image
+from asset3d_gen.data.utils import (
+    get_images_from_file,
+    normalize_vertices_array,
+    post_process_texture,
+    save_mesh_with_mtl,
+)
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+__all__ = ["TextureBacker", "Image_Super_Net", "Image_GANNet"]
+
+
+import math
+import numpy as np
+
+
+def get_perspective_projection(
+    fov: float, aspect_wh: float, near: float = 0.01, far: float = 100
+) -> np.ndarray:
+    """Compute the perspective projection matrix for 3D rendering."""
+    fov_rad = math.radians(fov)
+    tan_half_fov = math.tan(fov_rad / 2.0)
+
+    return np.array(
+        [
+            [1.0 / (tan_half_fov * aspect_wh), 0.0, 0.0, 0.0],
+            [0.0, 1.0 / tan_half_fov, 0.0, 0.0],
+            [
+                0.0,
+                0.0,
+                -(far + near) / (far - near),
+                -(2.0 * far * near) / (far - near),
+            ],
+            [0.0, 0.0, -1.0, 0.0],
+        ],
+        dtype=np.float32,
+    )
+
+
+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 compute_w2c_matrix(
+    elev_deg: float, azim_deg: float, cam_dist: float
+) -> np.ndarray:
+    """Compute w2c 4x4 transformation matrix from spherical coordinates."""
+
+    elev_rad = math.radians(-elev_deg)
+    azim_rad = math.radians(azim_deg)
+
+    sin_elev = math.sin(elev_rad)
+    cos_elev = math.cos(elev_rad)
+    sin_azim = math.sin(azim_rad)
+    cos_azim = math.cos(azim_rad)
+
+    cam_pos = np.array(
+        [
+            cam_dist * cos_elev * cos_azim,
+            cam_dist * cos_elev * sin_azim,
+            cam_dist * sin_elev,
+        ]
+    )
+
+    look_dir = -cam_pos / np.linalg.norm(cam_pos)
+    right_dir = np.cross(look_dir, [0, 0, 1])
+    right_dir /= np.linalg.norm(right_dir)
+    up_dir = np.cross(right_dir, look_dir)
+
+    c2w = np.eye(4)
+    c2w[:3, 0] = right_dir
+    c2w[:3, 1] = up_dir
+    c2w[:3, 2] = -look_dir
+    c2w[:3, 3] = cam_pos
+
+    try:
+        w2c = np.linalg.inv(c2w)
+    except np.linalg.LinAlgError as e:
+        raise ArithmeticError("Failed to invert camera-to-world matrix") from e
+
+    return w2c.astype(np.float32)
+
+
+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."""
+
+    def __init__(
+        self,
+        camera_elevs: list[float],
+        camera_azims: list[float],
+        camera_distance: int,
+        camera_fov: float,
+        view_weights: list[float] = None,
+        render_wh: tuple[int, int] = (2048, 2048),
+        texture_wh: tuple[int, int] = (2048, 2048),
+        use_antialias: bool = True,
+        bake_angle_thres: int = 75,
+        device="cuda",
+    ):
+        self.camera_elevs = camera_elevs
+        self.camera_azims = camera_azims
+        self.view_weights = (
+            view_weights
+            if view_weights is not None
+            else [1] * len(camera_elevs)
+        )
+        self.device = device
+        self.render_wh = render_wh
+        self.texture_wh = texture_wh
+
+        self.camera_distance = camera_distance
+        self.use_antialias = use_antialias
+
+        self.bake_angle_thres = bake_angle_thres
+        self.bake_unreliable_kernel_size = int(
+            (2 / 512) * max(self.render_wh[0], self.render_wh[1])
+        )
+
+        self.camera_proj_mat = get_perspective_projection(
+            camera_fov,
+            self.render_wh[1] / self.render_wh[0],
+        )
+        self.cnt = 0
+
+    def rasterize_mesh(
+        self,
+        vertex: torch.Tensor,
+        face: torch.Tensor,
+        resolution: tuple[int, int],
+    ) -> torch.Tensor:
+        vertex = vertex[None] if vertex.ndim == 2 else vertex
+        indices, weights = cr.rasterize(vertex, face, resolution)
+
+        return torch.cat(
+            [weights, indices.unsqueeze(-1).to(weights.dtype)], dim=-1
+        ).unsqueeze(0)
+
+    def raster_interpolate(
+        self, uv: torch.Tensor, rast_out: torch.Tensor, faces: torch.Tensor
+    ) -> torch.Tensor:
+        barycentric = rast_out[0, ..., :-1]
+        findices = rast_out[0, ..., -1]
+        if uv.dim() == 2:
+            uv = uv.unsqueeze(0)
+
+        return cr.interpolate(uv, findices, barycentric, faces)[0]
+
+    def load_mesh(self, mesh_path: str) -> None:
+        mesh = trimesh.load(mesh_path)
+        if isinstance(mesh, trimesh.Scene):
+            mesh = mesh.dump(concatenate=True)
+
+        mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
+        self.scale, self.center = scale, center
+
+        vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
+        mesh.vertices = mesh.vertices[vmapping]
+        mesh.faces = indices
+        mesh.visual.uv = uvs
+
+        self.vertices = torch.from_numpy(mesh.vertices).to(self.device).float()
+        self.faces = torch.from_numpy(mesh.faces).to(self.device).to(torch.int)
+        self.uv_map = torch.from_numpy(mesh.visual.uv).to(self.device).float()
+
+        # Transformation of coordinate system
+        self.vertices[:, [0, 1]] = -self.vertices[:, [0, 1]]
+        self.vertices[:, [1, 2]] = self.vertices[:, [2, 1]]
+        self.uv_map[:, 1] = 1 - self.uv_map[:, 1]
+
+    def get_mesh_attrs(
+        self,
+        scale: float = None,
+        center: np.ndarray = None,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        vertices = self.vertices.cpu().numpy()
+        faces = self.faces.cpu().numpy()
+        uv_map = self.uv_map.cpu().numpy()
+
+        # Inverse transformation of coordinate system
+        vertices[:, [1, 2]] = vertices[:, [2, 1]]
+        vertices[:, [0, 1]] = -vertices[:, [0, 1]]
+        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)
+        combined_edges = depth_edges
+        sketch_image = (
+            torch.from_numpy(combined_edges).to(depth_image.device).float()
+            / 255
+        )
+        sketch_image = sketch_image.unsqueeze(-1)
+
+        return sketch_image
+
+    def back_project(
+        self, image: Image.Image, elev: float, azim: float
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if isinstance(image, Image.Image):
+            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.0
+
+        view_mat = compute_w2c_matrix(elev, azim, self.camera_distance)
+        import pdb
+
+        pdb.set_trace()
+        pos_cam = transform_vertices(view_mat, self.vertices, keepdim=True)
+        pos_clip = transform_vertices(self.camera_proj_mat, pos_cam)
+        pos_cam = pos_cam[:, :3] / pos_cam[:, 3:]
+
+        v0, v1, v2 = (pos_cam[self.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), self.faces.cpu(), face_norm.cpu()
+                )
+            )
+            .to(self.device)
+            .contiguous()
+        )
+
+        rast_out = self.rasterize_mesh(pos_clip, self.faces, image.shape[:2])
+        vis_mask = torch.clamp(rast_out[..., -1:], 0, 1)[0]
+
+        interp_data = {
+            "normal": self.raster_interpolate(
+                vertex_norm[None], rast_out, self.faces
+            ),
+            "uv": self.raster_interpolate(
+                self.uv_map[None], rast_out, self.faces
+            ),
+            "depth": self.raster_interpolate(
+                pos_cam[:, 2].reshape(1, -1, 1), rast_out, self.faces
+            ),
+        }
+
+        valid_depth = interp_data["depth"][vis_mask > 0]
+        depth_norm = (interp_data["depth"] - valid_depth.min()) / (
+            valid_depth.max() - valid_depth.min()
+        )
+        # depth_norm[vis_mask <= 0] = 0
+        sketch_image = self._render_depth_edges(depth_norm * vis_mask)
+
+        # ddd = depth_norm * vis_mask
+        # cv2.imwrite(f"v2_depth_d{self.cnt}.png", (ddd.cpu().numpy() * 255).astype(np.uint8))
+
+        cv2.imwrite(
+            f"v2_vis_mask{self.cnt}.png",
+            (vis_mask.cpu().numpy() * 255).astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v2_normal{self.cnt}.png",
+            (interp_data["normal"].cpu().numpy() * 255).astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v2_depth{self.cnt}.png",
+            (depth_norm.cpu().numpy() * 255).astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v2_uv{self.cnt}.png",
+            (interp_data["uv"][..., 0].cpu().numpy() * 255).astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v2_sketch{self.cnt}.png",
+            (sketch_image.cpu().numpy() * 255).astype(np.uint8),
+        )
+
+        self.cnt += 1
+
+        cos = F.cosine_similarity(
+            torch.tensor([[0, 0, -1]], device=self.device),
+            interp_data["normal"].view(-1, 3),
+        ).view_as(interp_data["normal"][..., :1])
+        cos[cos < np.cos(np.radians(self.bake_angle_thres))] = 0
+
+        cv2.imwrite(
+            f"v2_cos{self.cnt}.png", (cos.cpu().numpy() * 255).astype(np.uint8)
+        )
+
+        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
+
+        vis_mask = cv2.imread(
+            f"v3_db_mask{self.cnt}.png", cv2.IMREAD_GRAYSCALE
+        )
+        vis_mask = (
+            torch.from_numpy(vis_mask[..., None]).to(self.device).float() / 255
+        )
+        # cos2 = cv2.imread(f"v3_db_cos{self.cnt}.png", cv2.IMREAD_GRAYSCALE)
+        # cos2 = torch.from_numpy(cos2[..., None]).to(self.device).float() / 255
+        # cos = cos2
+
+        valid_pixels = (vis_mask != 0).view(-1)
+        # import pdb; pdb.set_trace()
+
+        cv2.imwrite(
+            f"v2_db_sketch{self.cnt}.png",
+            (sketch_image.cpu().numpy() * 255).astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v2_db_uv{self.cnt}.png",
+            (interp_data["uv"][..., 0].cpu().numpy() * 255).astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v2_db_uv2{self.cnt}.png",
+            (interp_data["uv"][..., 1].cpu().numpy() * 255).astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v2_db_color{self.cnt}.png",
+            (image.cpu().numpy() * 255).astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v2_db_cos{self.cnt}.png",
+            (cos.cpu().numpy() * 255).astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v2_db_mask{self.cnt}.png",
+            (vis_mask.cpu().numpy() * 255).astype(np.uint8),
+        )
+        # import pdb; pdb.set_trace()
+        return (
+            self._scatter_texture(interp_data["uv"], image, valid_pixels),
+            self._scatter_texture(interp_data["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, texture: torch.Tensor, mask: torch.Tensor
+    ) -> np.ndarray:
+        texture_np = texture.cpu().numpy()
+        mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)
+        vertices, faces, uv_map = self.get_mesh_attrs()
+        # import pdb; pdb.set_trace()
+        texture_np, mask_np = _texture_inpaint_smooth(
+            texture_np, mask_np, vertices, faces, uv_map
+        )
+        texture_np = texture_np.clip(0, 1)
+        texture_np = cv2.inpaint(
+            (texture_np * 255).astype(np.uint8),
+            255 - mask_np,
+            3,
+            cv2.INPAINT_NS,
+        )
+
+        return texture_np
+
+    def __call__(
+        self, colors: list[Image.Image], input_mesh: str, output_path: str
+    ) -> trimesh.Trimesh:
+        self.load_mesh(input_mesh)
+
+        textures, weighted_cos_maps = [], []
+        for color, cam_elev, cam_azim, weight in zip(
+            colors, self.camera_elevs, self.camera_azims, self.view_weights
+        ):
+            texture, cos_map = self.back_project(color, cam_elev, cam_azim)
+            cv2.imwrite(
+                f"v2_texture{self.cnt}.png",
+                (texture.cpu().numpy() * 255).astype(np.uint8),
+            )
+            cv2.imwrite(
+                f"v2_texture_cos{self.cnt}.png",
+                (cos_map.cpu().numpy() * 255).astype(np.uint8),
+            )
+            # import pdb; pdb.set_trace()
+            textures.append(texture)
+            weighted_cos_maps.append(weight * (cos_map**4))
+
+        texture, mask = self.fast_bake_texture(textures, weighted_cos_maps)
+        texture_np = self.uv_inpaint(texture, mask)
+        texture_np = post_process_texture(texture_np)
+        vertices, faces, uvs = self.get_mesh_attrs(self.scale, self.center)
+        # import pdb; pdb.set_trace()
+        cv2.imwrite("v2_texture_np.png", texture_np)
+
+        textured_mesh = save_mesh_with_mtl(
+            vertices, faces, uvs, texture_np, output_path
+        )
+
+        return textured_mesh
+
+
+class Image_Super_Net:
+    def __init__(self, device="cuda"):
+        from diffusers import StableDiffusionUpscalePipeline
+
+        self.up_pipeline_x4 = StableDiffusionUpscalePipeline.from_pretrained(
+            "stabilityai/stable-diffusion-x4-upscaler",
+            torch_dtype=torch.float16,
+        ).to(device)
+        self.up_pipeline_x4.set_progress_bar_config(disable=True)
+
+    def __call__(self, image, prompt=""):
+        with torch.no_grad():
+            upscaled_image = self.up_pipeline_x4(
+                prompt=[prompt],
+                image=image,
+                num_inference_steps=10,
+            ).images[0]
+
+        return upscaled_image
+
+
+class Image_GANNet:
+    def __init__(self, outscale: int):
+        from basicsr.archs.rrdbnet_arch import RRDBNet
+        from realesrgan import RealESRGANer
+
+        self.outscale = outscale
+        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="/horizon-bucket/robot_lab/users/xinjie.wang/weights/super_resolution/RealESRGAN_x4plus.pth",  # noqa
+            model=model,
+            pre_pad=0,
+            half=True,
+        )
+
+    def __call__(self, image: Union[Image.Image, np.ndarray]) -> Image.Image:
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        output, _ = self.upsampler.enhance(image, outscale=self.outscale)
+
+        return Image.fromarray(output)
+
+
+if __name__ == "__main__":
+    device = "cuda"
+    color_path = "outputs/texture_mesh_gen/multi_view/color_sample0.png"
+    mesh_path = "outputs/texture_mesh_gen/texture_mesh/kettle_color.glb"
+    output_path = "robot_test_v2/robot.obj"
+    target_image_size = (2048, 2048)
+
+    super_model = Image_GANNet(outscale=4)
+    multiviews = get_images_from_file(color_path, img_size=512)
+
+    texture_backer = TextureBacker(
+        camera_elevs=[20, 20, 20, -10, -10, -10],
+        camera_azims=[-180, -60, 60, -120, 0, 120],
+        view_weights=[1, 0.2, 0.2, 0.2, 1, 0.2],
+        camera_distance=5,
+        camera_fov=30,
+        render_wh=(2048, 2048),
+        texture_wh=(2048, 2048),
+    )
+
+    multiviews = [super_model(img) for img in multiviews]
+    multiviews = [img.convert("RGB") for img in multiviews]
+    textured_mesh = texture_backer(multiviews, mesh_path, output_path)

asset3d_gen/data/backup/backproject_v3.py

@@ -0,0 +1,866 @@
+import logging
+import math
+from typing import Union
+
+import custom_rasterizer as cr
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+import trimesh
+import xatlas
+from PIL import Image
+from asset3d_gen.data.utils import (
+    get_images_from_file,
+    normalize_vertices_array,
+    post_process_texture,
+    save_mesh_with_mtl,
+)
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+__all__ = ["TextureBacker", "Image_Super_Net", "Image_GANNet"]
+
+
+import math
+import numpy as np
+
+
+def get_perspective_projection(
+    fov: float, aspect_wh: float, near: float = 0.01, far: float = 100
+) -> np.ndarray:
+    """Compute the perspective projection matrix for 3D rendering."""
+    fov_rad = math.radians(fov)
+    tan_half_fov = math.tan(fov_rad / 2.0)
+
+    return np.array(
+        [
+            [1.0 / (tan_half_fov * aspect_wh), 0.0, 0.0, 0.0],
+            [0.0, 1.0 / tan_half_fov, 0.0, 0.0],
+            [
+                0.0,
+                0.0,
+                -(far + near) / (far - near),
+                -(2.0 * far * near) / (far - near),
+            ],
+            [0.0, 0.0, -1.0, 0.0],
+        ],
+        dtype=np.float32,
+    )
+
+
+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 compute_w2c_matrix(
+    elev_deg: float, azim_deg: float, cam_dist: float
+) -> np.ndarray:
+    """Compute w2c 4x4 transformation matrix from spherical coordinates."""
+
+    elev_rad = math.radians(-elev_deg)
+    azim_rad = math.radians(azim_deg)
+
+    sin_elev = math.sin(elev_rad)
+    cos_elev = math.cos(elev_rad)
+    sin_azim = math.sin(azim_rad)
+    cos_azim = math.cos(azim_rad)
+
+    cam_pos = np.array(
+        [
+            cam_dist * cos_elev * cos_azim,
+            cam_dist * cos_elev * sin_azim,
+            cam_dist * sin_elev,
+        ]
+    )
+
+    look_dir = -cam_pos / np.linalg.norm(cam_pos)
+    right_dir = np.cross(look_dir, [0, 0, 1])
+    right_dir /= np.linalg.norm(right_dir)
+    up_dir = np.cross(right_dir, look_dir)
+
+    c2w = np.eye(4)
+    c2w[:3, 0] = right_dir
+    c2w[:3, 1] = up_dir
+    c2w[:3, 2] = -look_dir
+    c2w[:3, 3] = cam_pos
+
+    try:
+        w2c = np.linalg.inv(c2w)
+    except np.linalg.LinAlgError as e:
+        raise ArithmeticError("Failed to invert camera-to-world matrix") from e
+
+    return w2c.astype(np.float32)
+
+
+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."""
+
+    def __init__(
+        self,
+        camera_elevs: list[float],
+        camera_azims: list[float],
+        camera_distance: int,
+        camera_fov: float,
+        view_weights: list[float] = None,
+        render_wh: tuple[int, int] = (2048, 2048),
+        texture_wh: tuple[int, int] = (2048, 2048),
+        use_antialias: bool = True,
+        bake_angle_thresh: int = 75,
+        device="cuda",
+    ):
+        self.camera_elevs = camera_elevs
+        self.camera_azims = camera_azims
+        self.view_weights = (
+            view_weights
+            if view_weights is not None
+            else [1] * len(camera_elevs)
+        )
+        self.device = device
+        self.render_wh = render_wh
+        self.texture_wh = texture_wh
+
+        self.camera_distance = camera_distance
+        self.use_antialias = use_antialias
+
+        self.bake_angle_thresh = bake_angle_thresh
+        self.bake_unreliable_kernel_size = int(
+            (2 / 512) * max(self.render_wh[0], self.render_wh[1])
+        )
+
+        self.camera_proj_mat = get_perspective_projection(
+            camera_fov,
+            self.render_wh[1] / self.render_wh[0],
+        )
+        self.cnt = 0
+
+    def rasterize_mesh(
+        self,
+        vertex: torch.Tensor,
+        face: torch.Tensor,
+        resolution: tuple[int, int],
+    ) -> torch.Tensor:
+        vertex = vertex[None] if vertex.ndim == 2 else vertex
+        indices, weights = cr.rasterize(vertex, face, resolution)
+
+        return torch.cat(
+            [weights, indices.unsqueeze(-1).to(weights.dtype)], dim=-1
+        ).unsqueeze(0)
+
+    def raster_interpolate(
+        self, uv: torch.Tensor, rast_out: torch.Tensor, faces: torch.Tensor
+    ) -> torch.Tensor:
+        barycentric = rast_out[0, ..., :-1]
+        findices = rast_out[0, ..., -1]
+        if uv.dim() == 2:
+            uv = uv.unsqueeze(0)
+
+        return cr.interpolate(uv, findices, barycentric, faces)[0]
+
+    def load_mesh(self, mesh_path: str) -> None:
+        mesh = trimesh.load(mesh_path)
+        if isinstance(mesh, trimesh.Scene):
+            mesh = mesh.dump(concatenate=True)
+
+        mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
+        self.scale, self.center = scale, center
+
+        vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
+        mesh.vertices = mesh.vertices[vmapping]
+        mesh.faces = indices
+        mesh.visual.uv = uvs
+
+        self.vertices = torch.from_numpy(mesh.vertices).to(self.device).float()
+        self.faces = torch.from_numpy(mesh.faces).to(self.device).to(torch.int)
+        self.uv_map = torch.from_numpy(mesh.visual.uv).to(self.device).float()
+
+        # Transformation of coordinate system
+        self.vertices[:, [0, 1]] = -self.vertices[:, [0, 1]]
+        self.vertices[:, [1, 2]] = self.vertices[:, [2, 1]]
+        self.uv_map[:, 1] = 1 - self.uv_map[:, 1]
+
+    def get_mesh_attrs(
+        self,
+        scale: float = None,
+        center: np.ndarray = None,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        vertices = self.vertices.cpu().numpy()
+        faces = self.faces.cpu().numpy()
+        uv_map = self.uv_map.cpu().numpy()
+
+        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 back_project(
+        self, image: Image.Image, elev: float, azim: float
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if isinstance(image, Image.Image):
+            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.0
+
+        view_mat = compute_w2c_matrix(elev, azim, self.camera_distance)
+        pos_cam = transform_vertices(view_mat, self.vertices, keepdim=True)
+        pos_clip = transform_vertices(self.camera_proj_mat, pos_cam)
+        pos_cam = pos_cam[:, :3] / pos_cam[:, 3:]
+
+        v0, v1, v2 = (pos_cam[self.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), self.faces.cpu(), face_norm.cpu()
+                )
+            )
+            .to(self.device)
+            .contiguous()
+        )
+
+        rast_out = self.rasterize_mesh(pos_clip, self.faces, image.shape[:2])
+        vis_mask = torch.clamp(rast_out[..., -1:], 0, 1)[0]
+
+        interp_data = {
+            "normal": self.raster_interpolate(
+                vertex_norm[None], rast_out, self.faces
+            ),
+            "uv": self.raster_interpolate(
+                self.uv_map[None], rast_out, self.faces
+            ),
+            "depth": self.raster_interpolate(
+                pos_cam[:, 2].reshape(1, -1, 1), rast_out, self.faces
+            ),
+        }
+
+        valid_depth = interp_data["depth"][vis_mask > 0]
+        depth_norm = (interp_data["depth"] - valid_depth.min()) / (
+            valid_depth.max() - valid_depth.min()
+        )
+        depth_norm[vis_mask <= 0] = 0
+        sketch_image = self._render_depth_edges(depth_norm * vis_mask)
+
+        # cv2.imwrite("vis_mask.png", (vis_mask.cpu().numpy() * 255).astype(np.uint8))
+        # cv2.imwrite("normal.png", (interp_data['normal'].cpu().numpy() * 255).astype(np.uint8))
+        # cv2.imwrite("depth.png", (depth_norm.cpu().numpy() * 255).astype(np.uint8))
+        # cv2.imwrite("uv.png", (interp_data['uv'][..., 0].cpu().numpy() * 255).astype(np.uint8))
+        # import pdb; pdb.set_trace()
+
+        cos = F.cosine_similarity(
+            torch.tensor([[0, 0, -1]], device=self.device),
+            interp_data["normal"].view(-1, 3),
+        ).view_as(interp_data["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(interp_data["uv"], image, valid_pixels),
+            self._scatter_texture(interp_data["uv"], cos, valid_pixels),
+        )
+
+    def back_project2(
+        self, image, vis_mask, depth, normal, uv
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if isinstance(image, Image.Image):
+            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.0
+
+        depth_inv = (1.0 - depth) * vis_mask
+        sketch_image = self._render_depth_edges(depth_inv)
+
+        cv2.imwrite(
+            f"v3_depth_inv{self.cnt}.png",
+            (depth_inv.cpu().numpy() * 255).astype(np.uint8),
+        )
+
+        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
+        # import pdb; pdb.set_trace()
+        # cv2.imwrite(f"v3_cos{self.cnt}.png", (cos.cpu().numpy() * 255).astype(np.uint8))
+        # cv2.imwrite(f"v3_sketch{self.cnt}.png", (sketch_image.cpu().numpy() * 255).astype(np.uint8))
+
+        # cos2 = cv2.imread(f"v2_cos{self.cnt+1}.png", cv2.IMREAD_GRAYSCALE)
+        # cos2 = torch.from_numpy(cos2[..., None]).to(self.device).float() / 255
+        # cos = cos2
+
+        self.cnt += 1
+
+        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)
+        # import pdb; pdb.set_trace()
+        cv2.imwrite(
+            f"v3_db_sketch{self.cnt}.png",
+            (sketch_image.cpu().numpy() * 255).astype(np.uint8),
+        )
+
+        cos[vis_mask == 0] = 0
+        # import pdb; pdb.set_trace()
+        # vis_mask = cv2.imread(f"v2_db_mask{self.cnt}.png", cv2.IMREAD_GRAYSCALE)
+        # vis_mask = torch.from_numpy(vis_mask[..., None]).to(self.device).float() / 255
+        # cos2 = cv2.imread(f"v2_db_cos{self.cnt}.png", cv2.IMREAD_GRAYSCALE)
+        # cos2 = torch.from_numpy(cos2[..., None]).to(self.device).float() / 255
+        # cos = cos2
+
+        valid_pixels = (vis_mask != 0).view(-1)
+        # import pdb; pdb.set_trace()
+        cv2.imwrite(
+            f"v3_db_uv{self.cnt}.png",
+            (uv[..., 0].cpu().numpy() * 255).astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v3_db_uv2{self.cnt}.png",
+            (uv[..., 1].cpu().numpy() * 255).astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v3_db_color{self.cnt}.png",
+            (image.cpu().numpy() * 255).astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v3_db_cos{self.cnt}.png",
+            (cos.cpu().numpy() * 255).astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v3_db_mask{self.cnt}.png",
+            (vis_mask.cpu().numpy() * 255).astype(np.uint8),
+        )
+
+        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, texture: torch.Tensor, mask: torch.Tensor
+    ) -> np.ndarray:
+        texture_np = texture.cpu().numpy()
+        mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)
+        vertices, faces, uv_map = self.get_mesh_attrs()
+        # import pdb; pdb.set_trace()
+        texture_np, mask_np = _texture_inpaint_smooth(
+            texture_np, mask_np, vertices, faces, uv_map
+        )
+        texture_np = texture_np.clip(0, 1)
+        texture_np = cv2.inpaint(
+            (texture_np * 255).astype(np.uint8),
+            255 - mask_np,
+            3,
+            cv2.INPAINT_NS,
+        )
+
+        return texture_np
+
+    def __call__(
+        self, colors: list[Image.Image], input_mesh: str, output_path: str
+    ) -> trimesh.Trimesh:
+        self.load_mesh(input_mesh)
+
+        textures, weighted_cos_maps = [], []
+        for color, cam_elev, cam_azim, weight in zip(
+            colors, self.camera_elevs, self.camera_azims, self.view_weights
+        ):
+            texture, cos_map = self.back_project(color, cam_elev, cam_azim)
+            textures.append(texture)
+            weighted_cos_maps.append(weight * (cos_map**4))
+
+        texture, mask = self.fast_bake_texture(textures, weighted_cos_maps)
+        texture_np = self.uv_inpaint(texture, mask)
+        texture_np = post_process_texture(texture_np)
+        vertices, faces, uv_map = self.get_mesh_attrs(self.scale, self.center)
+        # import pdb; pdb.set_trace()
+        textured_mesh = save_mesh_with_mtl(
+            vertices, faces, uv_map, texture_np, output_path
+        )
+
+        return textured_mesh
+
+    def forward(
+        self,
+        colors: list[Image.Image],
+        masks,
+        depths,
+        normals,
+        uvs,
+    ) -> trimesh.Trimesh:
+        textures, weighted_cos_maps = [], []
+        for color, mask, depth, normal, uv, weight in zip(
+            colors, masks, depths, normals, uvs, self.view_weights
+        ):
+            texture, cos_map = self.back_project2(
+                color, mask, depth, normal, uv
+            )
+            cv2.imwrite(
+                f"v3_texture{self.cnt}.png",
+                (texture.cpu().numpy() * 255).astype(np.uint8),
+            )
+            cv2.imwrite(
+                f"v3_texture_cos{self.cnt}.png",
+                (cos_map.cpu().numpy() * 255).astype(np.uint8),
+            )
+
+            textures.append(texture)
+            weighted_cos_maps.append(weight * (cos_map**4))
+
+        texture, mask = self.fast_bake_texture(textures, weighted_cos_maps)
+        texture_np = self.uv_inpaint(texture, mask)
+        texture_np = post_process_texture(texture_np)
+        vertices, faces, uv_map = self.get_mesh_attrs(self.scale, self.center)
+        # import pdb; pdb.set_trace()
+        cv2.imwrite("v3_texture_np.png", texture_np)
+        textured_mesh = save_mesh_with_mtl(
+            vertices, faces, uv_map, texture_np, output_path
+        )
+
+        return textured_mesh
+
+
+class Image_Super_Net:
+    def __init__(self, device="cuda"):
+        from diffusers import StableDiffusionUpscalePipeline
+
+        self.up_pipeline_x4 = StableDiffusionUpscalePipeline.from_pretrained(
+            "stabilityai/stable-diffusion-x4-upscaler",
+            torch_dtype=torch.float16,
+        ).to(device)
+        self.up_pipeline_x4.set_progress_bar_config(disable=True)
+
+    def __call__(self, image, prompt=""):
+        with torch.no_grad():
+            upscaled_image = self.up_pipeline_x4(
+                prompt=[prompt],
+                image=image,
+                num_inference_steps=10,
+            ).images[0]
+
+        return upscaled_image
+
+
+class Image_GANNet:
+    def __init__(self, outscale: int):
+        from basicsr.archs.rrdbnet_arch import RRDBNet
+        from realesrgan import RealESRGANer
+
+        self.outscale = outscale
+        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="/horizon-bucket/robot_lab/users/xinjie.wang/weights/super_resolution/RealESRGAN_x4plus.pth",  # noqa
+            model=model,
+            pre_pad=0,
+            half=True,
+        )
+
+    def __call__(self, image: Union[Image.Image, np.ndarray]) -> Image.Image:
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        output, _ = self.upsampler.enhance(image, outscale=self.outscale)
+
+        return Image.fromarray(output)
+
+
+if __name__ == "__main__":
+    device = "cuda"
+    color_path = "outputs/texture_mesh_gen/multi_view/color_sample0.png"
+    mesh_path = "outputs/texture_mesh_gen/texture_mesh/kettle_color.glb"
+    output_path = "robot_test_v6/robot.obj"
+    target_image_size = (2048, 2048)
+
+    super_model = Image_GANNet(outscale=4)
+    multiviews = get_images_from_file(color_path, img_size=512)
+    multiviews = [super_model(img) for img in multiviews]
+    multiviews = [img.convert("RGB") for img in multiviews]
+
+    from asset3d_gen.data.utils import (
+        CameraSetting,
+        init_kal_camera,
+        DiffrastRender,
+    )
+    import nvdiffrast.torch as dr
+
+    camera_params = CameraSetting(
+        num_images=6,
+        elevation=[20.0, -10.0],
+        distance=5,
+        resolution_hw=(2048, 2048),
+        fov=math.radians(30),
+        device="cuda",
+    )
+    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]
+    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,
+    )
+
+    mesh = trimesh.load(mesh_path)
+    if isinstance(mesh, trimesh.Scene):
+        mesh = mesh.dump(concatenate=True)
+
+    mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
+    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
+
+    vertices = torch.from_numpy(mesh.vertices).to(camera_params.device).float()
+    faces = (
+        torch.from_numpy(mesh.faces).to(camera_params.device).to(torch.int64)
+    )
+    uvs = torch.from_numpy(mesh.visual.uv).to(camera_params.device).float()
+
+    rendered_view_normals = []
+    rast, vertices_clip = renderer.compute_dr_raster(vertices, faces)
+    for idx in range(len(mv)):
+        pos_cam = transform_vertices(mv[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(camera_params.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)
+
+    rendered_depth, masks = renderer.render_depth(vertices, faces)
+    norm_depths = []
+    for idx in range(len(rendered_depth)):
+        norm_depth = renderer.normalize_map_by_mask(
+            rendered_depth[idx : idx + 1], masks[idx : idx + 1]
+        )
+        norm_depths.append(norm_depth)
+    norm_depths = torch.cat(norm_depths, dim=0)
+    render_uvs, _ = renderer.render_uv(vertices, faces, uvs)
+
+    for index in range(6):
+        cv2.imwrite(
+            f"v3_mask{index}.png",
+            (masks[index] * 255).cpu().numpy().astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v3_normalv2{index}.png",
+            (rendered_view_normals[index] * 255)
+            .cpu()
+            .numpy()
+            .astype(np.uint8)[..., ::-1],
+        )
+        cv2.imwrite(
+            f"v3_depth{index}.png",
+            (norm_depths[index] * 255).cpu().numpy().astype(np.uint8),
+        )
+        cv2.imwrite(
+            f"v3_uv{index}.png",
+            (render_uvs[index, ..., 0] * 255).cpu().numpy().astype(np.uint8),
+        )
+        multiviews[index].save(f"v3_color{index}.png")
+
+    texture_backer = TextureBacker(
+        camera_elevs=[20, 20, 20, -10, -10, -10],
+        camera_azims=[-180, -60, 60, -120, 0, 120],
+        view_weights=[1, 0.2, 0.2, 0.2, 1, 0.2],
+        camera_distance=5,
+        camera_fov=30,
+        render_wh=(2048, 2048),
+        texture_wh=(2048, 2048),
+    )
+    texture_backer.vertices = vertices
+    texture_backer.faces = faces
+    uvs[:, 1] = 1.0 - uvs[:, 1]
+    texture_backer.uv_map = uvs
+    texture_backer.center = center
+    texture_backer.scale = scale
+
+    textured_mesh = texture_backer.forward(
+        multiviews, masks, norm_depths, rendered_view_normals, render_uvs
+    )
+
+    # multiviews = [super_model(img) for img in multiviews]
+    # multiviews = [img.convert("RGB") for img in multiviews]
+    # textured_mesh = texture_backer(multiviews, mesh_path, output_path)

asset3d_gen/data/backup/backprojectv2.py

@@ -0,0 +1,835 @@
+from PIL import Image
+import torch
+import torch.nn.functional as F
+import numpy as np
+import math
+import trimesh
+import cv2
+import xatlas
+from typing import Union
+
+
+def get_perspective_projection_matrix(fovy, aspect_wh, near, far):
+    fovy_rad = math.radians(fovy)
+    return np.array(
+        [
+            [1.0 / (math.tan(fovy_rad / 2.0) * aspect_wh), 0, 0, 0],
+            [0, 1.0 / math.tan(fovy_rad / 2.0), 0, 0],
+            [
+                0,
+                0,
+                -(far + near) / (far - near),
+                -2.0 * far * near / (far - near),
+            ],
+            [0, 0, -1, 0],
+        ]
+    ).astype(np.float32)
+
+
+def load_mesh(mesh):
+    vtx_pos = mesh.vertices if hasattr(mesh, "vertices") else None
+    pos_idx = mesh.faces if hasattr(mesh, "faces") else None
+
+    vtx_uv = mesh.visual.uv if hasattr(mesh.visual, "uv") else None
+    uv_idx = mesh.faces if hasattr(mesh, "faces") else None
+
+    texture_data = None
+
+    return vtx_pos, pos_idx, vtx_uv, uv_idx, texture_data
+
+
+def save_mesh(mesh, texture_data):
+    material = trimesh.visual.texture.SimpleMaterial(
+        image=texture_data, diffuse=(255, 255, 255)
+    )
+    texture_visuals = trimesh.visual.TextureVisuals(
+        uv=mesh.visual.uv, image=texture_data, material=material
+    )
+    mesh.visual = texture_visuals
+    return mesh
+
+
+def transform_pos(mtx, pos, keepdim=False):
+    t_mtx = (
+        torch.from_numpy(mtx).to(pos.device)
+        if isinstance(mtx, np.ndarray)
+        else mtx
+    )
+    if pos.shape[-1] == 3:
+        posw = torch.cat(
+            [pos, torch.ones([pos.shape[0], 1]).to(pos.device)], axis=1
+        )
+    else:
+        posw = pos
+
+    if keepdim:
+        return torch.matmul(posw, t_mtx.t())[...]
+    else:
+        return torch.matmul(posw, t_mtx.t())[None, ...]
+
+
+def get_mv_matrix(elev, azim, camera_distance, center=None):
+    elev = -elev
+
+    elev_rad = math.radians(elev)
+    azim_rad = math.radians(azim)
+
+    camera_position = np.array(
+        [
+            camera_distance * math.cos(elev_rad) * math.cos(azim_rad),
+            camera_distance * math.cos(elev_rad) * math.sin(azim_rad),
+            camera_distance * math.sin(elev_rad),
+        ]
+    )
+
+    if center is None:
+        center = np.array([0, 0, 0])
+    else:
+        center = np.array(center)
+
+    lookat = center - camera_position
+    lookat = lookat / np.linalg.norm(lookat)
+
+    up = np.array([0, 0, 1.0])
+    right = np.cross(lookat, up)
+    right = right / np.linalg.norm(right)
+    up = np.cross(right, lookat)
+    up = up / np.linalg.norm(up)
+
+    c2w = np.concatenate(
+        [np.stack([right, up, -lookat], axis=-1), camera_position[:, None]],
+        axis=-1,
+    )
+
+    w2c = np.zeros((4, 4))
+    w2c[:3, :3] = np.transpose(c2w[:3, :3], (1, 0))
+    w2c[:3, 3:] = -np.matmul(np.transpose(c2w[:3, :3], (1, 0)), c2w[:3, 3:])
+    w2c[3, 3] = 1.0
+
+    return w2c.astype(np.float32)
+
+
+def stride_from_shape(shape):
+    stride = [1]
+    for x in reversed(shape[1:]):
+        stride.append(stride[-1] * x)
+    return list(reversed(stride))
+
+
+def scatter_add_nd_with_count(input, count, indices, values, weights=None):
+    # input: [..., C], D dimension + C channel
+    # count: [..., 1], D dimension
+    # indices: [N, D], long
+    # values: [N, C]
+
+    D = indices.shape[-1]
+    C = input.shape[-1]
+    size = input.shape[:-1]
+    stride = stride_from_shape(size)
+
+    assert len(size) == D
+
+    input = input.view(-1, C)  # [HW, C]
+    count = count.view(-1, 1)
+
+    flatten_indices = (
+        indices * torch.tensor(stride, dtype=torch.long, device=indices.device)
+    ).sum(
+        -1
+    )  # [N]
+
+    if weights is None:
+        weights = torch.ones_like(values[..., :1])
+
+    input.scatter_add_(0, flatten_indices.unsqueeze(1).repeat(1, C), values)
+    count.scatter_add_(0, flatten_indices.unsqueeze(1), weights)
+
+    return input.view(*size, C), count.view(*size, 1)
+
+
+def linear_grid_put_2d(H, W, coords, values, return_count=False):
+    # coords: [N, 2], float in [0, 1]
+    # values: [N, C]
+
+    C = values.shape[-1]
+
+    indices = coords * torch.tensor(
+        [H - 1, W - 1], dtype=torch.float32, device=coords.device
+    )
+    indices_00 = indices.floor().long()  # [N, 2]
+    indices_00[:, 0].clamp_(0, H - 2)
+    indices_00[:, 1].clamp_(0, W - 2)
+    indices_01 = indices_00 + torch.tensor(
+        [0, 1], dtype=torch.long, device=indices.device
+    )
+    indices_10 = indices_00 + torch.tensor(
+        [1, 0], dtype=torch.long, device=indices.device
+    )
+    indices_11 = indices_00 + torch.tensor(
+        [1, 1], dtype=torch.long, device=indices.device
+    )
+
+    h = indices[..., 0] - indices_00[..., 0].float()
+    w = indices[..., 1] - indices_00[..., 1].float()
+    w_00 = (1 - h) * (1 - w)
+    w_01 = (1 - h) * w
+    w_10 = h * (1 - w)
+    w_11 = h * w
+
+    result = torch.zeros(
+        H, W, C, device=values.device, dtype=values.dtype
+    )  # [H, W, C]
+    count = torch.zeros(
+        H, W, 1, device=values.device, dtype=values.dtype
+    )  # [H, W, 1]
+    weights = torch.ones_like(values[..., :1])  # [N, 1]
+
+    result, count = scatter_add_nd_with_count(
+        result,
+        count,
+        indices_00,
+        values * w_00.unsqueeze(1),
+        weights * w_00.unsqueeze(1),
+    )
+    result, count = scatter_add_nd_with_count(
+        result,
+        count,
+        indices_01,
+        values * w_01.unsqueeze(1),
+        weights * w_01.unsqueeze(1),
+    )
+    result, count = scatter_add_nd_with_count(
+        result,
+        count,
+        indices_10,
+        values * w_10.unsqueeze(1),
+        weights * w_10.unsqueeze(1),
+    )
+    result, count = scatter_add_nd_with_count(
+        result,
+        count,
+        indices_11,
+        values * w_11.unsqueeze(1),
+        weights * w_11.unsqueeze(1),
+    )
+
+    if return_count:
+        return result, count
+
+    mask = count.squeeze(-1) > 0
+    result[mask] = result[mask] / count[mask].repeat(1, C)
+
+    return result
+
+
+def meshVerticeInpaint_smooth(texture, mask, vtx_pos, vtx_uv, pos_idx, uv_idx):
+    texture_height, texture_width, texture_channel = texture.shape
+    vtx_num = vtx_pos.shape[0]
+
+    vtx_mask = np.zeros(vtx_num, dtype=np.float32)
+    vtx_color = [
+        np.zeros(texture_channel, dtype=np.float32) for _ in range(vtx_num)
+    ]
+    uncolored_vtxs = []
+    G = [[] for _ in range(vtx_num)]
+
+    for i in range(uv_idx.shape[0]):
+        for k in range(3):
+            vtx_uv_idx = uv_idx[i, k]
+            vtx_idx = pos_idx[i, k]
+            uv_v = int(round(vtx_uv[vtx_uv_idx, 0] * (texture_width - 1)))
+            uv_u = int(
+                round((1.0 - vtx_uv[vtx_uv_idx, 1]) * (texture_height - 1))
+            )
+            if mask[uv_u, uv_v] > 0:
+                vtx_mask[vtx_idx] = 1.0
+                vtx_color[vtx_idx] = texture[uv_u, uv_v]
+            else:
+                uncolored_vtxs.append(vtx_idx)
+            G[pos_idx[i, k]].append(pos_idx[i, (k + 1) % 3])
+
+    smooth_count = 2
+    last_uncolored_vtx_count = 0
+    while smooth_count > 0:
+        uncolored_vtx_count = 0
+        for vtx_idx in uncolored_vtxs:
+            sum_color = np.zeros(texture_channel, dtype=np.float32)
+            total_weight = 0.0
+            vtx_0 = vtx_pos[vtx_idx]
+            for connected_idx in G[vtx_idx]:
+                if vtx_mask[connected_idx] > 0:
+                    vtx1 = vtx_pos[connected_idx]
+                    dist = np.sqrt(np.sum((vtx_0 - vtx1) ** 2))
+                    dist_weight = 1.0 / max(dist, 1e-4)
+                    dist_weight *= dist_weight
+                    sum_color += vtx_color[connected_idx] * dist_weight
+                    total_weight += dist_weight
+            if total_weight > 0:
+                vtx_color[vtx_idx] = sum_color / total_weight
+                vtx_mask[vtx_idx] = 1.0
+            else:
+                uncolored_vtx_count += 1
+
+        if last_uncolored_vtx_count == uncolored_vtx_count:
+            smooth_count -= 1
+        else:
+            smooth_count += 1
+        last_uncolored_vtx_count = uncolored_vtx_count
+
+    new_texture = texture.copy()
+    new_mask = mask.copy()
+    for face_idx in range(uv_idx.shape[0]):
+        for k in range(3):
+            vtx_uv_idx = uv_idx[face_idx, k]
+            vtx_idx = pos_idx[face_idx, k]
+            if vtx_mask[vtx_idx] == 1.0:
+                uv_v = int(round(vtx_uv[vtx_uv_idx, 0] * (texture_width - 1)))
+                uv_u = int(
+                    round((1.0 - vtx_uv[vtx_uv_idx, 1]) * (texture_height - 1))
+                )
+                new_texture[uv_u, uv_v] = vtx_color[vtx_idx]
+                new_mask[uv_u, uv_v] = 255
+
+    return new_texture, new_mask
+
+
+def mesh_uv_wrap(mesh):
+    if isinstance(mesh, trimesh.Scene):
+        mesh = mesh.dump(concatenate=True)
+
+    if len(mesh.faces) > 500000000:
+        raise ValueError(
+            "The mesh has more than 500,000,000 faces, which is not supported."
+        )
+
+    vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
+
+    mesh.vertices = mesh.vertices[vmapping]
+    mesh.faces = indices
+    mesh.visual.uv = uvs
+
+    return mesh
+
+
+class MeshRender:
+    def __init__(
+        self,
+        camera_distance=1.45,
+        default_resolution=1024,
+        texture_size=1024,
+        use_antialias=True,
+        max_mip_level=None,
+        filter_mode="linear",
+        bake_mode="linear",
+        raster_mode="cr",
+        device="cuda",
+    ):
+
+        self.device = device
+
+        self.set_default_render_resolution(default_resolution)
+        self.set_default_texture_resolution(texture_size)
+
+        self.camera_distance = camera_distance
+        self.use_antialias = use_antialias
+        self.max_mip_level = max_mip_level
+        self.filter_mode = filter_mode
+
+        self.bake_angle_thres = 75
+        self.bake_unreliable_kernel_size = int(
+            (2 / 512)
+            * max(self.default_resolution[0], self.default_resolution[1])
+        )
+        self.bake_mode = bake_mode
+
+        self.raster_mode = raster_mode
+        if self.raster_mode == "cr":
+            import custom_rasterizer as cr
+
+            self.raster = cr
+        else:
+            raise f"No raster named {self.raster_mode}"
+
+        fov = 30
+        self.camera_proj_mat = get_perspective_projection_matrix(
+            fov,
+            self.default_resolution[1] / self.default_resolution[0],
+            0.01,
+            100.0,
+        )
+
+    def raster_rasterize(
+        self, pos, tri, resolution, ranges=None, grad_db=True
+    ):
+
+        if self.raster_mode == "cr":
+            rast_out_db = None
+            if pos.dim() == 2:
+                pos = pos.unsqueeze(0)
+            findices, barycentric = self.raster.rasterize(pos, tri, resolution)
+            rast_out = torch.cat((barycentric, findices.unsqueeze(-1)), dim=-1)
+            rast_out = rast_out.unsqueeze(0)
+        else:
+            raise f"No raster named {self.raster_mode}"
+
+        return rast_out, rast_out_db
+
+    def raster_interpolate(
+        self, uv, rast_out, uv_idx, rast_db=None, diff_attrs=None
+    ):
+
+        if self.raster_mode == "cr":
+            textd = None
+            barycentric = rast_out[0, ..., :-1]
+            findices = rast_out[0, ..., -1]
+            if uv.dim() == 2:
+                uv = uv.unsqueeze(0)
+            textc = self.raster.interpolate(uv, findices, barycentric, uv_idx)
+        else:
+            raise f"No raster named {self.raster_mode}"
+
+        return textc, textd
+
+    def load_mesh(
+        self,
+        mesh,
+    ):
+        vtx_pos, pos_idx, vtx_uv, uv_idx, texture_data = load_mesh(mesh)
+        self.mesh_copy = mesh
+        self.set_mesh(
+            vtx_pos,
+            pos_idx,
+            vtx_uv=vtx_uv,
+            uv_idx=uv_idx,
+        )
+        if texture_data is not None:
+            self.set_texture(texture_data)
+
+    def save_mesh(self):
+        texture_data = self.get_texture()
+        texture_data = Image.fromarray((texture_data * 255).astype(np.uint8))
+        return save_mesh(self.mesh_copy, texture_data)
+
+    def set_mesh(
+        self,
+        vtx_pos,
+        pos_idx,
+        vtx_uv=None,
+        uv_idx=None,
+    ):
+
+        self.vtx_pos = torch.from_numpy(vtx_pos).to(self.device).float()
+        self.pos_idx = torch.from_numpy(pos_idx).to(self.device).to(torch.int)
+        if (vtx_uv is not None) and (uv_idx is not None):
+            self.vtx_uv = torch.from_numpy(vtx_uv).to(self.device).float()
+            self.uv_idx = (
+                torch.from_numpy(uv_idx).to(self.device).to(torch.int)
+            )
+        else:
+            self.vtx_uv = None
+            self.uv_idx = None
+
+        self.vtx_pos[:, [0, 1]] = -self.vtx_pos[:, [0, 1]]
+        self.vtx_pos[:, [1, 2]] = self.vtx_pos[:, [2, 1]]
+        if (vtx_uv is not None) and (uv_idx is not None):
+            self.vtx_uv[:, 1] = 1.0 - self.vtx_uv[:, 1]
+
+    def set_texture(self, tex):
+        if isinstance(tex, np.ndarray):
+            tex = Image.fromarray((tex * 255).astype(np.uint8))
+        elif isinstance(tex, torch.Tensor):
+            tex = tex.cpu().numpy()
+            tex = Image.fromarray((tex * 255).astype(np.uint8))
+
+        tex = tex.resize(self.texture_size).convert("RGB")
+        tex = np.array(tex) / 255.0
+        self.tex = torch.from_numpy(tex).to(self.device)
+        self.tex = self.tex.float()
+
+    def set_default_render_resolution(self, default_resolution):
+        if isinstance(default_resolution, int):
+            default_resolution = (default_resolution, default_resolution)
+        self.default_resolution = default_resolution
+
+    def set_default_texture_resolution(self, texture_size):
+        if isinstance(texture_size, int):
+            texture_size = (texture_size, texture_size)
+        self.texture_size = texture_size
+
+    def get_mesh(self):
+        vtx_pos = self.vtx_pos.cpu().numpy()
+        pos_idx = self.pos_idx.cpu().numpy()
+        vtx_uv = self.vtx_uv.cpu().numpy()
+        uv_idx = self.uv_idx.cpu().numpy()
+
+        # 坐标变换的逆变换
+        vtx_pos[:, [1, 2]] = vtx_pos[:, [2, 1]]
+        vtx_pos[:, [0, 1]] = -vtx_pos[:, [0, 1]]
+
+        vtx_uv[:, 1] = 1.0 - vtx_uv[:, 1]
+        return vtx_pos, pos_idx, vtx_uv, uv_idx
+
+    def get_texture(self):
+        return self.tex.cpu().numpy()
+
+    def render_sketch_from_depth(self, depth_image):
+        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)
+        combined_edges = depth_edges
+        sketch_image = (
+            torch.from_numpy(combined_edges).to(depth_image.device).float()
+            / 255.0
+        )
+        sketch_image = sketch_image.unsqueeze(-1)
+        return sketch_image
+
+    def back_project(
+        self, image, elev, azim, camera_distance=None, center=None, method=None
+    ):
+        if isinstance(image, Image.Image):
+            image = torch.tensor(np.array(image) / 255.0)
+        elif isinstance(image, np.ndarray):
+            image = torch.tensor(image)
+        if image.dim() == 2:
+            image = image.unsqueeze(-1)
+        image = image.float().to(self.device)
+        resolution = image.shape[:2]
+        channel = image.shape[-1]
+        texture = torch.zeros(self.texture_size + (channel,)).to(self.device)
+        cos_map = torch.zeros(self.texture_size + (1,)).to(self.device)
+
+        proj = self.camera_proj_mat
+        r_mv = get_mv_matrix(
+            elev=elev,
+            azim=azim,
+            camera_distance=(
+                self.camera_distance
+                if camera_distance is None
+                else camera_distance
+            ),
+            center=center,
+        )
+        pos_camera = transform_pos(r_mv, self.vtx_pos, keepdim=True)
+        pos_clip = transform_pos(proj, pos_camera)
+        pos_camera = pos_camera[:, :3] / pos_camera[:, 3:4]
+        v0 = pos_camera[self.pos_idx[:, 0], :]
+        v1 = pos_camera[self.pos_idx[:, 1], :]
+        v2 = pos_camera[self.pos_idx[:, 2], :]
+        face_normals = F.normalize(
+            torch.cross(v1 - v0, v2 - v0, dim=-1), dim=-1
+        )
+        vertex_normals = trimesh.geometry.mean_vertex_normals(
+            vertex_count=self.vtx_pos.shape[0],
+            faces=self.pos_idx.cpu(),
+            face_normals=face_normals.cpu(),
+        )
+        vertex_normals = (
+            torch.from_numpy(vertex_normals)
+            .float()
+            .to(self.device)
+            .contiguous()
+        )
+        tex_depth = pos_camera[:, 2].reshape(1, -1, 1).contiguous()
+        rast_out, rast_out_db = self.raster_rasterize(
+            pos_clip, self.pos_idx, resolution=resolution
+        )
+        visible_mask = torch.clamp(rast_out[..., -1:], 0, 1)[0, ...]
+
+        normal, _ = self.raster_interpolate(
+            vertex_normals[None, ...], rast_out, self.pos_idx
+        )
+        normal = normal[0, ...]
+
+        uv, _ = self.raster_interpolate(
+            self.vtx_uv[None, ...], rast_out, self.uv_idx
+        )
+        depth, _ = self.raster_interpolate(tex_depth, rast_out, self.pos_idx)
+        depth = depth[0, ...]
+
+        depth_max, depth_min = (
+            depth[visible_mask > 0].max(),
+            depth[visible_mask > 0].min(),
+        )
+        depth_normalized = (depth - depth_min) / (depth_max - depth_min)
+        depth_image = depth_normalized * visible_mask  # Mask out background.
+
+        sketch_image = self.render_sketch_from_depth(depth_image)
+
+        cv2.imwrite("d_depth.png", depth_image.cpu().numpy() * 255)
+        cv2.imwrite("d_normal.png", normal.cpu().numpy() * 255)
+        cv2.imwrite(
+            "d_image.png", image.cpu().numpy()[..., :3][..., ::-1] * 255
+        )
+        cv2.imwrite("d_sketch_image.png", sketch_image.cpu().numpy() * 255)
+        cv2.imwrite("d_uv1.png", uv.cpu().numpy()[0, ..., 0] * 255)
+        cv2.imwrite("d_uv2.png", uv.cpu().numpy()[0, ..., 1] * 255)
+        # p uv[0,...,0].mean(axis=0)
+        # import pdb; pdb.set_trace()
+
+        # depth_image = None
+        # normal = None
+        # image = None
+
+        sketch_image = self.render_sketch_from_depth(depth_image)
+        channel = image.shape[-1]
+
+        lookat = torch.tensor([[0, 0, -1]], device=self.device)
+        cos_image = torch.nn.functional.cosine_similarity(
+            lookat, normal.view(-1, 3)
+        )
+        cos_image = cos_image.view(normal.shape[0], normal.shape[1], 1)
+
+        cos_thres = np.cos(self.bake_angle_thres / 180 * np.pi)
+        cos_image[cos_image < cos_thres] = 0
+
+        # shrink
+        kernel_size = self.bake_unreliable_kernel_size * 2 + 1
+        kernel = torch.ones(
+            (1, 1, kernel_size, kernel_size), dtype=torch.float32
+        ).to(sketch_image.device)
+
+        visible_mask = visible_mask.permute(2, 0, 1).unsqueeze(0).float()
+        visible_mask = F.conv2d(
+            1.0 - visible_mask, kernel, padding=kernel_size // 2
+        )
+        visible_mask = 1.0 - (visible_mask > 0).float()  # 二值化
+        visible_mask = visible_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=kernel_size // 2)
+        sketch_image = (sketch_image > 0).float()  # 二值化
+        sketch_image = sketch_image.squeeze(0).permute(1, 2, 0)
+        visible_mask = visible_mask * (sketch_image < 0.5)
+
+        cos_image[visible_mask == 0] = 0
+        proj_mask = (visible_mask != 0).view(-1)
+        uv = uv.squeeze(0).contiguous().view(-1, 2)[proj_mask]
+        image = image.squeeze(0).contiguous().view(-1, channel)[proj_mask]
+        cos_image = cos_image.contiguous().view(-1, 1)[proj_mask]
+        sketch_image = sketch_image.contiguous().view(-1, 1)[proj_mask]
+        import pdb
+
+        pdb.set_trace()
+        texture = linear_grid_put_2d(
+            self.texture_size[1], self.texture_size[0], uv[..., [1, 0]], image
+        )
+        cos_map = linear_grid_put_2d(
+            self.texture_size[1],
+            self.texture_size[0],
+            uv[..., [1, 0]],
+            cos_image,
+        )
+        boundary_map = linear_grid_put_2d(
+            self.texture_size[1],
+            self.texture_size[0],
+            uv[..., [1, 0]],
+            sketch_image,
+        )
+
+        return texture, cos_map, boundary_map
+
+    @torch.no_grad()
+    def fast_bake_texture(self, textures, cos_maps):
+
+        channel = textures[0].shape[-1]
+        texture_merge = torch.zeros(self.texture_size + (channel,)).to(
+            self.device
+        )
+        trust_map_merge = torch.zeros(self.texture_size + (1,)).to(self.device)
+        for texture, cos_map in zip(textures, cos_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, texture, mask):
+
+        if isinstance(texture, torch.Tensor):
+            texture_np = texture.cpu().numpy()
+        elif isinstance(texture, np.ndarray):
+            texture_np = texture
+        elif isinstance(texture, Image.Image):
+            texture_np = np.array(texture) / 255.0
+
+        vtx_pos, pos_idx, vtx_uv, uv_idx = self.get_mesh()
+
+        texture_np, mask = meshVerticeInpaint_smooth(
+            texture_np, mask, vtx_pos, vtx_uv, pos_idx, uv_idx
+        )
+
+        texture_np = cv2.inpaint(
+            (texture_np * 255).astype(np.uint8), 255 - mask, 3, cv2.INPAINT_NS
+        )
+
+        return texture_np
+
+
+def get_images_from_file(img_path: str, img_size: int) -> list[np.array]:
+    input_image = Image.open(img_path)
+    view_images = np.array(input_image)
+    view_images = np.concatenate(
+        [view_images[:img_size, ...], view_images[img_size:, ...]], axis=1
+    )
+    images = np.split(view_images, view_images.shape[1] // img_size, axis=1)
+
+    return images
+
+
+def bake_from_multiview(
+    render, views, camera_elevs, camera_azims, view_weights, method="fast"
+):
+    project_textures, project_weighted_cos_maps = [], []
+    project_boundary_maps = []
+    for view, camera_elev, camera_azim, weight in zip(
+        views, camera_elevs, camera_azims, view_weights
+    ):
+        project_texture, project_cos_map, project_boundary_map = (
+            render.back_project(view, camera_elev, camera_azim)
+        )
+        project_cos_map = weight * (project_cos_map**4)
+        project_textures.append(project_texture)
+        project_weighted_cos_maps.append(project_cos_map)
+        project_boundary_maps.append(project_boundary_map)
+
+    if method == "fast":
+        texture, ori_trust_map = render.fast_bake_texture(
+            project_textures, project_weighted_cos_maps
+        )
+    else:
+        raise f"no method {method}"
+
+    return texture, ori_trust_map > 1e-8
+
+
+def post_process(texture: np.ndarray, iter: int = 2) -> np.ndarray:
+    for _ in range(iter):
+        texture = cv2.fastNlMeansDenoisingColored(texture, None, 11, 11, 9, 25)
+        texture = cv2.bilateralFilter(
+            texture, d=7, sigmaColor=80, sigmaSpace=80
+        )
+
+    return texture
+
+
+class Image_Super_Net:
+    def __init__(self, device="cuda"):
+        from diffusers import StableDiffusionUpscalePipeline
+
+        self.up_pipeline_x4 = StableDiffusionUpscalePipeline.from_pretrained(
+            "stabilityai/stable-diffusion-x4-upscaler",
+            torch_dtype=torch.float16,
+        ).to(device)
+        self.up_pipeline_x4.set_progress_bar_config(disable=True)
+
+    def __call__(self, image, prompt=""):
+        with torch.no_grad():
+            upscaled_image = self.up_pipeline_x4(
+                prompt=[prompt],
+                image=image,
+                num_inference_steps=10,
+            ).images[0]
+
+        return upscaled_image
+
+
+class Image_GANNet:
+    def __init__(self, outscale: int):
+        from realesrgan import RealESRGANer
+        from basicsr.archs.rrdbnet_arch import RRDBNet
+
+        self.outscale = outscale
+        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="/home/users/xinjie.wang/xinjie/Real-ESRGAN/weights/RealESRGAN_x4plus.pth",
+            model=model,
+            pre_pad=0,
+            half=True,
+        )
+
+    def __call__(self, image: Union[Image.Image, np.ndarray]) -> Image.Image:
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        output, _ = self.upsampler.enhance(image, outscale=self.outscale)
+
+        return Image.fromarray(output)
+
+
+if __name__ == "__main__":
+    device = "cuda"
+
+    # super_model = Image_Super_Net(device)
+    super_model = Image_GANNet(outscale=4)
+
+    selected_camera_elevs = [20, 20, 20, -10, -10, -10]
+    selected_camera_azims = [-180, -60, 60, -120, 0, 120]
+    selected_view_weights = [1, 0.2, 0.2, 0.2, 1, 0.2]
+    # selected_view_weights = [1, 0.1, 0.5, 0.1, 0.05, 0.05]
+
+    multiviews = get_images_from_file(
+        "scripts/apps/texture_sessions/mfq4e7u4ko/multi_view/color_sample1.png",
+        512,
+    )
+    target_image_size = (2048, 2048)
+
+    render = MeshRender(
+        camera_distance=5,
+        default_resolution=2048,
+        texture_size=2048,
+    )
+
+    mesh = trimesh.load("scripts/apps/assets/example_texture/meshes/robot.obj")
+    from asset3d_gen.data.utils import normalize_vertices_array
+
+    mesh.vertices, scale, center = normalize_vertices_array(mesh.vertices)
+    mesh = mesh_uv_wrap(mesh)
+    render.load_mesh(mesh)
+
+    # multiviews = [Image.fromarray(img) for img in multiviews]
+    # multiviews = [Image.fromarray(img).convert("RGB") for img in multiviews]
+    # for idx, img in enumerate(multiviews):
+    #     img.save(f"robot/raw/res_{idx}.png")
+
+    multiviews = [super_model(img) for img in multiviews]
+    multiviews = [img.convert("RGB") for img in multiviews]
+    for idx, img in enumerate(multiviews):
+        img.save(f"robot/super_gan_res_{idx}.png")
+
+    texture, mask = bake_from_multiview(
+        render,
+        multiviews,
+        selected_camera_elevs,
+        selected_camera_azims,
+        selected_view_weights,
+    )
+
+    texture_np = (texture.cpu().numpy() * 255).astype(np.uint8)[..., :3][
+        ..., ::-1
+    ]
+    cv2.imwrite("robot/raw_texture.png", texture_np)
+    print("texture done.")
+
+    mask_np = (mask.squeeze(-1).cpu().numpy() * 255).astype(np.uint8)
+    texture_np = render.uv_inpaint(texture, mask_np)
+    cv2.imwrite("robot/inpaint_texture.png", texture_np[..., ::-1])
+    # texture_np = post_process(texture_np, 2)
+    # cv2.imwrite("robot/inpaint_conv_texture.png", texture_np[..., ::-1])
+    print("inpaint done.")
+
+    texture = torch.tensor(texture_np / 255).float().to(texture.device)
+    render.set_texture(texture)
+    textured_mesh = render.save_mesh()
+    _ = textured_mesh.export("robot/robot.obj")

asset3d_gen/data/backup/gpt_qwen.py

@@ -0,0 +1,70 @@
+import torch
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
+from qwen_vl_utils import process_vision_info
+import os 
+os.environ["https_proxy"] = "10.9.0.31:8838"
+
+
+# # default: Load the model on the available device(s)
+# model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+#     "Qwen/Qwen2.5-VL-7B-Instruct", torch_dtype="auto", device_map="auto"
+# )
+
+# We recommend enabling flash_attention_2 for better acceleration and memory saving, especially in multi-image and video scenarios.
+model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    "Qwen/Qwen2.5-VL-7B-Instruct",
+    torch_dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2",
+    device_map="auto",
+)
+
+
+# default processer
+processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
+
+# The default range for the number of visual tokens per image in the model is 4-16384.
+# You can set min_pixels and max_pixels according to your needs, such as a token range of 256-1280, to balance performance and cost.
+# min_pixels = 256*28*28
+# max_pixels = 1280*28*28
+# processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
+
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "image",
+                "image": "outputs/text2image/demo_objects/bed/sample_0.jpg",
+            },
+            {
+                "type": "image",
+                "image": "outputs/imageto3d/v2/cups/sample_69/URDF_sample_69/qa_renders/image_color/003.png",
+            },
+            {"type": "text", "text": "Describe the secend image."},
+        ],
+    }
+]
+
+# Preparation for inference
+text = processor.apply_chat_template(
+    messages, tokenize=False, add_generation_prompt=True
+)
+image_inputs, video_inputs = process_vision_info(messages)
+inputs = processor(
+    text=[text],
+    images=image_inputs,
+    videos=video_inputs,
+    padding=True,
+    return_tensors="pt",
+)
+inputs = inputs.to("cuda")
+
+# Inference: Generation of the output
+generated_ids = model.generate(**inputs, max_new_tokens=128)
+generated_ids_trimmed = [
+    out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+]
+output_text = processor.batch_decode(
+    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+)
+print(output_text)

asset3d_gen/data/backup/quat.py

@@ -0,0 +1,49 @@
+import numpy as np
+
+def quaternion_rotation_x_counterclockwise(angle_degrees):
+    angle_radians = np.radians(angle_degrees)
+    w = np.cos(angle_radians / 2)
+    x = np.sin(angle_radians / 2)
+    y, z = 0.0, 0.0
+    return np.array([x, y, z, w]).round(4).tolist()
+
+
+def quaternion_rotation_y_counterclockwise(angle_degrees):
+    angle_radians = np.radians(angle_degrees)
+    w = np.cos(angle_radians / 2)
+    y = np.sin(angle_radians / 2)
+    x, z = 0.0, 0.0
+    return np.array([x, y, z, w]).round(4).tolist()
+
+
+def quaternion_rotation_z_counterclockwise(angle_degrees):
+    angle_radians = np.radians(angle_degrees)
+    w = np.cos(angle_radians / 2)
+    z = np.sin(angle_radians / 2)
+    x, y = 0.0, 0.0
+    return np.array([x, y, z, w]).round(4).tolist()
+
+
+def quaternion_multiply(q1, q2):
+    x1, y1, z1, w1 = q1
+    x2, y2, z2, w2 = q2
+    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 np.array([w, x, y, z])
+
+    
+    
+angle = 180
+
+print(f"X轴逆时针旋转{angle}度: {quaternion_rotation_x_counterclockwise(angle)}")
+print(f"Y轴逆时针旋转{angle}度: {quaternion_rotation_y_counterclockwise(angle)}")
+print(f"Z轴逆时针旋转{angle}度: {quaternion_rotation_z_counterclockwise(angle)}")
+
+
+q_1 = np.array([1.0, 0.0, 0.0, 0.0]) 
+q_2 = np.array([0.0, 0.0, 1.0, 0.0]) 
+
+q_total = quaternion_multiply(q_2, q_1)
+print(q_total.round(4).tolist())

asset3d_gen/data/datasets.py

@@ -0,0 +1,239 @@
+import json
+import logging
+import os
+import random
+from typing import Any, Callable, Dict, List, Tuple, Union
+
+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",
+]
+
+
+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
+
+
+if __name__ == "__main__":
+    index_file = "/horizon-bucket/robot_lab/users/xinjie.wang/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="./"
+    )

asset3d_gen/data/differentiable_render.py

@@ -0,0 +1,520 @@
+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 torch
+from PIL import Image
+from tqdm import tqdm
+from asset3d_gen.data.utils import (
+    CameraSetting,
+    DiffrastRender,
+    RenderItems,
+    as_list,
+    calc_vertex_normals,
+    import_kaolin_mesh,
+    init_kal_camera,
+    normalize_vertices_array,
+    render_pbr,
+    save_images,
+)
+
+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__)
+
+
+def create_gif_from_images(images, output_path, fps=10):
+    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}")
+
+
+def create_mp4_from_images(images, output_path, fps=10, prompt=None):
+    font = cv2.FONT_HERSHEY_SIMPLEX  # 字体样式
+    font_scale = 0.5  # 字体大小
+    font_thickness = 1  # 字体粗细
+    color = (255, 255, 255)  # 文字颜色（白色）
+    position = (20, 25)  # 左上角坐标 (x, y)
+
+    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}")
+
+
+class ImageRender(object):
+    def __init__(
+        self,
+        render_items: list[RenderItems],
+        camera_params: CameraSetting,
+        recompute_vtx_normal: bool = True,
+        device: str = "cuda",
+        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_params.device = device
+        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 = 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]:
+        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
+                )
+                # rendered_inv_view_normals = renderer.transform_normal(rendered_view_normals, torch.linalg.inv(mv), masks, to_view=False)  # noqa
+                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="+",
+        required=True,
+        help="Paths to the mesh files for rendering.",
+    )
+    parser.add_argument(
+        "--output_root",
+        type=str,
+        required=True,
+        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: 2.)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        choices=["cpu", "cuda"],
+        default="cuda",
+        help="Device to run on (default: 'cuda')",
+    )
+    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(
+        "--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(
+        "--prompts",
+        type=str,
+        nargs="+",
+        default=None,
+        help="Text prompts for the rendering.",
+    )
+
+    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_settings = 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,
+    )
+
+    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,
+    )
+    image_render.render_mesh(
+        mesh_path=args.mesh_path,
+        output_root=args.output_root,
+        uuid=args.uuid,
+        prompts=args.prompts,
+    )
+
+    return
+
+
+if __name__ == "__main__":
+    entrypoint()

asset3d_gen/data/mesh_operator.py

@@ -0,0 +1,425 @@
+import logging
+from typing import Tuple, Union
+
+import igraph
+import numpy as np
+import pyvista as pv
+import torch
+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_sequence(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):
+    """Reduce and postprocess 3D meshes, simplifying and filling holes."""
+
+    def __init__(
+        self,
+        vertices: Union[torch.Tensor, np.ndarray],
+        faces: Union[torch.Tensor, np.ndarray],
+        device: str = "cuda",
+    ) -> None:
+        self.device = device
+        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)
+        )
+
+    @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_sequence(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 = 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,
+        )
+
+    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.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

asset3d_gen/data/utils.py

@@ -0,0 +1,943 @@
+import math
+import os
+import random
+from glob import glob
+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
+
+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
+from enum import Enum
+
+import trimesh
+from kaolin.render.camera import Camera
+from torch import nn
+
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "center_points",
+    "get_points_stat",
+    "DiffrastRender",
+    "compute_cam_pts_by_az_el",
+    "compute_cam_pts_by_views",
+    "save_images",
+    "render_pbr",
+    "load_llm_models",
+    "prelabel_text_feature",
+    "calc_vertex_normals",
+    "normalize_vertices_array",
+    "load_mesh_to_unit_cube",
+    "as_list",
+    "CameraSetting",
+    "RenderItems",
+    "import_kaolin_mesh",
+    "save_mesh_with_mtl",
+    "get_images_from_grid",
+    "post_process_texture",
+]
+
+
+def get_points_stat(
+    points: torch.FloatTensor, eps: float = 1e-6
+) -> torch.FloatTensor:
+    assert (
+        len(points.shape) == 3
+    ), f"Points have unexpected shape {points.shape}"
+
+    vmin = points.min(dim=1, keepdim=True)[0]
+    vmax = points.max(dim=1, keepdim=True)[0]
+    pts_center = (vmin + vmax) / 2
+
+    pts_dim = (vmax - vmin).max(dim=-1, keepdim=True)[0].clip(min=eps)
+
+    return pts_center, pts_dim
+
+
+def center_points(
+    points: torch.FloatTensor, normalize: bool = False, eps: float = 1e-6
+) -> torch.FloatTensor:
+    vmid, den = get_points_stat(points)
+
+    res = points - vmid
+
+    if normalize:
+        res = res / den
+
+    return res
+
+
+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').
+
+    Methods:
+        _warmup(glctx): Warmup rasterization by rendering a simple triangle.
+        compute_dr_raster(vertices, faces): Rasterizes the mesh and returns
+            rasterized outputs and transformed vertices.
+        transform_vertices(vertices, matrix): Transforms the vertices using
+            the provided transformation matrix.
+        normalize_map_by_mask_separately(map, mask): Normalizes each map in
+            the batch separately using the mask.
+        normalize_map_by_mask(map, mask): Normalizes the entire map using the
+            mask, keeping the output in the range [0, 1].
+        render_position(vertices, faces): Renders the position map and
+            alpha mask from the given vertices and faces.
+        render_depth(vertices, faces): Renders the depth map and alpha
+            mask from the given vertices and faces.
+        _compute_mask(rast, vertices_clip, faces): Computes the mask from the
+            rasterization output.
+        render_global_normal(vertices, faces, vertice_normals): Renders the
+            normal map and alpha mask from the given vertices, faces, and
+            vertex normals.
+        transform_normal_to_view(normals, mat_w2c, masks): Transforms the normals
+            to the view space using the world-to-camera matrix.
+    """
+
+    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
+        # self._warmup(self._ctx)
+
+    def _warmup(self, glctx):
+        # Seem solved. https://github.com/NVlabs/nvdiffrast/issues/59
+        def tensor(*args, **kwargs):
+            return torch.tensor(*args, device=self.device, **kwargs)
+
+        pos = tensor(
+            [[[-0.8, -0.8, 0, 1], [0.8, -0.8, 0, 1], [-0.8, 0.8, 0, 1]]],
+            dtype=torch.float32,
+        )
+        tri = tensor([[0, 1, 2]], dtype=torch.int32)
+        dr.rasterize(glctx, pos, tri, resolution=[256, 256])
+
+    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)).to(vertices)
+        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
+
+    def normalize_map_by_mask(
+        self, 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 load_saved_normal(path: str) -> np.ndarray:
+    image_paths = glob(os.path.join(path, "*.jpg"))
+    images = []
+    for path in sorted(image_paths):
+        image = cv2.imread(path)
+        image = image[..., ::-1]  # rgb -> bgr
+        images.append(image)
+    images = np.stack(images, axis=0)
+
+    return images
+
+
+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)
+    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
+
+
+@dataclass
+class RenderItems(str, Enum):
+    IMAGE = "image_color"
+    ALPHA = "image_mask"
+    VIEW_NORMAL = "image_view_normal"
+    GLOBAL_NORMAL = "image_global_normal"
+    POSITION_MAP = "image_position"
+    DEPTH = "image_depth"
+    ALBEDO = "image_albedo"
+    DIFFUSE = "image_diffuse"
+
+
+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
+    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) -> Camera:
+    azimuths, elevations = compute_az_el_by_camera_params(camera_params)
+    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)
+    view_images = np.concatenate(
+        [view_images[:img_size, ...], view_images[img_size:, ...]], axis=1
+    )
+    images = np.split(view_images, view_images.shape[1] // img_size, axis=1)
+    images = [Image.fromarray(img) for img in images]
+
+    return images
+
+
+def post_process_texture(texture: np.ndarray, iter: int = 2) -> np.ndarray:
+    for _ in range(iter):
+        texture = cv2.fastNlMeansDenoisingColored(texture, None, 13, 13, 9, 27)
+        texture = cv2.bilateralFilter(
+            texture, d=9, sigmaColor=80, sigmaSpace=80
+        )
+
+    return texture

asset3d_gen/models/delight.py

@@ -0,0 +1,165 @@
+import os
+from typing import Union
+
+import cv2
+import numpy as np
+import torch
+from diffusers import (
+    EulerAncestralDiscreteScheduler,
+    StableDiffusionInstructPix2PixPipeline,
+)
+from huggingface_hub import snapshot_download
+from PIL import Image
+from asset3d_gen.models.segment import RembgRemover
+
+os.environ["https_proxy"] = "http://10.9.0.31:8838"
+
+
+__all__ = [
+    "DelightingModel",
+]
+
+
+class DelightingModel(object):
+    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.bg_remover = RembgRemover()
+
+        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)
+
+        pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+            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)
+        pipeline.enable_model_cpu_offload()
+        pipeline.enable_xformers_memory_efficient_attention()
+        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
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        image: Union[str, np.ndarray, Image.Image],
+        preprocess: bool = False,
+        target_wh: tuple[int, int] = None,
+    ) -> Image.Image:
+        if isinstance(image, str):
+            image = Image.open(image)
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+
+        if preprocess:
+            image = self.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(
+        # model_path="/horizon-bucket/robot_lab/users/xinjie.wang/weights/hunyuan3d-delight-v2-0"  # noqa
+    )
+    image_path = "scripts/apps/assets/example_image/room_bottle_002.jpeg"
+    image = delighting_model(
+        image_path, preprocess=True, target_wh=(512, 512)
+    )  # noqa
+    image.save("delight.png")
+
+    # image_path = "asset3d_gen/scripts/test_robot.png"
+    # image = delighting_model(image_path)
+    # image.save("delighting_image_a2.png")

asset3d_gen/models/gs_model.py

@@ -0,0 +1,540 @@
+import logging
+import os
+import struct
+from dataclasses import dataclass, field
+from typing import Optional, Union
+
+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 torch.nn import functional as F
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "RenderResult",
+    "GaussianOperator",
+]
+
+
+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
+
+
+@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.detach().cpu().numpy()
+            rgb = (rgb * 255).astype(np.uint8)
+            self.rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
+        if isinstance(self.depth, torch.Tensor):
+            self.depth = self.depth.detach().cpu().numpy()
+        if isinstance(self.opacity, torch.Tensor):
+            opacity = self.opacity.detach().cpu().numpy()
+            opacity = (opacity * 255).astype(np.uint8)
+            self.opacity = cv2.cvtColor(opacity, cv2.COLOR_GRAY2RGB)
+            mask = np.where(self.opacity > self.mask_threshold, 255, 0)
+            self.mask = mask[..., 0:1].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,
+    ) -> "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,
+        )
+
+    def save_to_ply(
+        self, path: str, colors: torch.Tensor = None, enable_mask: bool = False
+    ):
+        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)))
+
+        # 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")
+            f.write(b"property float nx\n")
+            f.write(b"property float ny\n")
+            f.write(b"property float nz\n")
+
+            if colors is not None:
+                for j in range(colors.shape[1]):
+                    f.write(f"property float f_dc_{j}\n".encode())
+            else:
+                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
+                f.write(struct.pack("<fff", 0, 0, 0))  # nx, ny, nz (zeros)
+
+                if colors is not None:
+                    color = colors.detach().cpu().numpy()
+                    for j in range(color.shape[1]):
+                        f_dc = (color[i, j] - 0.5) / 0.2820947917738781
+                        f.write(struct.pack("<f", f_dc))
+                else:
+                    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]))  # opacity
+
+                for data in [scales, quats]:
+                    for j in range(data.shape[1]):
+                        f.write(struct.pack("<f", data[i, j]))
+
+
+@dataclass
+class GaussianOperator(GaussianBase):
+
+    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,
+        )
+
+        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,
+        sh_degree: int = 0,
+    ) -> None:
+        gs_model = GaussianOperator.load_from_ply(in_ply, sh_degree)
+
+        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/test/debug.ply"
+    output_gs = "./debug_v3.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)

asset3d_gen/models/segment.py

@@ -0,0 +1,376 @@
+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 asset3d_gen.utils.process_media import filter_small_connected_components
+from asset3d_gen.validators.quality_checkers import ImageSegChecker
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+os.environ["https_proxy"] = "http://10.9.0.31:8838"
+
+__all__ = [
+    "resize_pil",
+    "trellis_preprocess",
+    "SAMRemover",
+    "SAMPredictor",
+    "RembgRemover",
+    "get_segmented_image",
+]
+
+
+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
+
+
+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
+
+
+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(
+    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(
+        # checkpoint="/horizon-bucket/robot_lab/users/xinjie.wang/weights/sam/sam_vit_h_4b8939.pth",  # noqa
+        model_type="vit_h",
+    )
+    remover = RembgRemover()
+    # clean_image = remover(input_image)
+    # clean_image.save(output_image)
+    get_segmented_image(
+        Image.open(input_image), remover, remover, None, "./test_seg.png"
+    )

asset3d_gen/models/super_resolution.py

@@ -0,0 +1,118 @@
+import logging
+import os
+from typing import Union
+
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download, snapshot_download
+from PIL import Image
+from asset3d_gen.data.utils import get_images_from_grid
+
+os.environ["https_proxy"] = "http://10.9.0.31:8838"
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "ImageStableSR",
+    "ImageRealESRGAN",
+]
+
+
+class ImageStableSR:
+    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()
+
+    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:
+    def __init__(self, outscale: int, model_path: str = None) -> None:
+        from basicsr.archs.rrdbnet_arch import RRDBNet
+        from realesrgan import RealESRGANer
+
+        self.outscale = outscale
+        model = RRDBNet(
+            num_in_ch=3,
+            num_out_ch=3,
+            num_feat=64,
+            num_block=23,
+            num_grow_ch=32,
+            scale=4,
+        )
+        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.upsampler = RealESRGANer(
+            scale=4,
+            model_path=model_path,
+            model=model,
+            pre_pad=0,
+            half=True,
+        )
+
+    def __call__(self, image: Union[Image.Image, np.ndarray]) -> Image.Image:
+        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.
+    # model_path = "/horizon-bucket/robot_lab/users/xinjie.wang/weights/super_resolution/RealESRGAN_x4plus.pth"  # noqa
+    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")

asset3d_gen/models/text_model.py

@@ -0,0 +1,143 @@
+import logging
+
+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",
+]
+
+
+def build_text2img_ip_pipeline(
+    ckpt_dir: str,
+    ref_scale: float,
+    device: str = "cuda",
+) -> StableDiffusionXLPipelineIP:
+    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.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:
+    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,
+) -> list[Image.Image]:
+    prompt = "Single " + prompt + ", in the center of the image"
+    prompt += ", high quality, high resolution, best quality, white background, 3D style,"  # noqa
+    logger.info(f"Processing prompt: {prompt}")
+
+    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,
+    )
+    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

asset3d_gen/models/texture_model.py

@@ -0,0 +1,91 @@
+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
+
+__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:
+    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 = "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

asset3d_gen/scripts/render_gs.py

@@ -0,0 +1,156 @@
+import argparse
+import logging
+import math
+import os
+
+import cv2
+import numpy as np
+import torch
+from tqdm import tqdm
+from asset3d_gen.data.utils import (
+    CameraSetting,
+    init_kal_camera,
+    normalize_vertices_array,
+)
+from asset3d_gen.models.gs_model import GaussianOperator
+
+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 = parser.parse_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
+
+
+def entrypoint(input_gs: str = None, output_path: str = None) -> None:
+    args = parse_args()
+    if isinstance(input_gs, str):
+        args.input_gs = input_gs
+    if isinstance(output_path, str):
+        args.output_path = output_path
+
+    # 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)
+    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,
+        )
+        images.append(color)
+
+    # Cat color images into grid image and save.
+    select_idxs = [[0, 2, 1], [5, 4, 3]]  # fix order for 6 views
+    grid_image = []
+    for row_idxs in select_idxs:
+        row_image = []
+        for row_idx in row_idxs:
+            row_image.append(images[row_idx])
+        row_image = np.concatenate(row_image, axis=1)
+        grid_image.append(row_image)
+
+    grid_image = np.concatenate(grid_image, axis=0)
+    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
+    cv2.imwrite(args.output_path, grid_image)
+    logger.info(f"Saved grid image to {args.output_path}")
+
+
+if __name__ == "__main__":
+    entrypoint()

asset3d_gen/scripts/render_mv.py

@@ -0,0 +1,185 @@
+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 asset3d_gen.data.datasets import Asset3dGenDataset
+from asset3d_gen.models.texture_model import build_texture_gen_pipe
+
+os.environ["https_proxy"] = "http://10.9.0.31:8838"
+
+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,脸部阴影,低分辨率,jpeg伪影、模糊、糟糕,黑脸,霓虹灯,高光,镜面反射"
+    )
+
+    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()

asset3d_gen/scripts/text2image.py

@@ -0,0 +1,145 @@
+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 asset3d_gen.models.text_model import (
+    build_text2img_ip_pipeline,
+    build_text2img_pipeline,
+    text2img_gen,
+)
+
+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,
+    )
+    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 = args.prompts
+    if len(prompts) == 1 and prompts[0].endswith(".txt"):
+        with open(prompts[0], "r") as f:
+            prompts = f.readlines()
+            prompts = [
+                prompt.strip() for prompt in prompts if prompt.strip() != ""
+            ]
+
+    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,
+        )
+
+        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()

asset3d_gen/utils/gpt_clients.py

@@ -0,0 +1,190 @@
+import base64
+import logging
+import os
+from io import BytesIO
+from typing import Optional
+
+from openai import AzureOpenAI, OpenAI  # pip install openai
+from PIL import Image
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    stop_after_delay,
+    wait_random_exponential,
+)
+from asset3d_gen.utils.process_media import combine_images_to_base64
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+os.environ["https_proxy"] = "10.9.0.31:8838"
+
+
+class GPTclient:
+    """A client to interact with the GPT model via OpenAI or Azure API."""
+
+    def __init__(
+        self,
+        endpoint: str,
+        api_key: str,
+        model_name: str = "yfb-gpt-4o",
+        api_version: str = None,
+        verbose: bool = False,
+    ):
+        if api_version is not None:
+            self.client = AzureOpenAI(
+                azure_endpoint=endpoint,
+                api_key=api_key,
+                api_version=api_version,
+            )
+        else:
+            self.client = OpenAI(
+                base_url=endpoint,
+                api_key=api_key,
+            )
+
+        self.endpoint = endpoint
+        self.model_name = model_name
+        self.image_formats = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif"}
+        self.verbose = verbose
+
+    @retry(
+        wait=wait_random_exponential(min=1, max=20),
+        stop=(stop_after_attempt(10) | stop_after_delay(30)),
+    )
+    def completion_with_backoff(self, **kwargs):
+        return self.client.chat.completions.create(**kwargs)
+
+    def query(
+        self,
+        text_prompt: str,
+        image_base64: Optional[list[str | Image.Image]] = None,
+        system_role: Optional[str] = None,
+    ) -> Optional[str]:
+        """Queries the GPT model with a text and optional image prompts.
+
+        Args:
+            text_prompt (str): The main text input that the model responds to.
+            image_base64 (Optional[List[str]]): A list of image base64 strings
+                or local image paths or PIL.Image to accompany the text prompt.
+            system_role (Optional[str]): Optional system-level instructions
+                that specify the behavior of the assistant.
+
+        Returns:
+            Optional[str]: The response content generated by the model based on
+                the prompt. Returns `None` if an error occurs.
+        """
+        if system_role is None:
+            system_role = "You are a highly knowledgeable assistant specializing in physics, engineering, and object properties."  # noqa
+
+        content_user = [
+            {
+                "type": "text",
+                "text": text_prompt,
+            },
+        ]
+
+        # Process images if provided
+        if image_base64 is not None:
+            image_base64 = (
+                image_base64
+                if isinstance(image_base64, list)
+                else [image_base64]
+            )
+            for img in image_base64:
+                if isinstance(img, Image.Image):
+                    buffer = BytesIO()
+                    img.save(buffer, format=img.format or "PNG")
+                    buffer.seek(0)
+                    image_binary = buffer.read()
+                    img = base64.b64encode(image_binary).decode("utf-8")
+                elif (
+                    len(os.path.splitext(img)) > 1
+                    and os.path.splitext(img)[-1].lower() in self.image_formats
+                ):
+                    if not os.path.exists(img):
+                        raise FileNotFoundError(f"Image file not found: {img}")
+                    with open(img, "rb") as f:
+                        img = base64.b64encode(f.read()).decode("utf-8")
+
+                content_user.append(
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/png;base64,{img}"},
+                    }
+                )
+
+        payload = {
+            "messages": [
+                {"role": "system", "content": system_role},
+                {"role": "user", "content": content_user},
+            ],
+            "temperature": 0.1,
+            "max_tokens": 500,
+            "top_p": 0.1,
+            "frequency_penalty": 0,
+            "presence_penalty": 0,
+            "stop": None,
+        }
+        payload.update({"model": self.model_name})
+
+        response = None
+        try:
+            response = self.completion_with_backoff(**payload)
+            response = response.choices[0].message.content
+        except Exception as e:
+            logger.error(f"Error GPTclint {self.endpoint} API call: {e}")
+            response = None
+
+        if self.verbose:
+            logger.info(f"Prompt: {text_prompt}")
+            logger.info(f"Response: {response}")
+
+        return response
+
+
+endpoint = os.environ.get("endpoint", None)
+api_key = os.environ.get("api_key", None)
+api_version = os.environ.get("api_version", None)
+if endpoint and api_key and api_version:
+    GPT_CLIENT = GPTclient(
+        endpoint=endpoint,
+        api_key=api_key,
+        api_version=api_version,
+        model_name="yfb-gpt-4o-sweden" if "sweden" in endpoint else None,
+    )
+else:
+    GPT_CLIENT = GPTclient(
+        endpoint="https://openrouter.ai/api/v1",
+        api_key="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",  # noqa
+        model_name="qwen/qwen2.5-vl-72b-instruct:free",
+    )
+
+
+if __name__ == "__main__":
+    if "openrouter" in GPT_CLIENT.endpoint:
+        response = GPT_CLIENT.query(
+            text_prompt="What is the content in each image?",
+            image_base64=combine_images_to_base64(
+                [
+                    "outputs/text2image/demo_objects/bed/sample_0.jpg",
+                    "outputs/imageto3d/v2/cups/sample_69/URDF_sample_69/qa_renders/image_color/003.png",  # noqa
+                    "outputs/text2image/demo_objects/cardboard/sample_1.jpg",
+                ]
+            ),  # input raw image_path if only one image
+        )
+        print(response)
+    else:
+        response = GPT_CLIENT.query(
+            text_prompt="What is the content in the images?",
+            image_base64=[
+                Image.open("outputs/text2image/demo_objects/bed/sample_0.jpg"),
+                Image.open(
+                    "outputs/imageto3d/v2/cups/sample_69/URDF_sample_69/qa_renders/image_color/003.png"  # noqa
+                ),
+            ],
+        )
+        print(response)
+
+        # test2: text prompt
+        response = GPT_CLIENT.query(text_prompt="What is the capital of China?")
+        print(response)

asset3d_gen/utils/process_media.py

@@ -0,0 +1,194 @@
+import base64
+import logging
+import math
+import os
+import subprocess
+from glob import glob
+from io import BytesIO
+from typing import Union
+
+import cv2
+import imageio
+import numpy as np
+import PIL.Image as Image
+from moviepy.editor import VideoFileClip, clips_array
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "render_asset3d",
+    "merge_images_video",
+    "filter_small_connected_components",
+    "filter_image_small_connected_components",
+    "combine_images_to_base64",
+]
+
+
+def render_asset3d(
+    mesh_path: str,
+    output_root: str,
+    distance: float = 5.0,
+    num_images: int = 1,
+    elevation: list[float] = (0.0,),
+    pbr_light_factor: float = 1.5,
+    return_key: str = "image_color/*",
+    output_subdir: str = "renders",
+    gen_color_mp4: bool = False,
+    gen_viewnormal_mp4: bool = False,
+    gen_glonormal_mp4: bool = False,
+    device: str = "cpu",
+) -> list[str]:
+    command = [
+        "python3",
+        "asset3d_gen/data/differentiable_render.py",
+        "--mesh_path",
+        mesh_path,
+        "--output_root",
+        output_root,
+        "--uuid",
+        output_subdir,
+        "--distance",
+        str(distance),
+        "--num_images",
+        str(num_images),
+        "--elevation",
+        *map(str, elevation),
+        "--pbr_light_factor",
+        str(pbr_light_factor),
+        "--with_mtl",
+        "--device",
+        device,
+    ]
+    if gen_color_mp4:
+        command.append("--gen_color_mp4")
+    if gen_viewnormal_mp4:
+        command.append("--gen_viewnormal_mp4")
+    if gen_glonormal_mp4:
+        command.append("--gen_glonormal_mp4")
+    try:
+        subprocess.run(command, check=True)
+    except subprocess.CalledProcessError as e:
+        logger.error(f"Error occurred during rendering: {e}.")
+
+    dst_paths = glob(os.path.join(output_root, output_subdir, return_key))
+
+    return dst_paths
+
+
+def merge_images_video(color_images, normal_images, output_path) -> None:
+    width = color_images[0].shape[1]
+    combined_video = [
+        np.hstack([rgb_img[:, : width // 2], normal_img[:, width // 2 :]])
+        for rgb_img, normal_img in zip(color_images, normal_images)
+    ]
+    imageio.mimsave(output_path, combined_video, fps=50)
+
+    return
+
+
+def merge_video_video(
+    video_path1: str, video_path2: str, output_path: str
+) -> None:
+    """Merge two videos by the left half and the right half of the videos."""
+    clip1 = VideoFileClip(video_path1)
+    clip2 = VideoFileClip(video_path2)
+
+    if clip1.size != clip2.size:
+        raise ValueError("The resolutions of the two videos do not match.")
+
+    width, height = clip1.size
+    clip1_half = clip1.crop(x1=0, y1=0, x2=width // 2, y2=height)
+    clip2_half = clip2.crop(x1=width // 2, y1=0, x2=width, y2=height)
+    final_clip = clips_array([[clip1_half, clip2_half]])
+    final_clip.write_videofile(output_path, codec="libx264")
+
+
+def filter_small_connected_components(
+    mask: Union[Image.Image, np.ndarray],
+    area_ratio: float,
+    connectivity: int = 8,
+) -> np.ndarray:
+    if isinstance(mask, Image.Image):
+        mask = np.array(mask)
+    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
+        mask,
+        connectivity=connectivity,
+    )
+
+    small_components = np.zeros_like(mask, dtype=np.uint8)
+    mask_area = (mask != 0).sum()
+    min_area = mask_area // area_ratio
+    for label in range(1, num_labels):
+        area = stats[label, cv2.CC_STAT_AREA]
+        if area < min_area:
+            small_components[labels == label] = 255
+
+    mask = cv2.bitwise_and(mask, cv2.bitwise_not(small_components))
+
+    return mask
+
+
+def filter_image_small_connected_components(
+    image: Union[Image.Image, np.ndarray],
+    area_ratio: float = 10,
+    connectivity: int = 8,
+) -> np.ndarray:
+    if isinstance(image, Image.Image):
+        image = image.convert("RGBA")
+        image = np.array(image)
+
+    mask = image[..., 3]
+    mask = filter_small_connected_components(mask, area_ratio, connectivity)
+    image[..., 3] = mask
+
+    return image
+
+
+def combine_images_to_base64(
+    images: list[str | Image.Image],
+    cat_row_col: tuple[int, int] = None,
+    target_wh: tuple[int, int] = (512, 512),
+) -> str:
+    n_images = len(images)
+    if cat_row_col is None:
+        n_col = math.ceil(math.sqrt(n_images))
+        n_row = math.ceil(n_images / n_col)
+    else:
+        n_row, n_col = cat_row_col
+
+    images = [
+        Image.open(p).convert("RGB") if isinstance(p, str) else p
+        for p in images[: n_row * n_col]
+    ]
+    images = [img.resize(target_wh) for img in images]
+
+    grid_w, grid_h = n_col * target_wh[0], n_row * target_wh[1]
+    grid = Image.new("RGB", (grid_w, grid_h), (255, 255, 255))
+
+    for idx, img in enumerate(images):
+        row, col = divmod(idx, n_col)
+        grid.paste(img, (col * target_wh[0], row * target_wh[1]))
+
+    buffer = BytesIO()
+    grid.save(buffer, format="PNG")
+
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+if __name__ == "__main__":
+    # Example usage:
+    merge_video_video(
+        "outputs/imageto3d/room_bottle7/room_bottle_007/URDF_room_bottle_007/mesh_glo_normal.mp4",  # noqa
+        "outputs/imageto3d/room_bottle7/room_bottle_007/URDF_room_bottle_007/mesh.mp4",  # noqa
+        "merge.mp4",
+    )
+
+    image_base64 = combine_images_to_base64(
+        [
+            "outputs/text2image/demo_objects/bed/sample_0.jpg",
+            "outputs/imageto3d/v2/cups/sample_69/URDF_sample_69/qa_renders/image_color/003.png",  # noqa
+            "outputs/text2image/demo_objects/cardboard/sample_1.jpg",
+        ]
+    )

asset3d_gen/utils/tags.py

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

asset3d_gen/validators/aesthetic_predictor.py

@@ -0,0 +1,136 @@
+import os
+
+import clip
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+from huggingface_hub import snapshot_download
+from PIL import Image
+
+os.environ["https_proxy"] = "http://10.9.0.31:8838"
+
+
+class AestheticPredictor:
+    """Aesthetic Score Predictor.
+
+    Args:
+        clip_model_dir (str): Path to the directory of the CLIP model.
+        sac_model_path (str): Path to the pre-trained SAC model.
+        device (str): Device to use for computation ("cuda" or "cpu").
+    """
+
+    def __init__(self, clip_model_dir=None, sac_model_path=None, device=None):
+
+        self.device = device or (
+            "cuda" if torch.cuda.is_available() else "cpu"
+        )
+
+        if clip_model_dir is None:
+            model_path = snapshot_download(
+                repo_id="xinjjj/RoboAssetGen", allow_patterns="aesthetic/*"
+            )
+            suffix = "aesthetic"
+            model_path = snapshot_download(
+                repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
+            )
+            clip_model_dir = os.path.join(model_path, suffix)
+
+        if sac_model_path is None:
+            model_path = snapshot_download(
+                repo_id="xinjjj/RoboAssetGen", allow_patterns="aesthetic/*"
+            )
+            suffix = "aesthetic"
+            model_path = snapshot_download(
+                repo_id="xinjjj/RoboAssetGen", allow_patterns=f"{suffix}/*"
+            )
+            sac_model_path = os.path.join(
+                model_path, suffix, "sac+logos+ava1-l14-linearMSE.pth"
+            )
+
+        self.clip_model, self.preprocess = self._load_clip_model(
+            clip_model_dir
+        )
+        self.sac_model = self._load_sac_model(sac_model_path, input_size=768)
+
+    class MLP(pl.LightningModule):  # noqa
+        def __init__(self, input_size):
+            super().__init__()
+            self.layers = nn.Sequential(
+                nn.Linear(input_size, 1024),
+                nn.Dropout(0.2),
+                nn.Linear(1024, 128),
+                nn.Dropout(0.2),
+                nn.Linear(128, 64),
+                nn.Dropout(0.1),
+                nn.Linear(64, 16),
+                nn.Linear(16, 1),
+            )
+
+        def forward(self, x):
+            return self.layers(x)
+
+    @staticmethod
+    def normalized(a, axis=-1, order=2):
+        """Normalize the array to unit norm."""
+        l2 = np.atleast_1d(np.linalg.norm(a, order, axis))
+        l2[l2 == 0] = 1
+        return a / np.expand_dims(l2, axis)
+
+    def _load_clip_model(self, model_dir: str, model_name: str = "ViT-L/14"):
+        """Load the CLIP model."""
+        model, preprocess = clip.load(
+            model_name, download_root=model_dir, device=self.device
+        )
+        return model, preprocess
+
+    def _load_sac_model(self, model_path, input_size):
+        """Load the SAC model."""
+        model = self.MLP(input_size)
+        ckpt = torch.load(model_path)
+        model.load_state_dict(ckpt)
+        model.to(self.device)
+        model.eval()
+        return model
+
+    def predict(self, image_path):
+        """Predict the aesthetic score for a given image.
+
+        Args:
+            image_path (str): Path to the image file.
+
+        Returns:
+            float: Predicted aesthetic score.
+        """
+        pil_image = Image.open(image_path)
+        image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            # Extract CLIP features
+            image_features = self.clip_model.encode_image(image)
+            # Normalize features
+            normalized_features = self.normalized(
+                image_features.cpu().detach().numpy()
+            )
+            # Predict score
+            prediction = self.sac_model(
+                torch.from_numpy(normalized_features)
+                .type(torch.FloatTensor)
+                .to(self.device)
+            )
+
+        return prediction.item()
+
+
+if __name__ == "__main__":
+    # Configuration
+    img_path = "/home/users/xinjie.wang/xinjie/asset3d-gen/outputs/imageto3d/demo_objects/bed/sample_0/sample_0_raw.png"  # noqa
+    # clip_model_dir = "/horizon-bucket/robot_lab/users/xinjie.wang/weights/clip"  # noqa
+    # sac_model_path = "/horizon-bucket/robot_lab/users/xinjie.wang/weights/sac/sac+logos+ava1-l14-linearMSE.pth"  # noqa
+
+    # Initialize the predictor
+    predictor = AestheticPredictor()
+
+    # Predict the aesthetic score
+    score = predictor.predict(img_path)
+    print("Aesthetic score predicted by the model:", score)

asset3d_gen/validators/quality_checkers.py

@@ -0,0 +1,195 @@
+import logging
+import os
+
+from tqdm import tqdm
+from asset3d_gen.utils.gpt_clients import GPT_CLIENT, GPTclient
+from asset3d_gen.utils.process_media import render_asset3d
+from asset3d_gen.validators.aesthetic_predictor import AestheticPredictor
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class BaseChecker:
+    def __init__(self, prompt: str = None, verbose: bool = False) -> None:
+        self.prompt = prompt
+        self.verbose = verbose
+
+    def query(self, *args, **kwargs):
+        raise NotImplementedError(
+            "Subclasses must implement the query method."
+        )
+
+    def __call__(self, *args, **kwargs) -> bool:
+        response = self.query(*args, **kwargs)
+        if response is None:
+            response = "Error when calling gpt api."
+
+        if self.verbose and response != "YES":
+            logger.info(response)
+
+        flag = "YES" in response
+        response = "YES" if flag else response
+
+        return flag, response
+
+    @staticmethod
+    def validate(
+        checkers: list["BaseChecker"], images_list: list[list[str]]
+    ) -> list:
+        assert len(checkers) == len(images_list)
+        results = []
+        overall_result = True
+        for checker, images in zip(checkers, images_list):
+            qa_flag, qa_info = checker(images)
+            if isinstance(qa_info, str):
+                qa_info = qa_info.replace("\n", ".")
+            results.append([checker.__class__.__name__, qa_info])
+            if qa_flag is False:
+                overall_result = False
+
+        results.append(["overall", "YES" if overall_result else "NO"])
+
+        return results
+
+
+class MeshGeoChecker(BaseChecker):
+    def __init__(
+        self,
+        gpt_client: GPTclient,
+        prompt: str = None,
+        verbose: bool = False,
+    ) -> None:
+        super().__init__(prompt, verbose)
+        self.gpt_client = gpt_client
+        if self.prompt is None:
+            self.prompt = """
+            Refer to the provided multi-view rendering images to evaluate
+            whether the geometry of the 3D object asset is complete and
+            whether the asset can be placed stably on the ground.
+            Return "YES" only if reach the requirments,
+            otherwise "NO" and explain the reason very briefly.
+            """
+
+    def query(self, image_paths: str) -> str:
+        # Hardcode tmp because of the openrouter can't input multi images.
+        if "openrouter" in self.gpt_client.endpoint:
+            from asset3d_gen.utils.process_media import (
+                combine_images_to_base64,
+            )
+
+            image_paths = combine_images_to_base64(image_paths)
+
+        return self.gpt_client.query(
+            text_prompt=self.prompt,
+            image_base64=image_paths,
+        )
+
+
+class ImageSegChecker(BaseChecker):
+    def __init__(
+        self,
+        gpt_client: GPTclient,
+        prompt: str = None,
+        verbose: bool = False,
+    ) -> None:
+        super().__init__(prompt, verbose)
+        self.gpt_client = gpt_client
+        if self.prompt is None:
+            self.prompt = """
+            The first image is the original, and the second image is the
+            result after segmenting the main object. Evaluate the segmentation
+            quality to ensure the main object is clearly segmented without
+            significant truncation. Note that the foreground of the object
+            needs to be extracted instead of the background.
+            Minor imperfections can be ignored. If segmentation is acceptable,
+            return "YES" only; otherwise, return "NO" with
+            very brief explanation.
+            """
+
+    def query(self, image_paths: list[str]) -> str:
+        if len(image_paths) != 2:
+            raise ValueError(
+                "ImageSegChecker requires exactly two images: [raw_image, seg_image]."  # noqa
+            )
+        # Hardcode tmp because of the openrouter can't input multi images.
+        if "openrouter" in self.gpt_client.endpoint:
+            from asset3d_gen.utils.process_media import (
+                combine_images_to_base64,
+            )
+
+            image_paths = combine_images_to_base64(image_paths)
+
+        return self.gpt_client.query(
+            text_prompt=self.prompt,
+            image_base64=image_paths,
+        )
+
+
+class ImageAestheticChecker(BaseChecker):
+    def __init__(
+        self,
+        clip_model_dir: str = None,
+        sac_model_path: str = None,
+        thresh: float = 4.50,
+        verbose: bool = False,
+    ) -> None:
+        super().__init__(verbose=verbose)
+        self.clip_model_dir = clip_model_dir
+        self.sac_model_path = sac_model_path
+        self.thresh = thresh
+        self.predictor = AestheticPredictor(clip_model_dir, sac_model_path)
+
+    def query(self, image_paths: list[str]) -> float:
+        scores = [self.predictor.predict(img_path) for img_path in image_paths]
+        return sum(scores) / len(scores)
+
+    def __call__(self, image_paths: list[str], **kwargs) -> bool:
+        avg_score = self.query(image_paths)
+        if self.verbose:
+            logger.info(f"Average aesthetic score: {avg_score}")
+        return avg_score > self.thresh, avg_score
+
+
+if __name__ == "__main__":
+    geo_checker = MeshGeoChecker(GPT_CLIENT)
+    seg_checker = ImageSegChecker(GPT_CLIENT)
+    aesthetic_checker = ImageAestheticChecker(
+        "/horizon-bucket/robot_lab/users/xinjie.wang/weights/clip",
+        "/horizon-bucket/robot_lab/users/xinjie.wang/weights/sac/sac+logos+ava1-l14-linearMSE.pth",  # noqa
+    )
+
+    checkers = [geo_checker, seg_checker, aesthetic_checker]
+
+    output_root = "outputs/test_gpt"
+
+    fails = []
+    for idx in tqdm(range(150)):
+        mesh_path = f"outputs/imageto3d/demo_objects/cups/sample_{idx}/sample_{idx}.obj"  # noqa
+        if not os.path.exists(mesh_path):
+            continue
+        image_paths = render_asset3d(
+            mesh_path,
+            f"{output_root}/{idx}",
+            num_images=8,
+            elevation=(30, -30),
+            distance=5.5,
+        )
+
+        for cid, checker in enumerate(checkers):
+            if isinstance(checker, ImageSegChecker):
+                images = [
+                    f"outputs/imageto3d/demo_objects/cups/sample_{idx}/sample_{idx}_raw.png",  # noqa
+                    f"outputs/imageto3d/demo_objects/cups/sample_{idx}/sample_{idx}_cond.png",  # noqa
+                ]
+            else:
+                images = image_paths
+            result, info = checker(images)
+            logger.info(
+                f"Checker {checker.__class__.__name__}: {result}, {info}, mesh {mesh_path}"  # noqa
+            )
+
+            if result is False:
+                fails.append((idx, cid, info))
+
+        break

asset3d_gen/validators/urdf_convertor.py

@@ -0,0 +1,423 @@
+import logging
+import os
+import shutil
+import xml.etree.ElementTree as ET
+import zipfile
+from datetime import datetime
+from xml.dom.minidom import parseString
+
+import numpy as np
+import trimesh
+from asset3d_gen.utils.gpt_clients import GPT_CLIENT, GPTclient
+from asset3d_gen.utils.process_media import render_asset3d
+from asset3d_gen.utils.tags import VERSION
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+__all__ = ["URDFGenerator"]
+
+
+URDF_TEMPLATE = """
+<robot name="template_robot">
+    <link name="template_link">
+        <visual>
+            <geometry>
+                <mesh filename="mesh.obj" scale="1.0 1.0 1.0"/>
+            </geometry>
+        </visual>
+        <collision>
+            <geometry>
+                <mesh filename="mesh.obj" scale="1.0 1.0 1.0"/>
+            </geometry>
+            <gazebo>
+                <mu1>0.8</mu1> <!-- 主摩擦系数 -->
+                <mu2>0.6</mu2> <!-- 次摩擦系数 -->
+            </gazebo>
+        </collision>
+        <inertial>
+            <mass value="1.0"/>
+            <origin xyz="0 0 0"/>
+            <inertia ixx="1.0" ixy="0.0" ixz="0.0" iyy="1.0" iyz="0.0" izz="1.0"/>
+        </inertial>
+        <extra_info>
+            <scale>1.0</scale>
+            <version>"0.0.0"</version>
+            <category>"unknown"</category>
+            <description>"unknown"</description>
+            <min_height>0.0</min_height>
+            <max_height>0.0</max_height>
+            <real_height>0.0</real_height>
+            <min_mass>0.0</min_mass>
+            <max_mass>0.0</max_mass>
+            <generate_time>"-1"</generate_time>
+            <gs_model>""</gs_model>
+        </extra_info>
+    </link>
+</robot>
+"""
+
+
+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
+
+
+class URDFGenerator(object):
+    def __init__(
+        self,
+        gpt_client: GPTclient,
+        mesh_file_list: list[str] = ["material_0.png", "material.mtl"],
+        prompt_template: str = None,
+        attrs_name: list[str] = None,
+        render_dir: str = "urdf_renders",
+        render_view_num: int = 4,
+    ) -> None:
+        if mesh_file_list is None:
+            mesh_file_list = []
+        self.mesh_file_list = mesh_file_list
+        self.output_mesh_dir = "mesh"
+        self.output_render_dir = render_dir
+        self.gpt_client = gpt_client
+        self.render_view_num = render_view_num
+        if render_view_num == 4:
+            view_desc = "This is orthographic projection showing the front, left, right and back views "  # noqa
+        else:
+            view_desc = "This is the rendered views "
+
+        if prompt_template is None:
+            prompt_template = (
+                view_desc
+                + """of the 3D object asset,
+                category: {category}.
+                Give the category of this object asset (within 3 words),
+                (if category is already provided, use it directly),
+                accurately describe this 3D object asset (within 15 words),
+                and give the recommended geometric height range (unit: meter),
+                weight range (unit: kilogram), the average static friction
+                coefficient of the object relative to rubber and the average
+                dynamic friction coefficient of the object relative to rubber.
+                Return response format as shown in Example.
+
+                Example:
+                Category: cup
+                Description: shiny golden cup with floral design
+                Height: 0.1-0.15 m
+                Weight: 0.3-0.6 kg
+                Static friction coefficient: 1.1
+                Dynamic friction coefficient: 0.9
+            """
+            )
+
+        self.prompt_template = prompt_template
+        if attrs_name is None:
+            attrs_name = [
+                "category",
+                "description",
+                "min_height",
+                "max_height",
+                "real_height",
+                "min_mass",
+                "max_mass",
+                "version",
+                "generate_time",
+                "gs_model",
+            ]
+        self.attrs_name = attrs_name
+
+    def parse_response(self, response: str) -> dict[str, any]:
+        lines = response.split("\n")
+        lines = [line.strip() for line in lines if line]
+        category = lines[0].split(": ")[1]
+        description = lines[1].split(": ")[1]
+        min_height, max_height = map(
+            lambda x: float(x.strip().replace(",", "").split()[0]),
+            lines[2].split(": ")[1].split("-"),
+        )
+        min_mass, max_mass = map(
+            lambda x: float(x.strip().replace(",", "").split()[0]),
+            lines[3].split(": ")[1].split("-"),
+        )
+        mu1 = float(lines[4].split(": ")[1].replace(",", ""))
+        mu2 = float(lines[5].split(": ")[1].replace(",", ""))
+
+        return {
+            "category": category.lower(),
+            "description": description.lower(),
+            "min_height": round(min_height, 4),
+            "max_height": round(max_height, 4),
+            "real_height": round((min_height + max_height) / 2, 4),
+            "min_mass": round(min_mass, 4),
+            "max_mass": round(max_mass, 4),
+            "mu1": round(mu1, 2),
+            "mu2": round(mu2, 2),
+            "version": VERSION,
+            "generate_time": datetime.now().strftime("%Y%m%d%H%M%S"),
+        }
+
+    def generate_urdf(
+        self,
+        input_mesh: str,
+        output_dir: str,
+        attr_dict: dict,
+        output_name: str = None,
+    ) -> str:
+        """Generate a URDF file for a given mesh with specified attributes.
+
+        Args:
+            input_mesh (str): Path to the input mesh file.
+            output_dir (str): Directory to store the generated URDF
+                and processed mesh.
+            attr_dict (dict): Dictionary containing attributes like height,
+                mass, and friction coefficients.
+            output_name (str, optional): Name for the generated URDF and robot.
+
+        Returns:
+            str: Path to the generated URDF file.
+        """
+
+        # 1. Load and normalize the mesh
+        mesh = trimesh.load(input_mesh)
+        mesh_scale = np.ptp(mesh.vertices, axis=0).max()
+        mesh.vertices /= mesh_scale  # Normalize to [-0.5, 0.5]
+        raw_height = np.ptp(mesh.vertices, axis=0)[1]
+
+        # 2. Scale the mesh to real height
+        real_height = attr_dict["real_height"]
+        scale = round(real_height / raw_height, 6)
+        mesh = mesh.apply_scale(scale)
+
+        # 3. Prepare output directories and save scaled mesh
+        mesh_folder = os.path.join(output_dir, self.output_mesh_dir)
+        os.makedirs(mesh_folder, exist_ok=True)
+
+        obj_name = os.path.basename(input_mesh)
+        mesh_output_path = os.path.join(mesh_folder, obj_name)
+        mesh.export(mesh_output_path)
+
+        # 4. Copy additional mesh files, if any
+        input_dir = os.path.dirname(input_mesh)
+        for file in self.mesh_file_list:
+            src_file = os.path.join(input_dir, file)
+            dest_file = os.path.join(mesh_folder, file)
+            if os.path.isfile(src_file):
+                shutil.copy(src_file, dest_file)
+
+        # 5. Determine output name
+        if output_name is None:
+            output_name = os.path.splitext(obj_name)[0]
+
+        # 6. Load URDF template and update attributes
+        robot = ET.fromstring(URDF_TEMPLATE)
+        robot.set("name", output_name)
+
+        link = robot.find("link")
+        if link is None:
+            raise ValueError("URDF template is missing 'link' element.")
+        link.set("name", output_name)
+
+        # Update visual geometry
+        visual = link.find("visual/geometry/mesh")
+        if visual is not None:
+            visual.set(
+                "filename", os.path.join(self.output_mesh_dir, obj_name)
+            )
+            visual.set("scale", "1.0 1.0 1.0")
+
+        # Update collision geometry
+        collision = link.find("collision/geometry/mesh")
+        if collision is not None:
+            collision.set(
+                "filename", os.path.join(self.output_mesh_dir, obj_name)
+            )
+            collision.set("scale", "1.0 1.0 1.0")
+
+        # Update friction coefficients
+        gazebo = link.find("collision/gazebo")
+        if gazebo is not None:
+            for param, key in zip(["mu1", "mu2"], ["mu1", "mu2"]):
+                element = gazebo.find(param)
+                if element is not None:
+                    element.text = f"{attr_dict[key]:.2f}"
+
+        # Update mass
+        inertial = link.find("inertial/mass")
+        if inertial is not None:
+            mass_value = (attr_dict["min_mass"] + attr_dict["max_mass"]) / 2
+            inertial.set("value", f"{mass_value:.4f}")
+
+        # Add extra_info element to the link
+        extra_info = link.find("extra_info/scale")
+        if extra_info is not None:
+            extra_info.text = f"{scale:.6f}"
+
+        for key in self.attrs_name:
+            extra_info = link.find(f"extra_info/{key}")
+            if extra_info is not None and key in attr_dict:
+                extra_info.text = f"{attr_dict[key]}"
+
+        # 7. Write URDF to file
+        os.makedirs(output_dir, exist_ok=True)
+        urdf_path = os.path.join(output_dir, f"{output_name}.urdf")
+        tree = ET.ElementTree(robot)
+        tree.write(urdf_path, encoding="utf-8", xml_declaration=True)
+
+        logger.info(f"URDF file saved to {urdf_path}")
+
+        return urdf_path
+
+    @staticmethod
+    def get_attr_from_urdf(
+        urdf_path: str,
+        attr_root: str = ".//link/extra_info",
+        attr_name: str = "scale",
+    ) -> float:
+        if not os.path.exists(urdf_path):
+            raise FileNotFoundError(f"URDF file not found: {urdf_path}")
+
+        mesh_scale = 1.0
+        tree = ET.parse(urdf_path)
+        root = tree.getroot()
+        extra_info = root.find(attr_root)
+        if extra_info is not None:
+            scale_element = extra_info.find(attr_name)
+            if scale_element is not None:
+                mesh_scale = float(scale_element.text)
+
+        return mesh_scale
+
+    @staticmethod
+    def add_quality_tag(
+        urdf_path: str, results, output_path: str = None
+    ) -> None:
+        if output_path is None:
+            output_path = urdf_path
+
+        tree = ET.parse(urdf_path)
+        root = tree.getroot()
+        custom_data = ET.SubElement(root, "custom_data")
+        quality = ET.SubElement(custom_data, "quality")
+        for key, value in results:
+            checker_tag = ET.SubElement(quality, key)
+            checker_tag.text = str(value)
+
+        rough_string = ET.tostring(root, encoding="utf-8")
+        formatted_string = parseString(rough_string).toprettyxml(indent="   ")
+        cleaned_string = "\n".join(
+            [line for line in formatted_string.splitlines() if line.strip()]
+        )
+
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        with open(output_path, "w", encoding="utf-8") as f:
+            f.write(cleaned_string)
+
+        logger.info(f"URDF files saved to {output_path}")
+
+    def get_estimated_attributes(self, asset_attrs: dict):
+        estimated_attrs = {
+            "height": round(
+                (asset_attrs["min_height"] + asset_attrs["max_height"]) / 2, 4
+            ),
+            "mass": round(
+                (asset_attrs["min_mass"] + asset_attrs["max_mass"]) / 2, 4
+            ),
+            "mu": round((asset_attrs["mu1"] + asset_attrs["mu2"]) / 2, 4),
+            "category": asset_attrs["category"],
+        }
+
+        return estimated_attrs
+
+    def __call__(
+        self,
+        mesh_path: str,
+        output_root: str,
+        text_prompt: str = None,
+        category: str = "unknown",
+        **kwargs,
+    ):
+        if text_prompt is None or len(text_prompt) == 0:
+            text_prompt = self.prompt_template
+            text_prompt = text_prompt.format(category=category.lower())
+
+        image_path = render_asset3d(
+            mesh_path,
+            output_root,
+            num_images=self.render_view_num,
+            output_subdir=self.output_render_dir,
+        )
+
+        # Hardcode tmp because of the openrouter can't input multi images.
+        if "openrouter" in self.gpt_client.endpoint:
+            from asset3d_gen.utils.process_media import (
+                combine_images_to_base64,
+            )
+
+            image_path = combine_images_to_base64(image_path)
+
+        response = self.gpt_client.query(text_prompt, image_path)
+        if response is None:
+            asset_attrs = {
+                "category": "unknown",
+                "description": "unknown",
+                "min_height": 1,
+                "max_height": 1,
+                "real_height": 1,
+                "min_mass": 1,
+                "max_mass": 1,
+                "mu1": 0.8,
+                "mu2": 0.6,
+                "version": VERSION,
+                "generate_time": datetime.now().strftime("%Y%m%d%H%M%S"),
+            }
+        else:
+            asset_attrs = self.parse_response(response)
+        for key in self.attrs_name:
+            if key in kwargs:
+                asset_attrs[key] = kwargs[key]
+
+        self.estimated_attrs = self.get_estimated_attributes(asset_attrs)
+
+        urdf_path = self.generate_urdf(mesh_path, output_root, asset_attrs)
+
+        logger.info(f"response: {response}")
+
+        return urdf_path
+
+
+if __name__ == "__main__":
+    urdf_gen = URDFGenerator(GPT_CLIENT, render_view_num=4)
+    urdf_path = urdf_gen(
+        mesh_path="scripts/apps/assets/example_texture/meshes/robot.obj",
+        output_root="outputs/test_urdf",
+        # category="coffee machine",
+        # min_height=1.0,
+        # max_height=1.2,
+        version=VERSION,
+    )
+
+    # zip_files(
+    #     input_paths=[
+    #         "scripts/apps/tmp/2umpdum3e5n/URDF_sample/mesh",
+    #         "scripts/apps/tmp/2umpdum3e5n/URDF_sample/sample.urdf"
+    #     ],
+    #     output_zip="zip.zip"
+    # )

common.py

@@ -0,0 +1,597 @@
+import gc
+import logging
+import os
+import sys
+from glob import glob
+from typing import Union
+
+import cv2
+import gradio as gr
+import numpy as np
+import torch
+import trimesh
+from easydict import EasyDict as edict
+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 tqdm import tqdm
+from asset3d_gen.data.backproject_v2 import entrypoint as backproject_api
+from asset3d_gen.models.delight import DelightingModel
+from asset3d_gen.models.gs_model import GaussianOperator
+from asset3d_gen.models.segment import (
+    RembgRemover,
+    SAMPredictor,
+    trellis_preprocess,
+)
+from asset3d_gen.models.super_resolution import ImageRealESRGAN, ImageStableSR
+from asset3d_gen.scripts.render_gs import entrypoint as render_gs_api
+from asset3d_gen.scripts.text2image import text2img_gen
+from asset3d_gen.utils.process_media import (
+    filter_image_small_connected_components,
+    merge_images_video,
+    render_asset3d,
+)
+from asset3d_gen.utils.tags import VERSION
+from asset3d_gen.validators.quality_checkers import (
+    BaseChecker,
+    ImageAestheticChecker,
+    ImageSegChecker,
+    MeshGeoChecker,
+)
+from asset3d_gen.validators.urdf_convertor import URDFGenerator, zip_files
+
+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.renderers.mesh_renderer import MeshRenderer
+from thirdparty.TRELLIS.trellis.representations import (
+    Gaussian,
+    MeshExtractResult,
+)
+from thirdparty.TRELLIS.trellis.utils import postprocessing_utils
+from thirdparty.TRELLIS.trellis.utils.render_utils import (
+    render_frames,
+    yaw_pitch_r_fov_to_extrinsics_intrinsics,
+)
+import spaces
+
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+MAX_SEED = 100000
+
+
+@spaces.GPU
+def render_mesh(sample, extrinsics, intrinsics, options={}, **kwargs):
+    renderer = MeshRenderer()
+    renderer.rendering_options.resolution = options.get("resolution", 512)
+    renderer.rendering_options.near = options.get("near", 1)
+    renderer.rendering_options.far = options.get("far", 100)
+    renderer.rendering_options.ssaa = options.get("ssaa", 4)
+    rets = {}
+    for extr, intr in tqdm(zip(extrinsics, intrinsics), desc="Rendering"):
+        res = renderer.render(sample, extr, intr)
+        if "normal" not in rets:
+            rets["normal"] = []
+        normal = torch.lerp(
+            torch.zeros_like(res["normal"]), res["normal"], res["mask"]
+        )
+        normal = np.clip(
+            normal.detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255
+        ).astype(np.uint8)
+        rets["normal"].append(normal)
+
+    return rets
+
+
+@spaces.GPU
+def render_video(
+    sample,
+    resolution=512,
+    bg_color=(0, 0, 0),
+    num_frames=300,
+    r=2,
+    fov=40,
+    **kwargs,
+):
+    yaws = torch.linspace(0, 2 * 3.1415, num_frames)
+    yaws = yaws.tolist()
+    pitch = [0.5] * num_frames
+    extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(
+        yaws, pitch, r, fov
+    )
+    render_fn = (
+        render_mesh if isinstance(sample, MeshExtractResult) else render_frames
+    )
+    result = render_fn(
+        sample,
+        extrinsics,
+        intrinsics,
+        {"resolution": resolution, "bg_color": bg_color},
+        **kwargs,
+    )
+
+    return result
+
+
+@spaces.GPU
+def preprocess_image_fn(
+    image: str | np.ndarray | Image.Image,
+    model: DelightingModel | RembgRemover,
+    buffer: dict = None,
+) -> Image.Image:
+    if isinstance(image, str):
+        image = Image.open(image)
+    elif isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+
+    if buffer is not None:
+        buffer["raw_image"] = image
+
+    if isinstance(model, DelightingModel):
+        image = model(image, preprocess=True, target_wh=(512, 512))
+    elif isinstance(model, RembgRemover):
+        image = model(image)
+    image = trellis_preprocess(image)
+
+    return image
+
+
+@spaces.GPU
+def preprocess_sam_image_fn(
+    image: Image.Image, buffer: dict, model: SAMPredictor
+) -> Image.Image:
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+
+    buffer["raw_image"] = image
+    sam_image = model.preprocess_image(image)
+    model.predictor.set_image(sam_image)
+
+    return sam_image
+
+
+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}")
+
+
+@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(),
+        },
+    }
+
+
+@spaces.GPU
+def unpack_state(state: dict) -> tuple[Gaussian, edict, str]:
+    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"],
+    )
+    gs._xyz = torch.tensor(state["gaussian"]["_xyz"], device="cuda")
+    gs._features_dc = torch.tensor(
+        state["gaussian"]["_features_dc"], device="cuda"
+    )
+    gs._scaling = torch.tensor(state["gaussian"]["_scaling"], device="cuda")
+    gs._rotation = torch.tensor(state["gaussian"]["_rotation"], device="cuda")
+    gs._opacity = torch.tensor(state["gaussian"]["_opacity"], device="cuda")
+
+    mesh = edict(
+        vertices=torch.tensor(state["mesh"]["vertices"], device="cuda"),
+        faces=torch.tensor(state["mesh"]["faces"], device="cuda"),
+    )
+
+    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
+
+
+@spaces.GPU
+def select_point(
+    image: np.ndarray,
+    sel_pix: list,
+    point_type: str,
+    model: SAMPredictor,
+    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 = model.generate_masks(image, sel_pix)
+    seg_image = model.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,
+    buffer: dict,
+    pipeline: TrellisImageTo3DPipeline,
+    output_root: str,
+    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)
+        # seg_image.save(f"{TMP_DIR}/seg_image_sam.png")
+    else:
+        seg_image = image
+
+    if isinstance(seg_image, np.ndarray):
+        seg_image = Image.fromarray(seg_image)
+    buffer["seg_image"] = seg_image
+
+    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"]
+    if req is not None:
+        output_root = os.path.join(output_root, str(req.session_hash))
+    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, output_root: str, req: gr.Request
+):
+    user_dir = os.path.join(output_root, str(req.session_hash))
+    gs_model, mesh_model = unpack_state(state)
+
+    mesh = postprocessing_utils.to_glb(
+        gs_model,
+        mesh_model,
+        simplify=0.9,
+        texture_size=1024,
+        verbose=True,
+    )
+    filename = "sample"
+    gs_path = os.path.join(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]]
+    # 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(user_dir, f"{filename}.obj")
+    mesh.export(mesh_obj_path)
+    mesh_glb_path = os.path.join(user_dir, f"{filename}.glb")
+    mesh.export(mesh_glb_path)
+
+    torch.cuda.empty_cache()
+
+    return mesh_glb_path, gs_path, mesh_obj_path, aligned_gs_path
+
+
+@spaces.GPU
+def extract_3d_representations_v2(
+    state: dict,
+    enable_delight: bool,
+    output_root: str,
+    delight_model: DelightingModel,
+    sr_model: Union[ImageRealESRGAN, ImageStableSR],
+    req: gr.Request,
+):
+    user_dir = os.path.join(output_root, str(req.session_hash))
+    gs_model, mesh_model = unpack_state(state)
+
+    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,
+    )
+    color_path = os.path.join(user_dir, "color.png")
+    render_gs_api(aligned_gs_path, color_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(user_dir, f"{filename}.obj")
+    mesh.export(mesh_obj_path)
+
+    mesh = backproject_api(
+        delight_model=delight_model,
+        imagesr_model=sr_model,
+        color_path=color_path,
+        mesh_path=mesh_obj_path,
+        output_path=mesh_obj_path,
+        skip_fix_mesh=False,
+        delight=enable_delight,
+    )
+
+    mesh_glb_path = os.path.join(user_dir, f"{filename}.glb")
+    mesh.export(mesh_glb_path)
+
+    torch.cuda.empty_cache()
+
+    return mesh_glb_path, gs_path, mesh_obj_path, aligned_gs_path
+
+
+@spaces.GPU
+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,
+    output_root: str,
+    urdf_convertor: URDFGenerator,
+    buffer: dict,
+    checkers: list[BaseChecker],
+    req: gr.Request = None,
+):
+    if req is not None:
+        output_root = os.path.join(output_root, str(req.session_hash))
+    # Convert to URDF and recover attrs by gpt4o
+    filename = "sample"
+    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,
+    )
+
+    # 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 = [buffer["raw_image"], buffer["seg_image"]]
+        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",
+    )
+
+    torch.cuda.empty_cache()
+
+    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,
+    output_root: str,
+    guidance_scale: float,
+    model_ip: StableDiffusionXLPipelineIP,
+    model_img: StableDiffusionXLPipeline,
+    bg_model: RembgRemover,
+    infer_step: int = 50,
+    ip_image: Image.Image | str = None,
+    ip_adapt_scale: float = 0.3,
+    image_wh: int | tuple[int, int] = [1024, 1024],
+    n_sample: int = 3,
+    postprocess: bool = True,
+    req: gr.Request = None,
+):
+    if isinstance(image_wh, int):
+        image_wh = (image_wh, image_wh)
+    if req is not None:
+        output_root = os.path.join(output_root, str(req.session_hash))
+        os.makedirs(output_root, exist_ok=True)
+
+    pipeline = model_img if ip_image is None else model_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,
+    )
+    if postprocess:
+        for idx in range(len(images)):
+            image = images[idx]
+            images[idx] = preprocess_image_fn(image, bg_model)
+
+    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

requirements.txt

@@ -1,9 +1,9 @@
---extra-index-url https://download.pytorch.org/whl/cu118
+# --extra-index-url https://download.pytorch.org/whl/cu118
 
-torch==2.1.0
-torchaudio==2.1.0
-torchvision==0.16.0
-xformers==0.0.22.post7
+torch==2.1.0+cu118
+torchaudio==2.1.0+cu118
+torchvision==0.16.0+cu118
+xformers==0.0.22.post7+cu118
 dataclasses_json
 easydict
 opencv-python>4.5
@@ -21,7 +21,7 @@ openai==1.58.1
 spconv-cu118==2.3.8
 transformers==4.42.4
 gradio_litmodel3d==0.0.1
-# gradio==5.12.0
+gradio==5.12.0
 sentencepiece==0.2.0
 diffusers==0.31.0
 xatlas==0.0.9
@@ -33,6 +33,7 @@ basicsr==1.4.2
 realesrgan==0.3.0
 pydantic==2.9.2
 vtk==9.3.1
+spaces
 utils3d@git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8
 clip@git+https://github.com/openai/CLIP.git
 kolors@git+https://github.com/Kwai-Kolors/Kolors.git#egg=038818d

thirdparty/TRELLIS/trellis/trellis/__init__.py

@@ -0,0 +1,6 @@
+from . import models
+from . import modules
+from . import pipelines
+from . import renderers
+from . import representations
+from . import utils

thirdparty/TRELLIS/trellis/trellis/models/__init__.py

@@ -0,0 +1,70 @@
+import importlib
+
+__attributes = {
+    'SparseStructureEncoder': 'sparse_structure_vae',
+    'SparseStructureDecoder': 'sparse_structure_vae',
+    'SparseStructureFlowModel': 'sparse_structure_flow',
+    'SLatEncoder': 'structured_latent_vae',
+    'SLatGaussianDecoder': 'structured_latent_vae',
+    'SLatRadianceFieldDecoder': 'structured_latent_vae',
+    'SLatMeshDecoder': 'structured_latent_vae',
+    'SLatFlowModel': 'structured_latent_flow',
+}
+
+__submodules = []
+
+__all__ = list(__attributes.keys()) + __submodules
+
+def __getattr__(name):
+    if name not in globals():
+        if name in __attributes:
+            module_name = __attributes[name]
+            module = importlib.import_module(f".{module_name}", __name__)
+            globals()[name] = getattr(module, name)
+        elif name in __submodules:
+            module = importlib.import_module(f".{name}", __name__)
+            globals()[name] = module
+        else:
+            raise AttributeError(f"module {__name__} has no attribute {name}")
+    return globals()[name]
+
+
+def from_pretrained(path: str, **kwargs):
+    """
+    Load a model from a pretrained checkpoint.
+
+    Args:
+        path: The path to the checkpoint. Can be either local path or a Hugging Face model name.
+              NOTE: config file and model file should take the name f'{path}.json' and f'{path}.safetensors' respectively.
+        **kwargs: Additional arguments for the model constructor.
+    """
+    import os
+    import json
+    from safetensors.torch import load_file
+    is_local = os.path.exists(f"{path}.json") and os.path.exists(f"{path}.safetensors")
+
+    if is_local:
+        config_file = f"{path}.json"
+        model_file = f"{path}.safetensors"
+    else:
+        from huggingface_hub import hf_hub_download
+        path_parts = path.split('/')
+        repo_id = f'{path_parts[0]}/{path_parts[1]}'
+        model_name = '/'.join(path_parts[2:])
+        config_file = hf_hub_download(repo_id, f"{model_name}.json")
+        model_file = hf_hub_download(repo_id, f"{model_name}.safetensors")
+
+    with open(config_file, 'r') as f:
+        config = json.load(f)
+    model = __getattr__(config['name'])(**config['args'], **kwargs)
+    model.load_state_dict(load_file(model_file))
+
+    return model
+
+
+# For Pylance
+if __name__ == '__main__':
+    from .sparse_structure_vae import SparseStructureEncoder, SparseStructureDecoder
+    from .sparse_structure_flow import SparseStructureFlowModel
+    from .structured_latent_vae import SLatEncoder, SLatGaussianDecoder, SLatRadianceFieldDecoder, SLatMeshDecoder
+    from .structured_latent_flow import SLatFlowModel

thirdparty/TRELLIS/trellis/trellis/models/sparse_structure_flow.py

@@ -0,0 +1,200 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from ..modules.utils import convert_module_to_f16, convert_module_to_f32
+from ..modules.transformer import AbsolutePositionEmbedder, ModulatedTransformerCrossBlock
+from ..modules.spatial import patchify, unpatchify
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+
+        Args:
+            t: a 1-D Tensor of N indices, one per batch element.
+                These may be fractional.
+            dim: the dimension of the output.
+            max_period: controls the minimum frequency of the embeddings.
+
+        Returns:
+            an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(
+            -np.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class SparseStructureFlowModel(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        model_channels: int,
+        cond_channels: int,
+        out_channels: int,
+        num_blocks: int,
+        num_heads: Optional[int] = None,
+        num_head_channels: Optional[int] = 64,
+        mlp_ratio: float = 4,
+        patch_size: int = 2,
+        pe_mode: Literal["ape", "rope"] = "ape",
+        use_fp16: bool = False,
+        use_checkpoint: bool = False,
+        share_mod: bool = False,
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+    ):
+        super().__init__()
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.cond_channels = cond_channels
+        self.out_channels = out_channels
+        self.num_blocks = num_blocks
+        self.num_heads = num_heads or model_channels // num_head_channels
+        self.mlp_ratio = mlp_ratio
+        self.patch_size = patch_size
+        self.pe_mode = pe_mode
+        self.use_fp16 = use_fp16
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.qk_rms_norm = qk_rms_norm
+        self.qk_rms_norm_cross = qk_rms_norm_cross
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+
+        self.t_embedder = TimestepEmbedder(model_channels)
+        if share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(model_channels, 6 * model_channels, bias=True)
+            )
+
+        if pe_mode == "ape":
+            pos_embedder = AbsolutePositionEmbedder(model_channels, 3)
+            coords = torch.meshgrid(*[torch.arange(res, device=self.device) for res in [resolution // patch_size] * 3], indexing='ij')
+            coords = torch.stack(coords, dim=-1).reshape(-1, 3)
+            pos_emb = pos_embedder(coords)
+            self.register_buffer("pos_emb", pos_emb)
+
+        self.input_layer = nn.Linear(in_channels * patch_size**3, model_channels)
+            
+        self.blocks = nn.ModuleList([
+            ModulatedTransformerCrossBlock(
+                model_channels,
+                cond_channels,
+                num_heads=self.num_heads,
+                mlp_ratio=self.mlp_ratio,
+                attn_mode='full',
+                use_checkpoint=self.use_checkpoint,
+                use_rope=(pe_mode == "rope"),
+                share_mod=share_mod,
+                qk_rms_norm=self.qk_rms_norm,
+                qk_rms_norm_cross=self.qk_rms_norm_cross,
+            )
+            for _ in range(num_blocks)
+        ])
+
+        self.out_layer = nn.Linear(model_channels, out_channels * patch_size**3)
+
+        self.initialize_weights()
+        if use_fp16:
+            self.convert_to_fp16()
+
+    @property
+    def device(self) -> torch.device:
+        """
+        Return the device of the model.
+        """
+        return next(self.parameters()).device
+
+    def convert_to_fp16(self) -> None:
+        """
+        Convert the torso of the model to float16.
+        """
+        self.blocks.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self) -> None:
+        """
+        Convert the torso of the model to float32.
+        """
+        self.blocks.apply(convert_module_to_f32)
+
+    def initialize_weights(self) -> None:
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers in DiT blocks:
+        if self.share_mod:
+            nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+        else:
+            for block in self.blocks:
+                nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+                nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.out_layer.weight, 0)
+        nn.init.constant_(self.out_layer.bias, 0)
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        assert [*x.shape] == [x.shape[0], self.in_channels, *[self.resolution] * 3], \
+                f"Input shape mismatch, got {x.shape}, expected {[x.shape[0], self.in_channels, *[self.resolution] * 3]}"
+
+        h = patchify(x, self.patch_size)
+        h = h.view(*h.shape[:2], -1).permute(0, 2, 1).contiguous()
+
+        h = self.input_layer(h)
+        h = h + self.pos_emb[None]
+        t_emb = self.t_embedder(t)
+        if self.share_mod:
+            t_emb = self.adaLN_modulation(t_emb)
+        t_emb = t_emb.type(self.dtype)
+        h = h.type(self.dtype)
+        cond = cond.type(self.dtype)
+        for block in self.blocks:
+            h = block(h, t_emb, cond)
+        h = h.type(x.dtype)
+        h = F.layer_norm(h, h.shape[-1:])
+        h = self.out_layer(h)
+
+        h = h.permute(0, 2, 1).view(h.shape[0], h.shape[2], *[self.resolution // self.patch_size] * 3)
+        h = unpatchify(h, self.patch_size).contiguous()
+
+        return h

thirdparty/TRELLIS/trellis/trellis/models/sparse_structure_vae.py

@@ -0,0 +1,306 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ..modules.norm import GroupNorm32, ChannelLayerNorm32
+from ..modules.spatial import pixel_shuffle_3d
+from ..modules.utils import zero_module, convert_module_to_f16, convert_module_to_f32
+
+
+def norm_layer(norm_type: str, *args, **kwargs) -> nn.Module:
+    """
+    Return a normalization layer.
+    """
+    if norm_type == "group":
+        return GroupNorm32(32, *args, **kwargs)
+    elif norm_type == "layer":
+        return ChannelLayerNorm32(*args, **kwargs)
+    else:
+        raise ValueError(f"Invalid norm type {norm_type}")
+
+
+class ResBlock3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        norm_type: Literal["group", "layer"] = "layer",
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+
+        self.norm1 = norm_layer(norm_type, channels)
+        self.norm2 = norm_layer(norm_type, self.out_channels)
+        self.conv1 = nn.Conv3d(channels, self.out_channels, 3, padding=1)
+        self.conv2 = zero_module(nn.Conv3d(self.out_channels, self.out_channels, 3, padding=1))
+        self.skip_connection = nn.Conv3d(channels, self.out_channels, 1) if channels != self.out_channels else nn.Identity()
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.norm1(x)
+        h = F.silu(h)
+        h = self.conv1(h)
+        h = self.norm2(h)
+        h = F.silu(h)
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        return h
+
+
+class DownsampleBlock3d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        mode: Literal["conv", "avgpool"] = "conv",
+    ):
+        assert mode in ["conv", "avgpool"], f"Invalid mode {mode}"
+
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        if mode == "conv":
+            self.conv = nn.Conv3d(in_channels, out_channels, 2, stride=2)
+        elif mode == "avgpool":
+            assert in_channels == out_channels, "Pooling mode requires in_channels to be equal to out_channels"
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, "conv"):
+            return self.conv(x)
+        else:
+            return F.avg_pool3d(x, 2)
+
+
+class UpsampleBlock3d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        mode: Literal["conv", "nearest"] = "conv",
+    ):
+        assert mode in ["conv", "nearest"], f"Invalid mode {mode}"
+
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        if mode == "conv":
+            self.conv = nn.Conv3d(in_channels, out_channels*8, 3, padding=1)
+        elif mode == "nearest":
+            assert in_channels == out_channels, "Nearest mode requires in_channels to be equal to out_channels"
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, "conv"):
+            x = self.conv(x)
+            return pixel_shuffle_3d(x, 2)
+        else:
+            return F.interpolate(x, scale_factor=2, mode="nearest")
+        
+
+class SparseStructureEncoder(nn.Module):
+    """
+    Encoder for Sparse Structure (\mathcal{E}_S in the paper Sec. 3.3).
+    
+    Args:
+        in_channels (int): Channels of the input.
+        latent_channels (int): Channels of the latent representation.
+        num_res_blocks (int): Number of residual blocks at each resolution.
+        channels (List[int]): Channels of the encoder blocks.
+        num_res_blocks_middle (int): Number of residual blocks in the middle.
+        norm_type (Literal["group", "layer"]): Type of normalization layer.
+        use_fp16 (bool): Whether to use FP16.
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        latent_channels: int,
+        num_res_blocks: int,
+        channels: List[int],
+        num_res_blocks_middle: int = 2,
+        norm_type: Literal["group", "layer"] = "layer",
+        use_fp16: bool = False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.latent_channels = latent_channels
+        self.num_res_blocks = num_res_blocks
+        self.channels = channels
+        self.num_res_blocks_middle = num_res_blocks_middle
+        self.norm_type = norm_type
+        self.use_fp16 = use_fp16
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+
+        self.input_layer = nn.Conv3d(in_channels, channels[0], 3, padding=1)
+
+        self.blocks = nn.ModuleList([])
+        for i, ch in enumerate(channels):
+            self.blocks.extend([
+                ResBlock3d(ch, ch)
+                for _ in range(num_res_blocks)
+            ])
+            if i < len(channels) - 1:
+                self.blocks.append(
+                    DownsampleBlock3d(ch, channels[i+1])
+                )
+        
+        self.middle_block = nn.Sequential(*[
+            ResBlock3d(channels[-1], channels[-1])
+            for _ in range(num_res_blocks_middle)
+        ])
+
+        self.out_layer = nn.Sequential(
+            norm_layer(norm_type, channels[-1]),
+            nn.SiLU(),
+            nn.Conv3d(channels[-1], latent_channels*2, 3, padding=1)
+        )
+
+        if use_fp16:
+            self.convert_to_fp16()
+
+    @property
+    def device(self) -> torch.device:
+        """
+        Return the device of the model.
+        """
+        return next(self.parameters()).device
+
+    def convert_to_fp16(self) -> None:
+        """
+        Convert the torso of the model to float16.
+        """
+        self.use_fp16 = True
+        self.dtype = torch.float16
+        self.blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self) -> None:
+        """
+        Convert the torso of the model to float32.
+        """
+        self.use_fp16 = False
+        self.dtype = torch.float32
+        self.blocks.apply(convert_module_to_f32)
+        self.middle_block.apply(convert_module_to_f32)
+
+    def forward(self, x: torch.Tensor, sample_posterior: bool = False, return_raw: bool = False) -> torch.Tensor:
+        h = self.input_layer(x)
+        h = h.type(self.dtype)
+
+        for block in self.blocks:
+            h = block(h)
+        h = self.middle_block(h)
+
+        h = h.type(x.dtype)
+        h = self.out_layer(h)
+
+        mean, logvar = h.chunk(2, dim=1)
+
+        if sample_posterior:
+            std = torch.exp(0.5 * logvar)
+            z = mean + std * torch.randn_like(std)
+        else:
+            z = mean
+            
+        if return_raw:
+            return z, mean, logvar
+        return z
+        
+
+class SparseStructureDecoder(nn.Module):
+    """
+    Decoder for Sparse Structure (\mathcal{D}_S in the paper Sec. 3.3).
+    
+    Args:
+        out_channels (int): Channels of the output.
+        latent_channels (int): Channels of the latent representation.
+        num_res_blocks (int): Number of residual blocks at each resolution.
+        channels (List[int]): Channels of the decoder blocks.
+        num_res_blocks_middle (int): Number of residual blocks in the middle.
+        norm_type (Literal["group", "layer"]): Type of normalization layer.
+        use_fp16 (bool): Whether to use FP16.
+    """ 
+    def __init__(
+        self,
+        out_channels: int,
+        latent_channels: int,
+        num_res_blocks: int,
+        channels: List[int],
+        num_res_blocks_middle: int = 2,
+        norm_type: Literal["group", "layer"] = "layer",
+        use_fp16: bool = False,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.latent_channels = latent_channels
+        self.num_res_blocks = num_res_blocks
+        self.channels = channels
+        self.num_res_blocks_middle = num_res_blocks_middle
+        self.norm_type = norm_type
+        self.use_fp16 = use_fp16
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+
+        self.input_layer = nn.Conv3d(latent_channels, channels[0], 3, padding=1)
+
+        self.middle_block = nn.Sequential(*[
+            ResBlock3d(channels[0], channels[0])
+            for _ in range(num_res_blocks_middle)
+        ])
+
+        self.blocks = nn.ModuleList([])
+        for i, ch in enumerate(channels):
+            self.blocks.extend([
+                ResBlock3d(ch, ch)
+                for _ in range(num_res_blocks)
+            ])
+            if i < len(channels) - 1:
+                self.blocks.append(
+                    UpsampleBlock3d(ch, channels[i+1])
+                )
+
+        self.out_layer = nn.Sequential(
+            norm_layer(norm_type, channels[-1]),
+            nn.SiLU(),
+            nn.Conv3d(channels[-1], out_channels, 3, padding=1)
+        )
+
+        if use_fp16:
+            self.convert_to_fp16()
+
+    @property
+    def device(self) -> torch.device:
+        """
+        Return the device of the model.
+        """
+        return next(self.parameters()).device
+    
+    def convert_to_fp16(self) -> None:
+        """
+        Convert the torso of the model to float16.
+        """
+        self.use_fp16 = True
+        self.dtype = torch.float16
+        self.blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self) -> None:
+        """
+        Convert the torso of the model to float32.
+        """
+        self.use_fp16 = False
+        self.dtype = torch.float32
+        self.blocks.apply(convert_module_to_f32)
+        self.middle_block.apply(convert_module_to_f32)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.input_layer(x)
+        
+        h = h.type(self.dtype)
+                
+        h = self.middle_block(h)
+        for block in self.blocks:
+            h = block(h)
+
+        h = h.type(x.dtype)
+        h = self.out_layer(h)
+        return h

thirdparty/TRELLIS/trellis/trellis/models/structured_latent_flow.py

@@ -0,0 +1,262 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from ..modules.utils import zero_module, convert_module_to_f16, convert_module_to_f32
+from ..modules.transformer import AbsolutePositionEmbedder
+from ..modules.norm import LayerNorm32
+from ..modules import sparse as sp
+from ..modules.sparse.transformer import ModulatedSparseTransformerCrossBlock
+from .sparse_structure_flow import TimestepEmbedder
+
+
+class SparseResBlock3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        emb_channels: int,
+        out_channels: Optional[int] = None,
+        downsample: bool = False,
+        upsample: bool = False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.out_channels = out_channels or channels
+        self.downsample = downsample
+        self.upsample = upsample
+        
+        assert not (downsample and upsample), "Cannot downsample and upsample at the same time"
+
+        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
+        self.conv1 = sp.SparseConv3d(channels, self.out_channels, 3)
+        self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(emb_channels, 2 * self.out_channels, bias=True),
+        )
+        self.skip_connection = sp.SparseLinear(channels, self.out_channels) if channels != self.out_channels else nn.Identity()
+        self.updown = None
+        if self.downsample:
+            self.updown = sp.SparseDownsample(2)
+        elif self.upsample:
+            self.updown = sp.SparseUpsample(2)
+
+    def _updown(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        if self.updown is not None:
+            x = self.updown(x)
+        return x
+
+    def forward(self, x: sp.SparseTensor, emb: torch.Tensor) -> sp.SparseTensor:
+        emb_out = self.emb_layers(emb).type(x.dtype)
+        scale, shift = torch.chunk(emb_out, 2, dim=1)
+
+        x = self._updown(x)
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv1(h)
+        h = h.replace(self.norm2(h.feats)) * (1 + scale) + shift
+        h = h.replace(F.silu(h.feats))
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+
+        return h
+    
+
+class SLatFlowModel(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        model_channels: int,
+        cond_channels: int,
+        out_channels: int,
+        num_blocks: int,
+        num_heads: Optional[int] = None,
+        num_head_channels: Optional[int] = 64,
+        mlp_ratio: float = 4,
+        patch_size: int = 2,
+        num_io_res_blocks: int = 2,
+        io_block_channels: List[int] = None,
+        pe_mode: Literal["ape", "rope"] = "ape",
+        use_fp16: bool = False,
+        use_checkpoint: bool = False,
+        use_skip_connection: bool = True,
+        share_mod: bool = False,
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+    ):
+        super().__init__()
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.cond_channels = cond_channels
+        self.out_channels = out_channels
+        self.num_blocks = num_blocks
+        self.num_heads = num_heads or model_channels // num_head_channels
+        self.mlp_ratio = mlp_ratio
+        self.patch_size = patch_size
+        self.num_io_res_blocks = num_io_res_blocks
+        self.io_block_channels = io_block_channels
+        self.pe_mode = pe_mode
+        self.use_fp16 = use_fp16
+        self.use_checkpoint = use_checkpoint
+        self.use_skip_connection = use_skip_connection
+        self.share_mod = share_mod
+        self.qk_rms_norm = qk_rms_norm
+        self.qk_rms_norm_cross = qk_rms_norm_cross
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+
+        assert int(np.log2(patch_size)) == np.log2(patch_size), "Patch size must be a power of 2"
+        assert np.log2(patch_size) == len(io_block_channels), "Number of IO ResBlocks must match the number of stages"
+
+        self.t_embedder = TimestepEmbedder(model_channels)
+        if share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(model_channels, 6 * model_channels, bias=True)
+            )
+
+        if pe_mode == "ape":
+            self.pos_embedder = AbsolutePositionEmbedder(model_channels)
+
+        self.input_layer = sp.SparseLinear(in_channels, io_block_channels[0])
+        self.input_blocks = nn.ModuleList([])
+        for chs, next_chs in zip(io_block_channels, io_block_channels[1:] + [model_channels]):
+            self.input_blocks.extend([
+                SparseResBlock3d(
+                    chs,
+                    model_channels,
+                    out_channels=chs,
+                )
+                for _ in range(num_io_res_blocks-1)
+            ])
+            self.input_blocks.append(
+                SparseResBlock3d(
+                    chs,
+                    model_channels,
+                    out_channels=next_chs,
+                    downsample=True,
+                )
+            )
+            
+        self.blocks = nn.ModuleList([
+            ModulatedSparseTransformerCrossBlock(
+                model_channels,
+                cond_channels,
+                num_heads=self.num_heads,
+                mlp_ratio=self.mlp_ratio,
+                attn_mode='full',
+                use_checkpoint=self.use_checkpoint,
+                use_rope=(pe_mode == "rope"),
+                share_mod=self.share_mod,
+                qk_rms_norm=self.qk_rms_norm,
+                qk_rms_norm_cross=self.qk_rms_norm_cross,
+            )
+            for _ in range(num_blocks)
+        ])
+
+        self.out_blocks = nn.ModuleList([])
+        for chs, prev_chs in zip(reversed(io_block_channels), [model_channels] + list(reversed(io_block_channels[1:]))):
+            self.out_blocks.append(
+                SparseResBlock3d(
+                    prev_chs * 2 if self.use_skip_connection else prev_chs,
+                    model_channels,
+                    out_channels=chs,
+                    upsample=True,
+                )
+            )
+            self.out_blocks.extend([
+                SparseResBlock3d(
+                    chs * 2 if self.use_skip_connection else chs,
+                    model_channels,
+                    out_channels=chs,
+                )
+                for _ in range(num_io_res_blocks-1)
+            ])
+        self.out_layer = sp.SparseLinear(io_block_channels[0], out_channels)
+
+        self.initialize_weights()
+        if use_fp16:
+            self.convert_to_fp16()
+
+    @property
+    def device(self) -> torch.device:
+        """
+        Return the device of the model.
+        """
+        return next(self.parameters()).device
+
+    def convert_to_fp16(self) -> None:
+        """
+        Convert the torso of the model to float16.
+        """
+        self.input_blocks.apply(convert_module_to_f16)
+        self.blocks.apply(convert_module_to_f16)
+        self.out_blocks.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self) -> None:
+        """
+        Convert the torso of the model to float32.
+        """
+        self.input_blocks.apply(convert_module_to_f32)
+        self.blocks.apply(convert_module_to_f32)
+        self.out_blocks.apply(convert_module_to_f32)
+
+    def initialize_weights(self) -> None:
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers in DiT blocks:
+        if self.share_mod:
+            nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+        else:
+            for block in self.blocks:
+                nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+                nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.out_layer.weight, 0)
+        nn.init.constant_(self.out_layer.bias, 0)
+
+    def forward(self, x: sp.SparseTensor, t: torch.Tensor, cond: torch.Tensor) -> sp.SparseTensor:
+        h = self.input_layer(x).type(self.dtype)
+        t_emb = self.t_embedder(t)
+        if self.share_mod:
+            t_emb = self.adaLN_modulation(t_emb)
+        t_emb = t_emb.type(self.dtype)
+        cond = cond.type(self.dtype)
+
+        skips = []
+        # pack with input blocks
+        for block in self.input_blocks:
+            h = block(h, t_emb)
+            skips.append(h.feats)
+        
+        if self.pe_mode == "ape":
+            h = h + self.pos_embedder(h.coords[:, 1:]).type(self.dtype)
+        for block in self.blocks:
+            h = block(h, t_emb, cond)
+
+        # unpack with output blocks
+        for block, skip in zip(self.out_blocks, reversed(skips)):
+            if self.use_skip_connection:
+                h = block(h.replace(torch.cat([h.feats, skip], dim=1)), t_emb)
+            else:
+                h = block(h, t_emb)
+
+        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+        h = self.out_layer(h.type(x.dtype))
+        return h

thirdparty/TRELLIS/trellis/trellis/models/structured_latent_vae/__init__.py

@@ -0,0 +1,4 @@
+from .encoder import SLatEncoder
+from .decoder_gs import SLatGaussianDecoder
+from .decoder_rf import SLatRadianceFieldDecoder
+from .decoder_mesh import SLatMeshDecoder

thirdparty/TRELLIS/trellis/trellis/models/structured_latent_vae/base.py

@@ -0,0 +1,117 @@
+from typing import *
+import torch
+import torch.nn as nn
+from ...modules.utils import convert_module_to_f16, convert_module_to_f32
+from ...modules import sparse as sp
+from ...modules.transformer import AbsolutePositionEmbedder
+from ...modules.sparse.transformer import SparseTransformerBlock
+
+
+def block_attn_config(self):
+    """
+    Return the attention configuration of the model.
+    """
+    for i in range(self.num_blocks):
+        if self.attn_mode == "shift_window":
+            yield "serialized", self.window_size, 0, (16 * (i % 2),) * 3, sp.SerializeMode.Z_ORDER
+        elif self.attn_mode == "shift_sequence":
+            yield "serialized", self.window_size, self.window_size // 2 * (i % 2), (0, 0, 0), sp.SerializeMode.Z_ORDER
+        elif self.attn_mode == "shift_order":
+            yield "serialized", self.window_size, 0, (0, 0, 0), sp.SerializeModes[i % 4]
+        elif self.attn_mode == "full":
+            yield "full", None, None, None, None
+        elif self.attn_mode == "swin":
+            yield "windowed", self.window_size, None, self.window_size // 2 * (i % 2), None
+
+
+class SparseTransformerBase(nn.Module):
+    """
+    Sparse Transformer without output layers.
+    Serve as the base class for encoder and decoder.
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        model_channels: int,
+        num_blocks: int,
+        num_heads: Optional[int] = None,
+        num_head_channels: Optional[int] = 64,
+        mlp_ratio: float = 4.0,
+        attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "full",
+        window_size: Optional[int] = None,
+        pe_mode: Literal["ape", "rope"] = "ape",
+        use_fp16: bool = False,
+        use_checkpoint: bool = False,
+        qk_rms_norm: bool = False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.num_blocks = num_blocks
+        self.window_size = window_size
+        self.num_heads = num_heads or model_channels // num_head_channels
+        self.mlp_ratio = mlp_ratio
+        self.attn_mode = attn_mode
+        self.pe_mode = pe_mode
+        self.use_fp16 = use_fp16
+        self.use_checkpoint = use_checkpoint
+        self.qk_rms_norm = qk_rms_norm
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+
+        if pe_mode == "ape":
+            self.pos_embedder = AbsolutePositionEmbedder(model_channels)
+
+        self.input_layer = sp.SparseLinear(in_channels, model_channels)
+        self.blocks = nn.ModuleList([
+            SparseTransformerBlock(
+                model_channels,
+                num_heads=self.num_heads,
+                mlp_ratio=self.mlp_ratio,
+                attn_mode=attn_mode,
+                window_size=window_size,
+                shift_sequence=shift_sequence,
+                shift_window=shift_window,
+                serialize_mode=serialize_mode,
+                use_checkpoint=self.use_checkpoint,
+                use_rope=(pe_mode == "rope"),
+                qk_rms_norm=self.qk_rms_norm,
+            )
+            for attn_mode, window_size, shift_sequence, shift_window, serialize_mode in block_attn_config(self)
+        ])
+
+    @property
+    def device(self) -> torch.device:
+        """
+        Return the device of the model.
+        """
+        return next(self.parameters()).device
+
+    def convert_to_fp16(self) -> None:
+        """
+        Convert the torso of the model to float16.
+        """
+        self.blocks.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self) -> None:
+        """
+        Convert the torso of the model to float32.
+        """
+        self.blocks.apply(convert_module_to_f32)
+
+    def initialize_weights(self) -> None:
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        h = self.input_layer(x)
+        if self.pe_mode == "ape":
+            h = h + self.pos_embedder(x.coords[:, 1:])
+        h = h.type(self.dtype)
+        for block in self.blocks:
+            h = block(h)
+        return h

thirdparty/TRELLIS/trellis/trellis/models/structured_latent_vae/decoder_gs.py

@@ -0,0 +1,122 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ...modules import sparse as sp
+from ...utils.random_utils import hammersley_sequence
+from .base import SparseTransformerBase
+from ...representations import Gaussian
+
+
+class SLatGaussianDecoder(SparseTransformerBase):
+    def __init__(
+        self,
+        resolution: int,
+        model_channels: int,
+        latent_channels: int,
+        num_blocks: int,
+        num_heads: Optional[int] = None,
+        num_head_channels: Optional[int] = 64,
+        mlp_ratio: float = 4,
+        attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "swin",
+        window_size: int = 8,
+        pe_mode: Literal["ape", "rope"] = "ape",
+        use_fp16: bool = False,
+        use_checkpoint: bool = False,
+        qk_rms_norm: bool = False,
+        representation_config: dict = None,
+    ):
+        super().__init__(
+            in_channels=latent_channels,
+            model_channels=model_channels,
+            num_blocks=num_blocks,
+            num_heads=num_heads,
+            num_head_channels=num_head_channels,
+            mlp_ratio=mlp_ratio,
+            attn_mode=attn_mode,
+            window_size=window_size,
+            pe_mode=pe_mode,
+            use_fp16=use_fp16,
+            use_checkpoint=use_checkpoint,
+            qk_rms_norm=qk_rms_norm,
+        )
+        self.resolution = resolution
+        self.rep_config = representation_config
+        self._calc_layout()
+        self.out_layer = sp.SparseLinear(model_channels, self.out_channels)
+        self._build_perturbation()
+
+        self.initialize_weights()
+        if use_fp16:
+            self.convert_to_fp16()
+
+    def initialize_weights(self) -> None:
+        super().initialize_weights()
+        # Zero-out output layers:
+        nn.init.constant_(self.out_layer.weight, 0)
+        nn.init.constant_(self.out_layer.bias, 0)
+
+    def _build_perturbation(self) -> None:
+        perturbation = [hammersley_sequence(3, i, self.rep_config['num_gaussians']) for i in range(self.rep_config['num_gaussians'])]
+        perturbation = torch.tensor(perturbation).float() * 2 - 1
+        perturbation = perturbation / self.rep_config['voxel_size']
+        perturbation = torch.atanh(perturbation).to(self.device)
+        self.register_buffer('offset_perturbation', perturbation)
+
+    def _calc_layout(self) -> None:
+        self.layout = {
+            '_xyz' : {'shape': (self.rep_config['num_gaussians'], 3), 'size': self.rep_config['num_gaussians'] * 3},
+            '_features_dc' : {'shape': (self.rep_config['num_gaussians'], 1, 3), 'size': self.rep_config['num_gaussians'] * 3},
+            '_scaling' : {'shape': (self.rep_config['num_gaussians'], 3), 'size': self.rep_config['num_gaussians'] * 3},
+            '_rotation' : {'shape': (self.rep_config['num_gaussians'], 4), 'size': self.rep_config['num_gaussians'] * 4},
+            '_opacity' : {'shape': (self.rep_config['num_gaussians'], 1), 'size': self.rep_config['num_gaussians']},
+        }
+        start = 0
+        for k, v in self.layout.items():
+            v['range'] = (start, start + v['size'])
+            start += v['size']
+        self.out_channels = start
+    
+    def to_representation(self, x: sp.SparseTensor) -> List[Gaussian]:
+        """
+        Convert a batch of network outputs to 3D representations.
+
+        Args:
+            x: The [N x * x C] sparse tensor output by the network.
+
+        Returns:
+            list of representations
+        """
+        ret = []
+        for i in range(x.shape[0]):
+            representation = Gaussian(
+                sh_degree=0,
+                aabb=[-0.5, -0.5, -0.5, 1.0, 1.0, 1.0],
+                mininum_kernel_size = self.rep_config['3d_filter_kernel_size'],
+                scaling_bias = self.rep_config['scaling_bias'],
+                opacity_bias = self.rep_config['opacity_bias'],
+                scaling_activation = self.rep_config['scaling_activation']
+            )
+            xyz = (x.coords[x.layout[i]][:, 1:].float() + 0.5) / self.resolution
+            for k, v in self.layout.items():
+                if k == '_xyz':
+                    offset = x.feats[x.layout[i]][:, v['range'][0]:v['range'][1]].reshape(-1, *v['shape'])
+                    offset = offset * self.rep_config['lr'][k]
+                    if self.rep_config['perturb_offset']:
+                        offset = offset + self.offset_perturbation
+                    offset = torch.tanh(offset) / self.resolution * 0.5 * self.rep_config['voxel_size']
+                    _xyz = xyz.unsqueeze(1) + offset
+                    setattr(representation, k, _xyz.flatten(0, 1))
+                else:
+                    feats = x.feats[x.layout[i]][:, v['range'][0]:v['range'][1]].reshape(-1, *v['shape']).flatten(0, 1)
+                    feats = feats * self.rep_config['lr'][k]
+                    setattr(representation, k, feats)
+            ret.append(representation)
+        return ret
+
+    def forward(self, x: sp.SparseTensor) -> List[Gaussian]:
+        h = super().forward(x)
+        h = h.type(x.dtype)
+        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+        h = self.out_layer(h)
+        return self.to_representation(h)

thirdparty/TRELLIS/trellis/trellis/models/structured_latent_vae/decoder_mesh.py

@@ -0,0 +1,167 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from ...modules.utils import zero_module, convert_module_to_f16, convert_module_to_f32
+from ...modules import sparse as sp
+from .base import SparseTransformerBase
+from ...representations import MeshExtractResult
+from ...representations.mesh import SparseFeatures2Mesh
+
+
+class SparseSubdivideBlock3d(nn.Module):
+    """
+    A 3D subdivide block that can subdivide the sparse tensor.
+
+    Args:
+        channels: channels in the inputs and outputs.
+        out_channels: if specified, the number of output channels.
+        num_groups: the number of groups for the group norm.
+    """
+    def __init__(
+        self,
+        channels: int,
+        resolution: int,
+        out_channels: Optional[int] = None,
+        num_groups: int = 32
+    ):
+        super().__init__()
+        self.channels = channels
+        self.resolution = resolution
+        self.out_resolution = resolution * 2
+        self.out_channels = out_channels or channels
+
+        self.act_layers = nn.Sequential(
+            sp.SparseGroupNorm32(num_groups, channels),
+            sp.SparseSiLU()
+        )
+        
+        self.sub = sp.SparseSubdivide()
+        
+        self.out_layers = nn.Sequential(
+            sp.SparseConv3d(channels, self.out_channels, 3, indice_key=f"res_{self.out_resolution}"),
+            sp.SparseGroupNorm32(num_groups, self.out_channels),
+            sp.SparseSiLU(),
+            zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3, indice_key=f"res_{self.out_resolution}")),
+        )
+        
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        else:
+            self.skip_connection = sp.SparseConv3d(channels, self.out_channels, 1, indice_key=f"res_{self.out_resolution}")
+        
+    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        """
+        Apply the block to a Tensor, conditioned on a timestep embedding.
+
+        Args:
+            x: an [N x C x ...] Tensor of features.
+        Returns:
+            an [N x C x ...] Tensor of outputs.
+        """
+        h = self.act_layers(x)
+        h = self.sub(h)
+        x = self.sub(x)
+        h = self.out_layers(h)
+        h = h + self.skip_connection(x)
+        return h
+
+
+class SLatMeshDecoder(SparseTransformerBase):
+    def __init__(
+        self,
+        resolution: int,
+        model_channels: int,
+        latent_channels: int,
+        num_blocks: int,
+        num_heads: Optional[int] = None,
+        num_head_channels: Optional[int] = 64,
+        mlp_ratio: float = 4,
+        attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "swin",
+        window_size: int = 8,
+        pe_mode: Literal["ape", "rope"] = "ape",
+        use_fp16: bool = False,
+        use_checkpoint: bool = False,
+        qk_rms_norm: bool = False,
+        representation_config: dict = None,
+    ):
+        super().__init__(
+            in_channels=latent_channels,
+            model_channels=model_channels,
+            num_blocks=num_blocks,
+            num_heads=num_heads,
+            num_head_channels=num_head_channels,
+            mlp_ratio=mlp_ratio,
+            attn_mode=attn_mode,
+            window_size=window_size,
+            pe_mode=pe_mode,
+            use_fp16=use_fp16,
+            use_checkpoint=use_checkpoint,
+            qk_rms_norm=qk_rms_norm,
+        )
+        self.resolution = resolution
+        self.rep_config = representation_config
+        self.mesh_extractor = SparseFeatures2Mesh(res=self.resolution*4, use_color=self.rep_config.get('use_color', False))
+        self.out_channels = self.mesh_extractor.feats_channels
+        self.upsample = nn.ModuleList([
+            SparseSubdivideBlock3d(
+                channels=model_channels,
+                resolution=resolution,
+                out_channels=model_channels // 4
+            ),
+            SparseSubdivideBlock3d(
+                channels=model_channels // 4,
+                resolution=resolution * 2,
+                out_channels=model_channels // 8
+            )
+        ])
+        self.out_layer = sp.SparseLinear(model_channels // 8, self.out_channels)
+
+        self.initialize_weights()
+        if use_fp16:
+            self.convert_to_fp16()
+
+    def initialize_weights(self) -> None:
+        super().initialize_weights()
+        # Zero-out output layers:
+        nn.init.constant_(self.out_layer.weight, 0)
+        nn.init.constant_(self.out_layer.bias, 0)
+
+    def convert_to_fp16(self) -> None:
+        """
+        Convert the torso of the model to float16.
+        """
+        super().convert_to_fp16()
+        self.upsample.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self) -> None:
+        """
+        Convert the torso of the model to float32.
+        """
+        super().convert_to_fp32()
+        self.upsample.apply(convert_module_to_f32)  
+    
+    def to_representation(self, x: sp.SparseTensor) -> List[MeshExtractResult]:
+        """
+        Convert a batch of network outputs to 3D representations.
+
+        Args:
+            x: The [N x * x C] sparse tensor output by the network.
+
+        Returns:
+            list of representations
+        """
+        ret = []
+        for i in range(x.shape[0]):
+            mesh = self.mesh_extractor(x[i], training=self.training)
+            ret.append(mesh)
+        return ret
+
+    def forward(self, x: sp.SparseTensor) -> List[MeshExtractResult]:
+        h = super().forward(x)
+        for block in self.upsample:
+            h = block(h)
+        h = h.type(x.dtype)
+        h = self.out_layer(h)
+        return self.to_representation(h)

thirdparty/TRELLIS/trellis/trellis/models/structured_latent_vae/decoder_rf.py

@@ -0,0 +1,104 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from ...modules import sparse as sp
+from .base import SparseTransformerBase
+from ...representations import Strivec
+
+
+class SLatRadianceFieldDecoder(SparseTransformerBase):
+    def __init__(
+        self,
+        resolution: int,
+        model_channels: int,
+        latent_channels: int,
+        num_blocks: int,
+        num_heads: Optional[int] = None,
+        num_head_channels: Optional[int] = 64,
+        mlp_ratio: float = 4,
+        attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "swin",
+        window_size: int = 8,
+        pe_mode: Literal["ape", "rope"] = "ape",
+        use_fp16: bool = False,
+        use_checkpoint: bool = False,
+        qk_rms_norm: bool = False,
+        representation_config: dict = None,
+    ):
+        super().__init__(
+            in_channels=latent_channels,
+            model_channels=model_channels,
+            num_blocks=num_blocks,
+            num_heads=num_heads,
+            num_head_channels=num_head_channels,
+            mlp_ratio=mlp_ratio,
+            attn_mode=attn_mode,
+            window_size=window_size,
+            pe_mode=pe_mode,
+            use_fp16=use_fp16,
+            use_checkpoint=use_checkpoint,
+            qk_rms_norm=qk_rms_norm,
+        )
+        self.resolution = resolution
+        self.rep_config = representation_config
+        self._calc_layout()
+        self.out_layer = sp.SparseLinear(model_channels, self.out_channels)
+
+        self.initialize_weights()
+        if use_fp16:
+            self.convert_to_fp16()
+
+    def initialize_weights(self) -> None:
+        super().initialize_weights()
+        # Zero-out output layers:
+        nn.init.constant_(self.out_layer.weight, 0)
+        nn.init.constant_(self.out_layer.bias, 0)
+
+    def _calc_layout(self) -> None:
+        self.layout = {
+            'trivec': {'shape': (self.rep_config['rank'], 3, self.rep_config['dim']), 'size': self.rep_config['rank'] * 3 * self.rep_config['dim']},
+            'density': {'shape': (self.rep_config['rank'],), 'size': self.rep_config['rank']},
+            'features_dc': {'shape': (self.rep_config['rank'], 1, 3), 'size': self.rep_config['rank'] * 3},
+        }
+        start = 0
+        for k, v in self.layout.items():
+            v['range'] = (start, start + v['size'])
+            start += v['size']
+        self.out_channels = start    
+    
+    def to_representation(self, x: sp.SparseTensor) -> List[Strivec]:
+        """
+        Convert a batch of network outputs to 3D representations.
+
+        Args:
+            x: The [N x * x C] sparse tensor output by the network.
+
+        Returns:
+            list of representations
+        """
+        ret = []
+        for i in range(x.shape[0]):
+            representation = Strivec(
+                sh_degree=0,
+                resolution=self.resolution,
+                aabb=[-0.5, -0.5, -0.5, 1, 1, 1],
+                rank=self.rep_config['rank'],
+                dim=self.rep_config['dim'],
+                device='cuda',
+            )
+            representation.density_shift = 0.0
+            representation.position = (x.coords[x.layout[i]][:, 1:].float() + 0.5) / self.resolution
+            representation.depth = torch.full((representation.position.shape[0], 1), int(np.log2(self.resolution)), dtype=torch.uint8, device='cuda')
+            for k, v in self.layout.items():
+                setattr(representation, k, x.feats[x.layout[i]][:, v['range'][0]:v['range'][1]].reshape(-1, *v['shape']))
+            representation.trivec = representation.trivec + 1
+            ret.append(representation)
+        return ret
+
+    def forward(self, x: sp.SparseTensor) -> List[Strivec]:
+        h = super().forward(x)
+        h = h.type(x.dtype)
+        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+        h = self.out_layer(h)
+        return self.to_representation(h)

thirdparty/TRELLIS/trellis/trellis/models/structured_latent_vae/encoder.py

@@ -0,0 +1,72 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ...modules import sparse as sp
+from .base import SparseTransformerBase
+
+
+class SLatEncoder(SparseTransformerBase):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        model_channels: int,
+        latent_channels: int,
+        num_blocks: int,
+        num_heads: Optional[int] = None,
+        num_head_channels: Optional[int] = 64,
+        mlp_ratio: float = 4,
+        attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "swin",
+        window_size: int = 8,
+        pe_mode: Literal["ape", "rope"] = "ape",
+        use_fp16: bool = False,
+        use_checkpoint: bool = False,
+        qk_rms_norm: bool = False,
+    ):
+        super().__init__(
+            in_channels=in_channels,
+            model_channels=model_channels,
+            num_blocks=num_blocks,
+            num_heads=num_heads,
+            num_head_channels=num_head_channels,
+            mlp_ratio=mlp_ratio,
+            attn_mode=attn_mode,
+            window_size=window_size,
+            pe_mode=pe_mode,
+            use_fp16=use_fp16,
+            use_checkpoint=use_checkpoint,
+            qk_rms_norm=qk_rms_norm,
+        )
+        self.resolution = resolution
+        self.out_layer = sp.SparseLinear(model_channels, 2 * latent_channels)
+
+        self.initialize_weights()
+        if use_fp16:
+            self.convert_to_fp16()
+
+    def initialize_weights(self) -> None:
+        super().initialize_weights()
+        # Zero-out output layers:
+        nn.init.constant_(self.out_layer.weight, 0)
+        nn.init.constant_(self.out_layer.bias, 0)
+
+    def forward(self, x: sp.SparseTensor, sample_posterior=True, return_raw=False):
+        h = super().forward(x)
+        h = h.type(x.dtype)
+        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+        h = self.out_layer(h)
+        
+        # Sample from the posterior distribution
+        mean, logvar = h.feats.chunk(2, dim=-1)
+        if sample_posterior:
+            std = torch.exp(0.5 * logvar)
+            z = mean + std * torch.randn_like(std)
+        else:
+            z = mean
+        z = h.replace(z)
+            
+        if return_raw:
+            return z, mean, logvar
+        else:
+            return z

thirdparty/TRELLIS/trellis/trellis/modules/attention/__init__.py

@@ -0,0 +1,36 @@
+from typing import *
+
+BACKEND = 'flash_attn' 
+DEBUG = False
+
+def __from_env():
+    import os
+    
+    global BACKEND
+    global DEBUG
+    
+    env_attn_backend = os.environ.get('ATTN_BACKEND')
+    env_sttn_debug = os.environ.get('ATTN_DEBUG')
+    
+    if env_attn_backend is not None and env_attn_backend in ['xformers', 'flash_attn', 'sdpa', 'naive']:
+        BACKEND = env_attn_backend
+    if env_sttn_debug is not None:
+        DEBUG = env_sttn_debug == '1'
+
+    print(f"[ATTENTION] Using backend: {BACKEND}")
+        
+
+__from_env()
+    
+
+def set_backend(backend: Literal['xformers', 'flash_attn']):
+    global BACKEND
+    BACKEND = backend
+
+def set_debug(debug: bool):
+    global DEBUG
+    DEBUG = debug
+
+
+from .full_attn import *
+from .modules import *

thirdparty/TRELLIS/trellis/trellis/modules/attention/full_attn.py

@@ -0,0 +1,140 @@
+from typing import *
+import torch
+import math
+from . import DEBUG, BACKEND
+
+if BACKEND == 'xformers':
+    import xformers.ops as xops
+elif BACKEND == 'flash_attn':
+    import flash_attn
+elif BACKEND == 'sdpa':
+    from torch.nn.functional import scaled_dot_product_attention as sdpa
+elif BACKEND == 'naive':
+    pass
+else:
+    raise ValueError(f"Unknown attention backend: {BACKEND}")
+
+
+__all__ = [
+    'scaled_dot_product_attention',
+]
+
+
+def _naive_sdpa(q, k, v):
+    """
+    Naive implementation of scaled dot product attention.
+    """
+    q = q.permute(0, 2, 1, 3)   # [N, H, L, C]
+    k = k.permute(0, 2, 1, 3)   # [N, H, L, C]
+    v = v.permute(0, 2, 1, 3)   # [N, H, L, C]
+    scale_factor = 1 / math.sqrt(q.size(-1))
+    attn_weight = q @ k.transpose(-2, -1) * scale_factor
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+    out = attn_weight @ v
+    out = out.permute(0, 2, 1, 3)   # [N, L, H, C]
+    return out
+
+
+@overload
+def scaled_dot_product_attention(qkv: torch.Tensor) -> torch.Tensor:
+    """
+    Apply scaled dot product attention.
+
+    Args:
+        qkv (torch.Tensor): A [N, L, 3, H, C] tensor containing Qs, Ks, and Vs.
+    """
+    ...
+
+@overload
+def scaled_dot_product_attention(q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor:
+    """
+    Apply scaled dot product attention.
+
+    Args:
+        q (torch.Tensor): A [N, L, H, C] tensor containing Qs.
+        kv (torch.Tensor): A [N, L, 2, H, C] tensor containing Ks and Vs.
+    """
+    ...
+
+@overload
+def scaled_dot_product_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    """
+    Apply scaled dot product attention.
+
+    Args:
+        q (torch.Tensor): A [N, L, H, Ci] tensor containing Qs.
+        k (torch.Tensor): A [N, L, H, Ci] tensor containing Ks.
+        v (torch.Tensor): A [N, L, H, Co] tensor containing Vs.
+
+    Note:
+        k and v are assumed to have the same coordinate map.
+    """
+    ...
+
+def scaled_dot_product_attention(*args, **kwargs):
+    arg_names_dict = {
+        1: ['qkv'],
+        2: ['q', 'kv'],
+        3: ['q', 'k', 'v']
+    }
+    num_all_args = len(args) + len(kwargs)
+    assert num_all_args in arg_names_dict, f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3"
+    for key in arg_names_dict[num_all_args][len(args):]:
+        assert key in kwargs, f"Missing argument {key}"
+
+    if num_all_args == 1:
+        qkv = args[0] if len(args) > 0 else kwargs['qkv']
+        assert len(qkv.shape) == 5 and qkv.shape[2] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, L, 3, H, C]"
+        device = qkv.device
+
+    elif num_all_args == 2:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        kv = args[1] if len(args) > 1 else kwargs['kv']
+        assert q.shape[0] == kv.shape[0], f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}"
+        assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]"
+        assert len(kv.shape) == 5, f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]"
+        device = q.device
+
+    elif num_all_args == 3:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        k = args[1] if len(args) > 1 else kwargs['k']
+        v = args[2] if len(args) > 2 else kwargs['v']
+        assert q.shape[0] == k.shape[0] == v.shape[0], f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}"
+        assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]"
+        assert len(k.shape) == 4, f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]"
+        assert len(v.shape) == 4, f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]"
+        device = q.device    
+
+    if BACKEND == 'xformers':
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=2)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=2)
+        out = xops.memory_efficient_attention(q, k, v)
+    elif BACKEND == 'flash_attn':
+        if num_all_args == 1:
+            out = flash_attn.flash_attn_qkvpacked_func(qkv)
+        elif num_all_args == 2:
+            out = flash_attn.flash_attn_kvpacked_func(q, kv)
+        elif num_all_args == 3:
+            out = flash_attn.flash_attn_func(q, k, v)
+    elif BACKEND == 'sdpa':
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=2)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=2)
+        q = q.permute(0, 2, 1, 3)   # [N, H, L, C]
+        k = k.permute(0, 2, 1, 3)   # [N, H, L, C]
+        v = v.permute(0, 2, 1, 3)   # [N, H, L, C]
+        out = sdpa(q, k, v)         # [N, H, L, C]
+        out = out.permute(0, 2, 1, 3)   # [N, L, H, C]
+    elif BACKEND == 'naive':
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=2)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=2)
+        out = _naive_sdpa(q, k, v)
+    else:
+        raise ValueError(f"Unknown attention module: {BACKEND}")
+    
+    return out

thirdparty/TRELLIS/trellis/trellis/modules/attention/modules.py

@@ -0,0 +1,146 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .full_attn import scaled_dot_product_attention
+
+
+class MultiHeadRMSNorm(nn.Module):
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(heads, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return (F.normalize(x.float(), dim = -1) * self.gamma * self.scale).to(x.dtype)
+
+
+class RotaryPositionEmbedder(nn.Module):
+    def __init__(self, hidden_size: int, in_channels: int = 3):
+        super().__init__()
+        assert hidden_size % 2 == 0, "Hidden size must be divisible by 2"
+        self.hidden_size = hidden_size
+        self.in_channels = in_channels
+        self.freq_dim = hidden_size // in_channels // 2
+        self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim
+        self.freqs = 1.0 / (10000 ** self.freqs)
+        
+    def _get_phases(self, indices: torch.Tensor) -> torch.Tensor:
+        self.freqs = self.freqs.to(indices.device)
+        phases = torch.outer(indices, self.freqs)
+        phases = torch.polar(torch.ones_like(phases), phases)
+        return phases
+        
+    def _rotary_embedding(self, x: torch.Tensor, phases: torch.Tensor) -> torch.Tensor:
+        x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        x_rotated = x_complex * phases
+        x_embed = torch.view_as_real(x_rotated).reshape(*x_rotated.shape[:-1], -1).to(x.dtype)
+        return x_embed
+        
+    def forward(self, q: torch.Tensor, k: torch.Tensor, indices: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            q (sp.SparseTensor): [..., N, D] tensor of queries
+            k (sp.SparseTensor): [..., N, D] tensor of keys
+            indices (torch.Tensor): [..., N, C] tensor of spatial positions
+        """
+        if indices is None:
+            indices = torch.arange(q.shape[-2], device=q.device)
+            if len(q.shape) > 2:
+                indices = indices.unsqueeze(0).expand(q.shape[:-2] + (-1,))
+        
+        phases = self._get_phases(indices.reshape(-1)).reshape(*indices.shape[:-1], -1)
+        if phases.shape[1] < self.hidden_size // 2:
+            phases = torch.cat([phases, torch.polar(
+                torch.ones(*phases.shape[:-1], self.hidden_size // 2 - phases.shape[1], device=phases.device),
+                torch.zeros(*phases.shape[:-1], self.hidden_size // 2 - phases.shape[1], device=phases.device)
+            )], dim=-1)
+        q_embed = self._rotary_embedding(q, phases)
+        k_embed = self._rotary_embedding(k, phases)
+        return q_embed, k_embed
+    
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        ctx_channels: Optional[int]=None,
+        type: Literal["self", "cross"] = "self",
+        attn_mode: Literal["full", "windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        qkv_bias: bool = True,
+        use_rope: bool = False,
+        qk_rms_norm: bool = False,
+    ):
+        super().__init__()
+        assert channels % num_heads == 0
+        assert type in ["self", "cross"], f"Invalid attention type: {type}"
+        assert attn_mode in ["full", "windowed"], f"Invalid attention mode: {attn_mode}"
+        assert type == "self" or attn_mode == "full", "Cross-attention only supports full attention"
+        
+        if attn_mode == "windowed":
+            raise NotImplementedError("Windowed attention is not yet implemented")
+        
+        self.channels = channels
+        self.head_dim = channels // num_heads
+        self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+        self.num_heads = num_heads
+        self._type = type
+        self.attn_mode = attn_mode
+        self.window_size = window_size
+        self.shift_window = shift_window
+        self.use_rope = use_rope
+        self.qk_rms_norm = qk_rms_norm
+
+        if self._type == "self":
+            self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+        else:
+            self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
+            self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
+            
+        if self.qk_rms_norm:
+            self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+            self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+            
+        self.to_out = nn.Linear(channels, channels)
+
+        if use_rope:
+            self.rope = RotaryPositionEmbedder(channels)
+    
+    def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None, indices: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, L, C = x.shape
+        if self._type == "self":
+            qkv = self.to_qkv(x)
+            qkv = qkv.reshape(B, L, 3, self.num_heads, -1)
+            if self.use_rope:
+                q, k, v = qkv.unbind(dim=2)
+                q, k = self.rope(q, k, indices)
+                qkv = torch.stack([q, k, v], dim=2)
+            if self.attn_mode == "full":
+                if self.qk_rms_norm:
+                    q, k, v = qkv.unbind(dim=2)
+                    q = self.q_rms_norm(q)
+                    k = self.k_rms_norm(k)
+                    h = scaled_dot_product_attention(q, k, v)
+                else:
+                    h = scaled_dot_product_attention(qkv)
+            elif self.attn_mode == "windowed":
+                raise NotImplementedError("Windowed attention is not yet implemented")
+        else:
+            Lkv = context.shape[1]
+            q = self.to_q(x)
+            kv = self.to_kv(context)
+            q = q.reshape(B, L, self.num_heads, -1)
+            kv = kv.reshape(B, Lkv, 2, self.num_heads, -1)
+            if self.qk_rms_norm:
+                q = self.q_rms_norm(q)
+                k, v = kv.unbind(dim=2)
+                k = self.k_rms_norm(k)
+                h = scaled_dot_product_attention(q, k, v)
+            else:
+                h = scaled_dot_product_attention(q, kv)
+        h = h.reshape(B, L, -1)
+        h = self.to_out(h)
+        return h

thirdparty/TRELLIS/trellis/trellis/modules/norm.py

@@ -0,0 +1,25 @@
+import torch
+import torch.nn as nn
+
+
+class LayerNorm32(nn.LayerNorm):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x.float()).type(x.dtype)
+    
+
+class GroupNorm32(nn.GroupNorm):
+    """
+    A GroupNorm layer that converts to float32 before the forward pass.
+    """
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x.float()).type(x.dtype)
+    
+    
+class ChannelLayerNorm32(LayerNorm32):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        DIM = x.dim()
+        x = x.permute(0, *range(2, DIM), 1).contiguous()
+        x = super().forward(x)
+        x = x.permute(0, DIM-1, *range(1, DIM-1)).contiguous()
+        return x
+    

thirdparty/TRELLIS/trellis/trellis/modules/sparse/__init__.py

@@ -0,0 +1,102 @@
+from typing import *
+
+BACKEND = 'spconv' 
+DEBUG = False
+ATTN = 'flash_attn'
+
+def __from_env():
+    import os
+    
+    global BACKEND
+    global DEBUG
+    global ATTN
+    
+    env_sparse_backend = os.environ.get('SPARSE_BACKEND')
+    env_sparse_debug = os.environ.get('SPARSE_DEBUG')
+    env_sparse_attn = os.environ.get('SPARSE_ATTN_BACKEND')
+    if env_sparse_attn is None:
+        env_sparse_attn = os.environ.get('ATTN_BACKEND')
+
+    if env_sparse_backend is not None and env_sparse_backend in ['spconv', 'torchsparse']:
+        BACKEND = env_sparse_backend
+    if env_sparse_debug is not None:
+        DEBUG = env_sparse_debug == '1'
+    if env_sparse_attn is not None and env_sparse_attn in ['xformers', 'flash_attn']:
+        ATTN = env_sparse_attn
+        
+    print(f"[SPARSE] Backend: {BACKEND}, Attention: {ATTN}")
+        
+
+__from_env()
+    
+
+def set_backend(backend: Literal['spconv', 'torchsparse']):
+    global BACKEND
+    BACKEND = backend
+
+def set_debug(debug: bool):
+    global DEBUG
+    DEBUG = debug
+
+def set_attn(attn: Literal['xformers', 'flash_attn']):
+    global ATTN
+    ATTN = attn
+    
+    
+import importlib
+
+__attributes = {
+    'SparseTensor': 'basic',
+    'sparse_batch_broadcast': 'basic',
+    'sparse_batch_op': 'basic',
+    'sparse_cat': 'basic',
+    'sparse_unbind': 'basic',
+    'SparseGroupNorm': 'norm',
+    'SparseLayerNorm': 'norm',
+    'SparseGroupNorm32': 'norm',
+    'SparseLayerNorm32': 'norm',
+    'SparseReLU': 'nonlinearity',
+    'SparseSiLU': 'nonlinearity',
+    'SparseGELU': 'nonlinearity',
+    'SparseActivation': 'nonlinearity',
+    'SparseLinear': 'linear',
+    'sparse_scaled_dot_product_attention': 'attention',
+    'SerializeMode': 'attention',
+    'sparse_serialized_scaled_dot_product_self_attention': 'attention',
+    'sparse_windowed_scaled_dot_product_self_attention': 'attention',
+    'SparseMultiHeadAttention': 'attention',
+    'SparseConv3d': 'conv',
+    'SparseInverseConv3d': 'conv',
+    'SparseDownsample': 'spatial',
+    'SparseUpsample': 'spatial',
+    'SparseSubdivide' : 'spatial'
+}
+
+__submodules = ['transformer']
+
+__all__ = list(__attributes.keys()) + __submodules
+
+def __getattr__(name):
+    if name not in globals():
+        if name in __attributes:
+            module_name = __attributes[name]
+            module = importlib.import_module(f".{module_name}", __name__)
+            globals()[name] = getattr(module, name)
+        elif name in __submodules:
+            module = importlib.import_module(f".{name}", __name__)
+            globals()[name] = module
+        else:
+            raise AttributeError(f"module {__name__} has no attribute {name}")
+    return globals()[name]
+
+
+# For Pylance
+if __name__ == '__main__':
+    from .basic import *
+    from .norm import *
+    from .nonlinearity import *
+    from .linear import *
+    from .attention import *
+    from .conv import *
+    from .spatial import *
+    import transformer

thirdparty/TRELLIS/trellis/trellis/modules/sparse/attention/__init__.py

@@ -0,0 +1,4 @@
+from .full_attn import *
+from .serialized_attn import *
+from .windowed_attn import *
+from .modules import *

thirdparty/TRELLIS/trellis/trellis/modules/sparse/attention/full_attn.py

@@ -0,0 +1,215 @@
+from typing import *
+import torch
+from .. import SparseTensor
+from .. import DEBUG, ATTN
+
+if ATTN == 'xformers':
+    import xformers.ops as xops
+elif ATTN == 'flash_attn':
+    import flash_attn
+else:
+    raise ValueError(f"Unknown attention module: {ATTN}")
+
+
+__all__ = [
+    'sparse_scaled_dot_product_attention',
+]
+
+
+@overload
+def sparse_scaled_dot_product_attention(qkv: SparseTensor) -> SparseTensor:
+    """
+    Apply scaled dot product attention to a sparse tensor.
+
+    Args:
+        qkv (SparseTensor): A [N, *, 3, H, C] sparse tensor containing Qs, Ks, and Vs.
+    """
+    ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: SparseTensor, kv: Union[SparseTensor, torch.Tensor]) -> SparseTensor:
+    """
+    Apply scaled dot product attention to a sparse tensor.
+
+    Args:
+        q (SparseTensor): A [N, *, H, C] sparse tensor containing Qs.
+        kv (SparseTensor or torch.Tensor): A [N, *, 2, H, C] sparse tensor or a [N, L, 2, H, C] dense tensor containing Ks and Vs.
+    """
+    ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: torch.Tensor, kv: SparseTensor) -> torch.Tensor:
+    """
+    Apply scaled dot product attention to a sparse tensor.
+
+    Args:
+        q (SparseTensor): A [N, L, H, C] dense tensor containing Qs.
+        kv (SparseTensor or torch.Tensor): A [N, *, 2, H, C] sparse tensor containing Ks and Vs.
+    """
+    ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: SparseTensor, k: SparseTensor, v: SparseTensor) -> SparseTensor:
+    """
+    Apply scaled dot product attention to a sparse tensor.
+
+    Args:
+        q (SparseTensor): A [N, *, H, Ci] sparse tensor containing Qs.
+        k (SparseTensor): A [N, *, H, Ci] sparse tensor containing Ks.
+        v (SparseTensor): A [N, *, H, Co] sparse tensor containing Vs.
+
+    Note:
+        k and v are assumed to have the same coordinate map.
+    """
+    ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: SparseTensor, k: torch.Tensor, v: torch.Tensor) -> SparseTensor:
+    """
+    Apply scaled dot product attention to a sparse tensor.
+
+    Args:
+        q (SparseTensor): A [N, *, H, Ci] sparse tensor containing Qs.
+        k (torch.Tensor): A [N, L, H, Ci] dense tensor containing Ks.
+        v (torch.Tensor): A [N, L, H, Co] dense tensor containing Vs.
+    """
+    ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: torch.Tensor, k: SparseTensor, v: SparseTensor) -> torch.Tensor:
+    """
+    Apply scaled dot product attention to a sparse tensor.
+
+    Args:
+        q (torch.Tensor): A [N, L, H, Ci] dense tensor containing Qs.
+        k (SparseTensor): A [N, *, H, Ci] sparse tensor containing Ks.
+        v (SparseTensor): A [N, *, H, Co] sparse tensor containing Vs.
+    """
+    ...
+
+def sparse_scaled_dot_product_attention(*args, **kwargs):
+    arg_names_dict = {
+        1: ['qkv'],
+        2: ['q', 'kv'],
+        3: ['q', 'k', 'v']
+    }
+    num_all_args = len(args) + len(kwargs)
+    assert num_all_args in arg_names_dict, f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3"
+    for key in arg_names_dict[num_all_args][len(args):]:
+        assert key in kwargs, f"Missing argument {key}"
+
+    if num_all_args == 1:
+        qkv = args[0] if len(args) > 0 else kwargs['qkv']
+        assert isinstance(qkv, SparseTensor), f"qkv must be a SparseTensor, got {type(qkv)}"
+        assert len(qkv.shape) == 4 and qkv.shape[1] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]"
+        device = qkv.device
+
+        s = qkv
+        q_seqlen = [qkv.layout[i].stop - qkv.layout[i].start for i in range(qkv.shape[0])]
+        kv_seqlen = q_seqlen
+        qkv = qkv.feats     # [T, 3, H, C]
+
+    elif num_all_args == 2:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        kv = args[1] if len(args) > 1 else kwargs['kv']
+        assert isinstance(q, SparseTensor) and isinstance(kv, (SparseTensor, torch.Tensor)) or \
+               isinstance(q, torch.Tensor) and isinstance(kv, SparseTensor), \
+               f"Invalid types, got {type(q)} and {type(kv)}"
+        assert q.shape[0] == kv.shape[0], f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}"
+        device = q.device
+
+        if isinstance(q, SparseTensor):
+            assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, C]"
+            s = q
+            q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
+            q = q.feats     # [T_Q, H, C]
+        else:
+            assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]"
+            s = None
+            N, L, H, C = q.shape
+            q_seqlen = [L] * N
+            q = q.reshape(N * L, H, C)   # [T_Q, H, C]
+
+        if isinstance(kv, SparseTensor):
+            assert len(kv.shape) == 4 and kv.shape[1] == 2, f"Invalid shape for kv, got {kv.shape}, expected [N, *, 2, H, C]"
+            kv_seqlen = [kv.layout[i].stop - kv.layout[i].start for i in range(kv.shape[0])]
+            kv = kv.feats     # [T_KV, 2, H, C]
+        else:
+            assert len(kv.shape) == 5, f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]"
+            N, L, _, H, C = kv.shape
+            kv_seqlen = [L] * N
+            kv = kv.reshape(N * L, 2, H, C)   # [T_KV, 2, H, C]
+
+    elif num_all_args == 3:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        k = args[1] if len(args) > 1 else kwargs['k']
+        v = args[2] if len(args) > 2 else kwargs['v']
+        assert isinstance(q, SparseTensor) and isinstance(k, (SparseTensor, torch.Tensor)) and type(k) == type(v) or \
+               isinstance(q, torch.Tensor) and isinstance(k, SparseTensor) and isinstance(v, SparseTensor), \
+               f"Invalid types, got {type(q)}, {type(k)}, and {type(v)}"
+        assert q.shape[0] == k.shape[0] == v.shape[0], f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}"
+        device = q.device
+
+        if isinstance(q, SparseTensor):
+            assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, Ci]"
+            s = q
+            q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
+            q = q.feats     # [T_Q, H, Ci]
+        else:
+            assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]"
+            s = None
+            N, L, H, CI = q.shape
+            q_seqlen = [L] * N
+            q = q.reshape(N * L, H, CI)  # [T_Q, H, Ci]
+
+        if isinstance(k, SparseTensor):
+            assert len(k.shape) == 3, f"Invalid shape for k, got {k.shape}, expected [N, *, H, Ci]"
+            assert len(v.shape) == 3, f"Invalid shape for v, got {v.shape}, expected [N, *, H, Co]"
+            kv_seqlen = [k.layout[i].stop - k.layout[i].start for i in range(k.shape[0])]
+            k = k.feats     # [T_KV, H, Ci]
+            v = v.feats     # [T_KV, H, Co]
+        else:
+            assert len(k.shape) == 4, f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]"
+            assert len(v.shape) == 4, f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]"
+            N, L, H, CI, CO = *k.shape, v.shape[-1]
+            kv_seqlen = [L] * N
+            k = k.reshape(N * L, H, CI)     # [T_KV, H, Ci]
+            v = v.reshape(N * L, H, CO)     # [T_KV, H, Co]
+
+    if DEBUG:
+        if s is not None:
+            for i in range(s.shape[0]):
+                assert (s.coords[s.layout[i]] == i).all(), f"SparseScaledDotProductSelfAttention: batch index mismatch"
+        if num_all_args in [2, 3]:
+            assert q.shape[:2] == [1, sum(q_seqlen)], f"SparseScaledDotProductSelfAttention: q shape mismatch"
+        if num_all_args == 3:
+            assert k.shape[:2] == [1, sum(kv_seqlen)], f"SparseScaledDotProductSelfAttention: k shape mismatch"
+            assert v.shape[:2] == [1, sum(kv_seqlen)], f"SparseScaledDotProductSelfAttention: v shape mismatch"
+
+    if ATTN == 'xformers':
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=1)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=1)
+        q = q.unsqueeze(0)
+        k = k.unsqueeze(0)
+        v = v.unsqueeze(0)
+        mask = xops.fmha.BlockDiagonalMask.from_seqlens(q_seqlen, kv_seqlen)
+        out = xops.memory_efficient_attention(q, k, v, mask)[0]
+    elif ATTN == 'flash_attn':
+        cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device)
+        if num_all_args in [2, 3]:
+            cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+        if num_all_args == 1:
+            out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv, cu_seqlens_q, max(q_seqlen))
+        elif num_all_args == 2:
+            out = flash_attn.flash_attn_varlen_kvpacked_func(q, kv, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+        elif num_all_args == 3:
+            out = flash_attn.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+    else:
+        raise ValueError(f"Unknown attention module: {ATTN}")
+    
+    if s is not None:
+        return s.replace(out)
+    else:
+        return out.reshape(N, L, H, -1)

thirdparty/TRELLIS/trellis/trellis/modules/sparse/attention/modules.py

@@ -0,0 +1,139 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .. import SparseTensor
+from .full_attn import sparse_scaled_dot_product_attention
+from .serialized_attn import SerializeMode, sparse_serialized_scaled_dot_product_self_attention
+from .windowed_attn import sparse_windowed_scaled_dot_product_self_attention
+from ...attention import RotaryPositionEmbedder
+
+
+class SparseMultiHeadRMSNorm(nn.Module):
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(heads, dim))
+
+    def forward(self, x: Union[SparseTensor, torch.Tensor]) -> Union[SparseTensor, torch.Tensor]:
+        x_type = x.dtype
+        x = x.float()
+        if isinstance(x, SparseTensor):
+            x = x.replace(F.normalize(x.feats, dim=-1))
+        else:
+            x = F.normalize(x, dim=-1)            
+        return (x * self.gamma * self.scale).to(x_type)
+
+
+class SparseMultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        ctx_channels: Optional[int] = None,
+        type: Literal["self", "cross"] = "self",
+        attn_mode: Literal["full", "serialized", "windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_sequence: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        serialize_mode: Optional[SerializeMode] = None,
+        qkv_bias: bool = True,
+        use_rope: bool = False,
+        qk_rms_norm: bool = False,
+    ):
+        super().__init__()
+        assert channels % num_heads == 0
+        assert type in ["self", "cross"], f"Invalid attention type: {type}"
+        assert attn_mode in ["full", "serialized", "windowed"], f"Invalid attention mode: {attn_mode}"
+        assert type == "self" or attn_mode == "full", "Cross-attention only supports full attention"
+        assert type == "self" or use_rope is False, "Rotary position embeddings only supported for self-attention"
+        self.channels = channels
+        self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+        self.num_heads = num_heads
+        self._type = type
+        self.attn_mode = attn_mode
+        self.window_size = window_size
+        self.shift_sequence = shift_sequence
+        self.shift_window = shift_window
+        self.serialize_mode = serialize_mode
+        self.use_rope = use_rope
+        self.qk_rms_norm = qk_rms_norm
+
+        if self._type == "self":
+            self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+        else:
+            self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
+            self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
+        
+        if self.qk_rms_norm:
+            self.q_rms_norm = SparseMultiHeadRMSNorm(channels // num_heads, num_heads)
+            self.k_rms_norm = SparseMultiHeadRMSNorm(channels // num_heads, num_heads)
+            
+        self.to_out = nn.Linear(channels, channels)
+
+        if use_rope:
+            self.rope = RotaryPositionEmbedder(channels)
+
+    @staticmethod
+    def _linear(module: nn.Linear, x: Union[SparseTensor, torch.Tensor]) -> Union[SparseTensor, torch.Tensor]:
+        if isinstance(x, SparseTensor):
+            return x.replace(module(x.feats))
+        else:
+            return module(x)
+
+    @staticmethod
+    def _reshape_chs(x: Union[SparseTensor, torch.Tensor], shape: Tuple[int, ...]) -> Union[SparseTensor, torch.Tensor]:
+        if isinstance(x, SparseTensor):
+            return x.reshape(*shape)
+        else:
+            return x.reshape(*x.shape[:2], *shape)
+
+    def _fused_pre(self, x: Union[SparseTensor, torch.Tensor], num_fused: int) -> Union[SparseTensor, torch.Tensor]:
+        if isinstance(x, SparseTensor):
+            x_feats = x.feats.unsqueeze(0)
+        else:
+            x_feats = x
+        x_feats = x_feats.reshape(*x_feats.shape[:2], num_fused, self.num_heads, -1)
+        return x.replace(x_feats.squeeze(0)) if isinstance(x, SparseTensor) else x_feats
+
+    def _rope(self, qkv: SparseTensor) -> SparseTensor:
+        q, k, v = qkv.feats.unbind(dim=1)   # [T, H, C]
+        q, k = self.rope(q, k, qkv.coords[:, 1:])
+        qkv = qkv.replace(torch.stack([q, k, v], dim=1)) 
+        return qkv
+    
+    def forward(self, x: Union[SparseTensor, torch.Tensor], context: Optional[Union[SparseTensor, torch.Tensor]] = None) -> Union[SparseTensor, torch.Tensor]:
+        if self._type == "self":
+            qkv = self._linear(self.to_qkv, x)
+            qkv = self._fused_pre(qkv, num_fused=3)
+            if self.use_rope:
+                qkv = self._rope(qkv)
+            if self.qk_rms_norm:
+                q, k, v = qkv.unbind(dim=1)
+                q = self.q_rms_norm(q)
+                k = self.k_rms_norm(k)
+                qkv = qkv.replace(torch.stack([q.feats, k.feats, v.feats], dim=1))
+            if self.attn_mode == "full":
+                h = sparse_scaled_dot_product_attention(qkv)
+            elif self.attn_mode == "serialized":
+                h = sparse_serialized_scaled_dot_product_self_attention(
+                    qkv, self.window_size, serialize_mode=self.serialize_mode, shift_sequence=self.shift_sequence, shift_window=self.shift_window
+                )
+            elif self.attn_mode == "windowed":
+                h = sparse_windowed_scaled_dot_product_self_attention(
+                    qkv, self.window_size, shift_window=self.shift_window
+                )
+        else:
+            q = self._linear(self.to_q, x)
+            q = self._reshape_chs(q, (self.num_heads, -1))
+            kv = self._linear(self.to_kv, context)
+            kv = self._fused_pre(kv, num_fused=2)
+            if self.qk_rms_norm:
+                q = self.q_rms_norm(q)
+                k, v = kv.unbind(dim=1)
+                k = self.k_rms_norm(k)
+                kv = kv.replace(torch.stack([k.feats, v.feats], dim=1))
+            h = sparse_scaled_dot_product_attention(q, kv)
+        h = self._reshape_chs(h, (-1,))
+        h = self._linear(self.to_out, h)
+        return h

thirdparty/TRELLIS/trellis/trellis/modules/sparse/attention/serialized_attn.py

@@ -0,0 +1,193 @@
+from typing import *
+from enum import Enum
+import torch
+import math
+from .. import SparseTensor
+from .. import DEBUG, ATTN
+
+if ATTN == 'xformers':
+    import xformers.ops as xops
+elif ATTN == 'flash_attn':
+    import flash_attn
+else:
+    raise ValueError(f"Unknown attention module: {ATTN}")
+
+
+__all__ = [
+    'sparse_serialized_scaled_dot_product_self_attention',
+]
+
+
+class SerializeMode(Enum):
+    Z_ORDER = 0
+    Z_ORDER_TRANSPOSED = 1
+    HILBERT = 2
+    HILBERT_TRANSPOSED = 3
+
+
+SerializeModes = [
+    SerializeMode.Z_ORDER,
+    SerializeMode.Z_ORDER_TRANSPOSED,
+    SerializeMode.HILBERT,
+    SerializeMode.HILBERT_TRANSPOSED
+]
+
+
+def calc_serialization(
+    tensor: SparseTensor,
+    window_size: int,
+    serialize_mode: SerializeMode = SerializeMode.Z_ORDER,
+    shift_sequence: int = 0,
+    shift_window: Tuple[int, int, int] = (0, 0, 0)
+) -> Tuple[torch.Tensor, torch.Tensor, List[int]]:
+    """
+    Calculate serialization and partitioning for a set of coordinates.
+
+    Args:
+        tensor (SparseTensor): The input tensor.
+        window_size (int): The window size to use.
+        serialize_mode (SerializeMode): The serialization mode to use.
+        shift_sequence (int): The shift of serialized sequence.
+        shift_window (Tuple[int, int, int]): The shift of serialized coordinates.
+
+    Returns:
+        (torch.Tensor, torch.Tensor): Forwards and backwards indices.
+    """
+    fwd_indices = []
+    bwd_indices = []
+    seq_lens = []
+    seq_batch_indices = []
+    offsets = [0]
+    
+    if 'vox2seq' not in globals():
+        import vox2seq
+
+    # Serialize the input
+    serialize_coords = tensor.coords[:, 1:].clone()
+    serialize_coords += torch.tensor(shift_window, dtype=torch.int32, device=tensor.device).reshape(1, 3)
+    if serialize_mode == SerializeMode.Z_ORDER:
+        code = vox2seq.encode(serialize_coords, mode='z_order', permute=[0, 1, 2])
+    elif serialize_mode == SerializeMode.Z_ORDER_TRANSPOSED:
+        code = vox2seq.encode(serialize_coords, mode='z_order', permute=[1, 0, 2])
+    elif serialize_mode == SerializeMode.HILBERT:
+        code = vox2seq.encode(serialize_coords, mode='hilbert', permute=[0, 1, 2])
+    elif serialize_mode == SerializeMode.HILBERT_TRANSPOSED:
+        code = vox2seq.encode(serialize_coords, mode='hilbert', permute=[1, 0, 2])
+    else:
+        raise ValueError(f"Unknown serialize mode: {serialize_mode}")
+    
+    for bi, s in enumerate(tensor.layout):
+        num_points = s.stop - s.start
+        num_windows = (num_points + window_size - 1) // window_size
+        valid_window_size = num_points / num_windows
+        to_ordered = torch.argsort(code[s.start:s.stop])
+        if num_windows == 1:
+            fwd_indices.append(to_ordered)
+            bwd_indices.append(torch.zeros_like(to_ordered).scatter_(0, to_ordered, torch.arange(num_points, device=tensor.device)))
+            fwd_indices[-1] += s.start
+            bwd_indices[-1] += offsets[-1]
+            seq_lens.append(num_points)
+            seq_batch_indices.append(bi)
+            offsets.append(offsets[-1] + seq_lens[-1])
+        else:
+            # Partition the input
+            offset = 0
+            mids = [(i + 0.5) * valid_window_size + shift_sequence for i in range(num_windows)]
+            split = [math.floor(i * valid_window_size + shift_sequence) for i in range(num_windows + 1)]
+            bwd_index = torch.zeros((num_points,), dtype=torch.int64, device=tensor.device)
+            for i in range(num_windows):
+                mid = mids[i]
+                valid_start = split[i]
+                valid_end = split[i + 1]
+                padded_start = math.floor(mid - 0.5 * window_size)
+                padded_end = padded_start + window_size
+                fwd_indices.append(to_ordered[torch.arange(padded_start, padded_end, device=tensor.device) % num_points])
+                offset += valid_start - padded_start
+                bwd_index.scatter_(0, fwd_indices[-1][valid_start-padded_start:valid_end-padded_start], torch.arange(offset, offset + valid_end - valid_start, device=tensor.device))
+                offset += padded_end - valid_start
+                fwd_indices[-1] += s.start
+            seq_lens.extend([window_size] * num_windows)
+            seq_batch_indices.extend([bi] * num_windows)
+            bwd_indices.append(bwd_index + offsets[-1])
+            offsets.append(offsets[-1] + num_windows * window_size)
+
+    fwd_indices = torch.cat(fwd_indices)
+    bwd_indices = torch.cat(bwd_indices)
+
+    return fwd_indices, bwd_indices, seq_lens, seq_batch_indices
+    
+
+def sparse_serialized_scaled_dot_product_self_attention(
+    qkv: SparseTensor,
+    window_size: int,
+    serialize_mode: SerializeMode = SerializeMode.Z_ORDER,
+    shift_sequence: int = 0,
+    shift_window: Tuple[int, int, int] = (0, 0, 0)
+) -> SparseTensor:
+    """
+    Apply serialized scaled dot product self attention to a sparse tensor.
+
+    Args:
+        qkv (SparseTensor): [N, *, 3, H, C] sparse tensor containing Qs, Ks, and Vs.
+        window_size (int): The window size to use.
+        serialize_mode (SerializeMode): The serialization mode to use.
+        shift_sequence (int): The shift of serialized sequence.
+        shift_window (Tuple[int, int, int]): The shift of serialized coordinates.
+        shift (int): The shift to use.
+    """
+    assert len(qkv.shape) == 4 and qkv.shape[1] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]"
+
+    serialization_spatial_cache_name = f'serialization_{serialize_mode}_{window_size}_{shift_sequence}_{shift_window}'
+    serialization_spatial_cache = qkv.get_spatial_cache(serialization_spatial_cache_name)
+    if serialization_spatial_cache is None:
+        fwd_indices, bwd_indices, seq_lens, seq_batch_indices = calc_serialization(qkv, window_size, serialize_mode, shift_sequence, shift_window)
+        qkv.register_spatial_cache(serialization_spatial_cache_name, (fwd_indices, bwd_indices, seq_lens, seq_batch_indices))
+    else:
+        fwd_indices, bwd_indices, seq_lens, seq_batch_indices = serialization_spatial_cache
+
+    M = fwd_indices.shape[0]
+    T = qkv.feats.shape[0]
+    H = qkv.feats.shape[2]
+    C = qkv.feats.shape[3]
+    
+    qkv_feats = qkv.feats[fwd_indices]      # [M, 3, H, C]
+
+    if DEBUG:
+        start = 0
+        qkv_coords = qkv.coords[fwd_indices]
+        for i in range(len(seq_lens)):
+            assert (qkv_coords[start:start+seq_lens[i], 0] == seq_batch_indices[i]).all(), f"SparseWindowedScaledDotProductSelfAttention: batch index mismatch"
+            start += seq_lens[i]
+
+    if all([seq_len == window_size for seq_len in seq_lens]):
+        B = len(seq_lens)
+        N = window_size
+        qkv_feats = qkv_feats.reshape(B, N, 3, H, C)
+        if ATTN == 'xformers':
+            q, k, v = qkv_feats.unbind(dim=2)                       # [B, N, H, C]
+            out = xops.memory_efficient_attention(q, k, v)          # [B, N, H, C]
+        elif ATTN == 'flash_attn':
+            out = flash_attn.flash_attn_qkvpacked_func(qkv_feats)   # [B, N, H, C]
+        else:
+            raise ValueError(f"Unknown attention module: {ATTN}")
+        out = out.reshape(B * N, H, C)                              # [M, H, C]
+    else:
+        if ATTN == 'xformers':
+            q, k, v = qkv_feats.unbind(dim=1)                       # [M, H, C]
+            q = q.unsqueeze(0)                                      # [1, M, H, C]
+            k = k.unsqueeze(0)                                      # [1, M, H, C]
+            v = v.unsqueeze(0)                                      # [1, M, H, C]
+            mask = xops.fmha.BlockDiagonalMask.from_seqlens(seq_lens)
+            out = xops.memory_efficient_attention(q, k, v, mask)[0] # [M, H, C]
+        elif ATTN == 'flash_attn':
+            cu_seqlens = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(seq_lens), dim=0)], dim=0) \
+                        .to(qkv.device).int()
+            out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, cu_seqlens, max(seq_lens)) # [M, H, C]
+
+    out = out[bwd_indices]      # [T, H, C]
+
+    if DEBUG:
+        qkv_coords = qkv_coords[bwd_indices]
+        assert torch.equal(qkv_coords, qkv.coords), "SparseWindowedScaledDotProductSelfAttention: coordinate mismatch"
+
+    return qkv.replace(out)