XPart/partgen/partformer_pipeline.py

import torch
from .utils.misc import logger, synchronize_timer
import inspect
from typing import List, Optional
import trimesh
import numpy as np
from tqdm import tqdm
import copy
from typing import List, Optional, Union
import os
from safetensors.torch import load_file
from .utils.mesh_utils import (
    SampleMesh,
    load_surface_points,
    sample_bbox_points_from_trimesh,
    explode_mesh,
    fix_mesh,
)
from .utils.misc import (
    init_from_ckpt,
    instantiate_from_config,
    get_config_from_file,
    smart_load_model,
)
from easydict import EasyDict

import json
from diffusers.utils.torch_utils import randn_tensor
from pathlib import Path


@synchronize_timer("Export to trimesh")
def export_to_trimesh(mesh_output):
    if isinstance(mesh_output, list):
        outputs = []
        for mesh in mesh_output:
            if mesh is None:
                outputs.append(None)
            else:
                mesh.mesh_f = mesh.mesh_f[:, ::-1]
                mesh_output = trimesh.Trimesh(mesh.mesh_v, mesh.mesh_f)
                mesh_output = fix_mesh(mesh_output)
                outputs.append(mesh_output)
        return outputs
    else:
        mesh_output.mesh_f = mesh_output.mesh_f[:, ::-1]
        mesh_output = trimesh.Trimesh(mesh_output.mesh_v, mesh_output.mesh_f)
        mesh_output = fix_mesh(mesh_output)
        return mesh_output


def retrieve_timesteps(
    scheduler,
    num_inference_steps: Optional[int] = None,
    device: Optional[Union[str, torch.device]] = None,
    timesteps: Optional[List[int]] = None,
    sigmas: Optional[Union[List[float], np.ndarray]] = None,
    **kwargs,
):
    """
    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.

    Args:
        scheduler (`SchedulerMixin`):
            The scheduler to get timesteps from.
        num_inference_steps (`int`):
            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
            must be `None`.
        device (`str` or `torch.device`, *optional*):
            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
        timesteps (`List[int]`, *optional*):
            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
            `num_inference_steps` and `sigmas` must be `None`.
        sigmas (`List[float]`, *optional*):
            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
            `num_inference_steps` and `timesteps` must be `None`.

    Returns:
        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
        second element is the number of inference steps.
    """
    if timesteps is not None and sigmas is not None:
        raise ValueError(
            "Only one of `timesteps` or `sigmas` can be passed. Please choose one to"
            " set custom values"
        )
    if timesteps is not None:
        accepts_timesteps = "timesteps" in set(
            inspect.signature(scheduler.set_timesteps).parameters.keys()
        )
        if not accepts_timesteps:
            raise ValueError(
                f"The current scheduler class {scheduler.__class__}'s `set_timesteps`"
                " does not support custom timestep schedules. Please check whether you"
                " are using the correct scheduler."
            )
        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
        timesteps = scheduler.timesteps
        num_inference_steps = len(timesteps)
    elif sigmas is not None:
        accept_sigmas = "sigmas" in set(
            inspect.signature(scheduler.set_timesteps).parameters.keys()
        )
        if not accept_sigmas:
            raise ValueError(
                f"The current scheduler class {scheduler.__class__}'s `set_timesteps`"
                " does not support custom sigmas schedules. Please check whether you"
                " are using the correct scheduler."
            )
        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
        timesteps = scheduler.timesteps
        num_inference_steps = len(timesteps)
    else:
        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
        timesteps = scheduler.timesteps
    return timesteps, num_inference_steps


class TokenAllocMixin:

    def allocate_tokens(self, bboxes, num_latents=512):
        return np.array([num_latents] * bboxes.shape[0])


class PartFormerPipeline(TokenAllocMixin):

    def __init__(
        self,
        vae,
        model,
        scheduler,
        conditioner,
        bbox_predictor=None,
        verbose=False,
        **kwargs,
    ):
        self.vae = vae
        self.model = model
        self.scheduler = scheduler
        self.conditioner = conditioner
        self.kwargs = kwargs
        self.bbox_predictor = bbox_predictor
        self.verbose = verbose
        self.kwargs = kwargs

    @classmethod
    @synchronize_timer("Hunyuan3D PartGen Pipeline Model Loading")
    def from_single_file(
        cls,
        ckpt_path=None,
        config=None,
        device="cuda",
        dtype=torch.float32,
        use_safetensors=None,
        ignore_keys=(),
        **kwargs,
    ):
        # prepare config
        if config is None:
            config = get_config_from_file(
                str(
                    Path(__file__).parent.parent
                    / "config"
                    / "partformer_full_pipeline_512_with_sonata.yaml"
                )
            )
        # TODO:
        if ckpt_path is None:
            ckpt_path = str(
                Path(__file__).parent
                / "ckpts"
                / "partformer_full_pipeline_512_with_sonata.ckpt"
            )
        # load ckpt
        if use_safetensors:
            ckpt_path = ckpt_path.replace(".ckpt", ".safetensors")
        if not os.path.exists(ckpt_path):
            raise FileNotFoundError(f"Model file {ckpt_path} not found")
        logger.info(f"Loading model from {ckpt_path}")

        if use_safetensors:
            # parse safetensors
            import safetensors.torch

            safetensors_ckpt = safetensors.torch.load_file(ckpt_path, device="cpu")
            ckpt = {}
            for key, value in safetensors_ckpt.items():
                model_name = key.split(".")[0]
                new_key = key[len(model_name) + 1 :]
                if model_name not in ckpt:
                    ckpt[model_name] = {}
                ckpt[model_name][new_key] = value
        else:
            # ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
            ckpt = torch.load(ckpt_path, map_location="cpu")
        # load model
        model = instantiate_from_config(config["model"])
        # model.load_state_dict(ckpt["model"])
        init_from_ckpt(model, ckpt, prefix="model", ignore_keys=ignore_keys)
        vae = instantiate_from_config(config["shapevae"])
        # vae.load_state_dict(ckpt["shapevae"], strict=False)
        init_from_ckpt(vae, ckpt, prefix="shapevae", ignore_keys=ignore_keys)
        if config.get("conditioner", None) is not None:
            conditioner = instantiate_from_config(config["conditioner"])
            init_from_ckpt(
                conditioner, ckpt, prefix="conditioner", ignore_keys=ignore_keys
            )
        else:
            conditioner = vae
        scheduler = instantiate_from_config(config["scheduler"])
        bbox_predictor = instantiate_from_config(config.get("bbox_predictor", None))
        model_kwargs = dict(
            vae=vae,
            model=model,
            scheduler=scheduler,
            conditioner=conditioner,
            bbox_predictor=bbox_predictor,  # TODO: add bbox predictor
            device=device,
            dtype=dtype,
        )
        model_kwargs.update(kwargs)
        return cls(**model_kwargs)

    @classmethod
    def from_pretrained(
        cls,
        model_path="tencent/Hunyuan3D-Part",
        dtype=torch.float32,
        device="cuda",
        **kwargs,
    ):
        model_dir = smart_load_model(
            model_path=model_path,
        )
        model_ckpt = load_file(os.path.join(model_dir, "model/model.safetensors"))
        conditioner_ckpt = load_file(
            os.path.join(model_dir, "conditioner/conditioner.safetensors")
        )
        shapevae_ckpt = load_file(
            os.path.join(model_dir, "shapevae/shapevae.safetensors")
        )
        p3sam_path = os.path.join(model_dir, "p3sam/p3sam.safetensors")
        with open(os.path.join(model_dir, "model/config.json"), "r") as f:
            model_config = EasyDict(json.load(f))
        with open(os.path.join(model_dir, "conditioner/config.json"), "r") as f:
            conditioner_config = EasyDict(json.load(f))
        with open(os.path.join(model_dir, "shapevae/config.json"), "r") as f:
            shapevae_config = EasyDict(json.load(f))
        with open(os.path.join(model_dir, "scheduler/config.json"), "r") as f:
            scheduler_config = EasyDict(json.load(f))
        with open(os.path.join(model_dir, "p3sam/config.json"), "r") as f:
            bbox_predictor_config = EasyDict(json.load(f))
            bbox_predictor_config["params"]["ckpt_path"] = p3sam_path
        # load model
        model = instantiate_from_config(model_config)
        model.load_state_dict(model_ckpt)
        vae = instantiate_from_config(shapevae_config)
        vae.load_state_dict(shapevae_ckpt)
        conditioner = instantiate_from_config(conditioner_config)
        conditioner.load_state_dict(conditioner_ckpt)
        scheduler = instantiate_from_config(scheduler_config)
        bbox_predictor = instantiate_from_config(bbox_predictor_config)
        model_kwargs = dict(
            vae=vae,
            model=model,
            scheduler=scheduler,
            conditioner=conditioner,
            bbox_predictor=bbox_predictor,  # TODO: add bbox predictor
            device=device,
            dtype=dtype,
        )
        model_kwargs.update(kwargs)
        return cls(**model_kwargs)

    def compile(self):
        self.vae = torch.compile(self.vae)
        self.model = torch.compile(self.model)
        self.conditioner = torch.compile(self.conditioner)

    def to(self, device=None, dtype=None):
        if dtype is not None:
            self.dtype = dtype
            self.vae.to(dtype=dtype)
            self.model.to(dtype=dtype)
            self.conditioner.to(dtype=dtype)
        if device is not None:
            self.device = torch.device(device)
            self.vae.to(device)
            self.model.to(device)
            self.conditioner.to(device)

    def prepare_extra_step_kwargs(self, generator, eta):
        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
        # and should be between [0, 1]

        accepts_eta = "eta" in set(
            inspect.signature(self.scheduler.step).parameters.keys()
        )
        extra_step_kwargs = {}
        if accepts_eta:
            extra_step_kwargs["eta"] = eta

        # check if the scheduler accepts generator
        accepts_generator = "generator" in set(
            inspect.signature(self.scheduler.step).parameters.keys()
        )
        if accepts_generator:
            extra_step_kwargs["generator"] = generator
        return extra_step_kwargs

    def predict_bbox(
        self, mesh: trimesh.Trimesh, scale_box=1.0, drop_normal=True, seed=42
    ):
        """
        Predict the bounding box of the object surface.
        Args:
            obj_surface (`torch.Tensor`): [B, N, 3]
        Returns:
            `torch.Tensor`: [B, K, 2, 3] where K is the number of bounding boxes
        """
        if self.bbox_predictor is None:
            raise ValueError("bbox_predictor is not set.")
        aabb, face_ids, mesh = self.bbox_predictor.predict_aabb(
            mesh, post_process=True, seed=seed
        )
        # aabb, face_ids, mesh = self.bbox_predictor.predict_aabb(mesh, post_process=False)
        aabb = torch.from_numpy(aabb)
        return aabb

    def prepare_latents(
        self, batch_size, latent_shape, dtype, device, generator, latents=None
    ):
        # prepare latents for different parts
        shape = (batch_size, *latent_shape)
        if isinstance(generator, list) and len(generator) != batch_size:
            raise ValueError(
                f"You have passed a list of generators of length {len(generator)}, but"
                f" requested an effective batch size of {batch_size}. Make sure the"
                " batch size matches the length of the generators."
            )

        if latents is None:
            latents = randn_tensor(
                shape, generator=generator, device=device, dtype=dtype
            )
        else:
            latents = latents.to(device)

        # scale the initial noise by the standard deviation required by the scheduler
        latents = latents * getattr(self.scheduler, "init_noise_sigma", 1.0)
        return latents

    @synchronize_timer("Encode cond")
    def encode_cond(
        self,
        part_surface_inbbox,
        object_surface,
        do_classifier_free_guidance,
    ):
        bsz = object_surface.shape[0]
        cond = self.conditioner(part_surface_inbbox, object_surface)

        if do_classifier_free_guidance:
            # TODO: do_classifier_free_guidance, un_cond
            un_cond = {k: torch.zeros_like(v) for k, v in cond.items()}

            def cat_recursive(a, b):
                if isinstance(a, torch.Tensor):
                    return torch.cat([a, b], dim=0).to(self.dtype)
                out = {}
                for k in a.keys():
                    out[k] = cat_recursive(a[k], b[k])
                return out

            cond = cat_recursive(cond, un_cond)
        return cond

    def normalize_mesh(self, mesh):
        vertices = mesh.vertices
        min_xyz = np.min(vertices, axis=0)
        max_xyz = np.max(vertices, axis=0)
        center = (min_xyz + max_xyz) / 2.0
        # scale = np.max(np.linalg.norm(vertices - center, axis=1))
        scale = np.max(max_xyz - min_xyz) / 2 / 0.8
        vertices = (vertices - center) / scale
        mesh.vertices = vertices
        return mesh, center, scale

    def check_inputs(
        self,
        obj_surface=None,
        obj_surface_raw=None,
        mesh_path=None,
        mesh=None,
        aabb=None,
        part_surface_inbbox=None,
        seed=42,
    ):
        """
        Check the inputs of the pipeline.
        Args:
            obj_surface (`torch.Tensor`): [B, N, 3+3+1]
            mesh_path (`str`): path to the mesh file
            mesh (`trimesh.Trimesh`): mesh object
            aabb (`torch.Tensor`): [B, K, 2, 3]
            part_surface_inbbox (`torch.Tensor`): [B, K,N, 3+3+1]
        """
        if obj_surface is None:
            if mesh_path is None and (mesh is None and obj_surface_raw is None):
                raise ValueError(
                    "obj_surface or mesh_path/mesh/obj_surface_raw must be provided."
                )
        elif aabb is None or part_surface_inbbox is None:
            raise ValueError(
                "aabb and part_surface_inbbox must be provided if obj_surface is"
                " provided."
            )
        else:
            assert aabb.shape[0] == part_surface_inbbox.shape[0], "Batch size mismatch."
        center = np.zeros(3)
        scale = 1.0
        # 1. Load object surface and sample
        if obj_surface is None:
            if obj_surface_raw is None:
                if mesh is not None:
                    obj_surface_raw = SampleMesh(
                        mesh.vertices, mesh.faces, -1, seed=seed
                    )
                elif mesh_path is not None:
                    mesh = trimesh.load(mesh_path, force="mesh")
                    mesh, center, scale = self.normalize_mesh(mesh)
                    print(f"Normalize mesh: {center}, {scale}")
                    obj_surface_raw = SampleMesh(
                        mesh.vertices, mesh.faces, -1, seed=seed
                    )
                else:
                    raise ValueError("obj_surface or mesh_path/mesh must be provided.")
            rng = np.random.default_rng(seed=seed)
            obj_surface, _ = load_surface_points(
                rng,
                obj_surface_raw["random_surface"],
                obj_surface_raw["sharp_surface"],
                pc_size=81920,
                pc_sharpedge_size=0,
                return_sharpedge_label=True,
                return_normal=True,
            )
            obj_surface = obj_surface.unsqueeze(0)
        # 2. load aabb
        if aabb is None:
            aabb = self.predict_bbox(mesh, seed=seed)
            print(f"Get bbox from bbox_predictor: {aabb.shape}")
        else:
            if isinstance(aabb, np.ndarray):
                aabb = torch.from_numpy(aabb)
            # normalize aabb by mesh scale and center
            aabb = aabb.float()
            aabb = (aabb - torch.from_numpy(center).float()) / scale

        # 3. load part surface in bbox
        if part_surface_inbbox is None:
            part_surface_inbbox, valid_parts_mask = sample_bbox_points_from_trimesh(
                mesh, aabb, num_points=81920, seed=seed
            )
            aabb = aabb[valid_parts_mask]
            aabb = aabb.unsqueeze(0)
            part_surface_inbbox = part_surface_inbbox.unsqueeze(0)
        return (
            obj_surface,
            aabb,
            part_surface_inbbox,
            mesh,
            center,
            scale,
        )

    def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=torch.float32):
        """
        See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298

        Args:
            timesteps (`torch.Tensor`):
                generate embedding vectors at these timesteps
            embedding_dim (`int`, *optional*, defaults to 512):
                dimension of the embeddings to generate
            dtype:
                data type of the generated embeddings

        Returns:
            `torch.FloatTensor`: Embedding vectors with shape `(len(timesteps), embedding_dim)`
        """
        assert len(w.shape) == 1
        w = w * 1000.0

        half_dim = embedding_dim // 2
        emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
        emb = w.to(dtype)[:, None] * emb[None, :]
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
        if embedding_dim % 2 == 1:  # zero pad
            emb = torch.nn.functional.pad(emb, (0, 1))
        assert emb.shape == (w.shape[0], embedding_dim)
        return emb

    def _export(
        self,
        latents,
        output_type="trimesh",
        box_v=1.01,
        mc_level=0.0,
        num_chunks=20000,
        octree_resolution=256,
        mc_algo="mc",
        enable_pbar=True,
        **kwargs,
    ):
        if not output_type == "latent":
            latents = 1.0 / self.vae.scale_factor * latents
            latents = self.vae(latents)
            outputs = self.vae.latent2mesh_2(
                # outputs = self.vae.latents2mesh(
                latents,
                bounds=box_v,
                mc_level=mc_level,
                octree_depth=8,
                num_chunks=num_chunks,
                octree_resolution=octree_resolution,
                mc_mode=mc_algo,
                # enable_pbar=enable_pbar,
                **kwargs,
            )
        else:
            outputs = latents

        if output_type == "trimesh":
            outputs = export_to_trimesh(outputs)

        return outputs

    @torch.no_grad()
    @torch.autocast("cuda", dtype=torch.bfloat16)
    def __call__(
        self,
        obj_surface=None,
        obj_surface_raw=None,
        mesh_path=None,
        mesh=None,
        aabb=None,
        part_surface_inbbox=None,
        num_inference_steps: int = 50,
        timesteps: Optional[List[int]] = None,
        sigmas: Optional[List[float]] = None,
        eta: float = 0.0,
        # guidance_scale: float = 7.5,
        guidance_scale: float = -1.0,
        dual_guidance_scale: float = 10.5,
        dual_guidance: bool = True,
        generator=None,
        seed=42,
        # marching cubes
        box_v=1.01,
        octree_resolution=512,
        mc_level=-1 / 512,
        num_chunks=400000,
        mc_algo="mc",
        output_type: Optional[str] = "trimesh",
        enable_pbar=True,
        **kwargs,
    ):
        """
        Args:
            obj_surface (`torch.Tensor`): [B, N, 3+3+1]
            aabb (`torch.Tensor`): [B, K, 2, 3]
            part_surface_inbbox (`torch.Tensor`): [B, K,N, 3+3+1]
        Returns:
            `trimesh.Scene` : single object composed of multiple parts
        """
        callback = kwargs.pop("callback", None)
        callback_steps = kwargs.pop("callback_steps", None)
        do_classifier_free_guidance = guidance_scale >= 0 and not (
            hasattr(self.model, "guidance_embed") and self.model.guidance_embed is True
        )
        # 1. Check inputs and predict bbox if not provided
        device = self.device
        dtype = self.dtype
        obj_surface, aabb, part_surface_inbbox, mesh, center, scale = self.check_inputs(
            obj_surface,
            obj_surface_raw,
            mesh_path,
            mesh,
            aabb,
            part_surface_inbbox,
            seed=seed,
        )
        if self.verbose:
            # return gt mesh with bbox
            mesh_bbox = trimesh.Scene()
            if mesh is not None:
                mesh_bbox.add_geometry(mesh)
            else:
                mesh = trimesh.points.PointCloud(
                    obj_surface[0, :, :3].float().cpu().numpy()
                )
                mesh_bbox.add_geometry(mesh)
            for bbox in aabb[0]:
                box = trimesh.path.creation.box_outline()
                box.vertices *= (bbox[1] - bbox[0]).float().cpu().numpy()
                box.vertices += (bbox[0] + bbox[1]).float().cpu().numpy() / 2
                mesh_bbox.add_geometry(box)
        #  Convert to device and dtype
        obj_surface = obj_surface.to(device=device, dtype=dtype)
        aabb = aabb.to(device=device, dtype=dtype)
        part_surface_inbbox = part_surface_inbbox.to(device=device, dtype=dtype)
        batch_size, num_parts, N, dim = part_surface_inbbox.shape
        # TODO: support batch size > 1
        assert batch_size == 1, "Batch size > 1 is not supported yet."
        # 2. Prepare latent variables
        # TODO:allocate tokens for each parts
        num_tokens = torch.tensor(
            [self.allocate_tokens(x, self.vae.latent_shape[0]) for x in aabb],
            device=device,
        )
        latent_shape = self.vae.latent_shape
        latents = self.prepare_latents(
            num_parts, latent_shape, dtype, device, generator
        )
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
        # 3. condition
        cond = self.encode_cond(
            part_surface_inbbox.reshape(batch_size * num_parts, N, dim),
            obj_surface.expand(batch_size * num_parts, -1, -1),
            do_classifier_free_guidance,
        )
        # 4. guidance_cond for controling sampling
        guidance_cond = None
        if getattr(self.model, "guidance_cond_proj_dim", None) is not None:
            logger.info("Using lcm guidance scale")
            guidance_scale_tensor = torch.tensor(guidance_scale - 1).repeat(batch_size)
            guidance_cond = self.get_guidance_scale_embedding(
                guidance_scale_tensor, embedding_dim=self.model.guidance_cond_proj_dim
            ).to(device=device, dtype=latents.dtype)

        # 5. Prepare timesteps
        # NOTE: this is slightly different from common usage, we start from 0.
        sigmas = np.linspace(0, 1, num_inference_steps) if sigmas is None else sigmas
        timesteps, num_inference_steps = retrieve_timesteps(
            self.scheduler,
            num_inference_steps,
            device,
            sigmas=sigmas,
        )

        torch.cuda.empty_cache()

        # 6. Denoising loop
        with synchronize_timer("Diffusion Sampling"):
            for i, t in enumerate(
                tqdm(timesteps, disable=not enable_pbar, desc="Diffusion Sampling:")
            ):
                # expand the latents if we are doing classifier free guidance
                if do_classifier_free_guidance:
                    latent_model_input = torch.cat([latents] * 2)
                    aabb = torch.repeat_interleave(aabb, 2, dim=0)
                else:
                    latent_model_input = latents

                # NOTE: we assume model get timesteps ranged from 0 to 1
                timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype)
                timestep = timestep / self.scheduler.config.num_train_timesteps
                noise_pred = self.model(
                    latent_model_input,
                    timestep,
                    cond,
                    aabb=aabb,
                    num_tokens=num_tokens,
                    guidance_cond=guidance_cond,
                )

                if do_classifier_free_guidance:
                    noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2)
                    noise_pred = noise_pred_uncond + guidance_scale * (
                        noise_pred_cond - noise_pred_uncond
                    )

                # compute the previous noisy sample x_t -> x_t-1
                outputs = self.scheduler.step(noise_pred, t, latents)
                latents = outputs.prev_sample

                if callback is not None and i % callback_steps == 0:
                    step_idx = i // getattr(self.scheduler, "order", 1)
                    callback(step_idx, t, outputs)

        # latents2mesh
        # part_latents = torch.split(latents, num_tokens[0].tolist(), dim=1)
        out = trimesh.Scene()
        for i, part_latent in enumerate(latents):
            try:
                part_mesh = self._export(
                    latents=part_latent.unsqueeze(0),
                    output_type=output_type,
                    box_v=box_v,
                    mc_level=mc_level,
                    num_chunks=num_chunks,
                    octree_resolution=octree_resolution,
                    mc_algo=mc_algo,
                    enable_pbar=enable_pbar,
                )[0]
                out.add_geometry(part_mesh)
                random_color = np.random.randint(0, 255, size=3)
                part_mesh.visual.face_colors = random_color
            except Exception as e:
                logger.error(f"Failed to export part {i} with error {e}")
        print(f"Denormalize mesh: {center}, {scale}")
        for key in out.geometry.keys():
            _v = out.geometry[key].vertices
            _v = _v * scale + center
            out.geometry[key].vertices = _v

        if self.verbose:
            explode_object = explode_mesh(copy.deepcopy(out), explosion_scale=0.2)
            # add bbox into out
            out_bbox = trimesh.Scene()
            out_bbox.add_geometry(out)
            for bbox in aabb[0]:
                box = trimesh.path.creation.box_outline()
                box.vertices *= (bbox[1] - bbox[0]).float().cpu().numpy()
                box.vertices += (bbox[0] + bbox[1]).float().cpu().numpy() / 2
                box.vertices = box.vertices * scale + center
                out_bbox.add_geometry(box)
            return out, (out_bbox, mesh_bbox, explode_object)
        else:
            # return only the generated mesh
            return out, None