step1x3d_geometry/models/pipelines/pipeline.py

# Some parts of this file are refer to Hugging Face Diffusers library.
import os
import json
import warnings
from typing import Callable, List, Optional, Union, Dict, Any
import PIL.Image
import trimesh
import rembg
import torch
import numpy as np
from huggingface_hub import hf_hub_download

from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from diffusers.utils import BaseOutput
from diffusers.utils.torch_utils import randn_tensor
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.loaders import (
    FluxIPAdapterMixin,
    FluxLoraLoaderMixin,
    FromSingleFileMixin,
    TextualInversionLoaderMixin,
)
from .pipeline_utils import (
    TransformerDiffusionMixin,
    preprocess_image,
    retrieve_timesteps,
    remove_floater,
    remove_degenerate_face,
    reduce_face,
    smart_load_model,
)
from transformers import (
    BitImageProcessor,
)

import step1x3d_geometry
from step1x3d_geometry.models.autoencoders.surface_extractors import MeshExtractResult
from step1x3d_geometry.utils.config import ExperimentConfig, load_config
from ..autoencoders.michelangelo_autoencoder import MichelangeloAutoencoder
from ..conditional_encoders.dinov2_encoder import Dinov2Encoder
from ..conditional_encoders.t5_encoder import T5Encoder
from ..conditional_encoders.label_encoder import LabelEncoder
from ..transformers.flux_transformer_1d import FluxDenoiser


class Step1X3DGeometryPipelineOutput(BaseOutput):
    """
    Output class for image pipelines.

    Args:
        images (`List[PIL.Image.Image]` or `torch.Tensor`):
            List of PIL images or a tensor representing the input images.
        meshes (`List[trimesh.Trimesh]` or `np.ndarray`)
            List of denoised trimesh meshes of length `batch_size` or a tuple of NumPy array with shape `((vertices, 3), (faces, 3)) of length `batch_size``.
    """

    image: PIL.Image.Image
    mesh: Union[trimesh.Trimesh, MeshExtractResult, np.ndarray]


class Step1X3DGeometryPipeline(
    DiffusionPipeline, FromSingleFileMixin, TransformerDiffusionMixin
):
    """
    Step1X-3D Geometry Pipeline, generate high-quality meshes conditioned on image/caption/label inputs

    Args:
        scheduler (FlowMatchEulerDiscreteScheduler):
            The diffusion scheduler controlling the denoising process
        vae (MichelangeloAutoencoder):
            Variational Autoencoder for latent space compression/reconstruction
        transformer (FluxDenoiser):
            Transformer-based denoising model
        visual_encoder (Dinov2Encoder):
            Pretrained visual encoder for image feature extraction
        caption_encoder (T5Encoder):
            Text encoder for processing natural language captions
        label_encoder (LabelEncoder):
            Auxiliary text encoder for label conditioning
        visual_eature_extractor (BitImageProcessor):
            Preprocessor for input images

    Note:
        - CPU offloading sequence: visual_encoder → caption_encoder → label_encoder → transformer → vae
        - Optional components: visual_encoder, visual_eature_extractor, caption_encoder, label_encoder
    """

    model_cpu_offload_seq = (
        "visual_encoder->caption_encoder->label_encoder->transformer->vae"
    )
    _optional_components = [
        "visual_encoder",
        "visual_eature_extractor",
        "caption_encoder",
        "label_encoder",
    ]

    @classmethod
    def from_pretrained(cls, model_path, subfolder='.', **kwargs):
        local_model_path = smart_load_model(model_path, subfolder)
        return super().from_pretrained(local_model_path, **kwargs)

    def __init__(
        self,
        scheduler: FlowMatchEulerDiscreteScheduler,
        vae: MichelangeloAutoencoder,
        transformer: FluxDenoiser,
        visual_encoder: Dinov2Encoder,
        caption_encoder: T5Encoder,
        label_encoder: LabelEncoder,
        visual_eature_extractor: BitImageProcessor,
    ):
        super().__init__()

        self.register_modules(
            vae=vae,
            transformer=transformer,
            scheduler=scheduler,
            visual_encoder=visual_encoder,
            caption_encoder=caption_encoder,
            label_encoder=label_encoder,
            visual_eature_extractor=visual_eature_extractor,
        )

    @property
    def guidance_scale(self):
        return self._guidance_scale

    @property
    def do_classifier_free_guidance(self):
        return self._guidance_scale > 1

    @property
    def num_timesteps(self):
        return self._num_timesteps

    def check_inputs(
        self,
        image,
    ):
        r"""
        Check if the inputs are valid. Raise an error if not.
        """
        if isinstance(image, str):
            assert os.path.isfile(image) or image.startswith(
                "http"
            ), "Input image must be a valid URL or a file path."
        elif isinstance(image, (torch.Tensor, PIL.Image.Image)):
            raise ValueError(
                "Input image must be a `torch.Tensor` or `PIL.Image.Image`."
            )

    def encode_image(self, image, device, num_meshes_per_prompt):
        dtype = next(self.visual_encoder.parameters()).dtype

        image_embeds = self.visual_encoder.encode_image(image)
        image_embeds = image_embeds.repeat_interleave(num_meshes_per_prompt, dim=0)

        uncond_image_embeds = self.visual_encoder.empty_image_embeds.repeat(
            image_embeds.shape[0], 1, 1
        ).to(image_embeds)

        return image_embeds, uncond_image_embeds

    def encode_caption(self, caption, device, num_meshes_per_prompt):
        dtype = next(self.label_encoder.parameters()).dtype

        caption_embeds = self.caption_encoder.encode_text([caption])
        caption_embeds = caption_embeds.repeat_interleave(num_meshes_per_prompt, dim=0)

        uncond_caption_embeds = self.caption_encoder.empty_text_embeds.repeat(
            caption_embeds.shape[0], 1, 1
        ).to(caption_embeds)

        return caption_embeds, uncond_caption_embeds

    def encode_label(self, label, device, num_meshes_per_prompt):
        dtype = next(self.label_encoder.parameters()).dtype

        label_embeds = self.label_encoder.encode_label([label])
        label_embeds = label_embeds.repeat_interleave(num_meshes_per_prompt, dim=0)

        uncond_label_embeds = self.label_encoder.empty_label_embeds.repeat(
            label_embeds.shape[0], 1, 1
        ).to(label_embeds)

        return label_embeds, uncond_label_embeds

    def prepare_latents(
        self,
        batch_size,
        num_tokens,
        num_channels_latents,
        dtype,
        device,
        generator,
        latents: Optional[torch.Tensor] = None,
    ):
        if latents is not None:
            return latents.to(device=device, dtype=dtype)

        shape = (batch_size, num_tokens, num_channels_latents)

        if isinstance(generator, list) and len(generator) != batch_size:
            raise ValueError(
                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
            )

        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)

        return latents

    @torch.no_grad()
    def __call__(
        self,
        image: Union[torch.FloatTensor, PIL.Image.Image, str],
        label: Optional[str] = None,
        caption: Optional[str] = None,
        num_inference_steps: int = 30,
        timesteps: List[int] = None,
        num_meshes_per_prompt: int = 1,
        guidance_scale: float = 7.5,
        generator: Optional[int] = None,
        latents: Optional[torch.FloatTensor] = None,
        force_remove_background: bool = False,
        background_color: List[int] = [255, 255, 255],
        foreground_ratio: float = 0.95,
        surface_extractor_type: Optional[str] = None,
        bounds: float = 1.05,
        mc_level: float = 0.0,
        octree_resolution: int = 384,
        output_type: str = "trimesh",
        do_remove_floater: bool = True,
        do_remove_degenerate_face: bool = False,
        do_reduce_face: bool = True,
        do_shade_smooth: bool = True,
        max_facenum: int = 200000,
        return_dict: bool = True,
        use_zero_init: Optional[bool] = True,
        zero_steps: Optional[int] = 0,
    ):
        r"""
        Function invoked when calling the pipeline for generation.

        Args:
            image (`torch.FloatTensor` or `PIL.Image.Image` or `str`):
                `Image`, or tensor representing an image batch, or path to an image file. The image will be encoded to
                its CLIP/DINO-v2 embedding which the DiT will be conditioned on.
            label (`str`):
                The label of the generated mesh, like {"symmetry": "asymmetry", "edge_type": "smooth"}
            num_inference_steps (`int`, *optional*, defaults to 30):
                The number of denoising steps. More denoising steps usually lead to a higher quality mesh at the expense
                of slower inference.
            timesteps (`List[int]`, *optional*):
                Custom timesteps to use for the denoising process. If not provided, will use equally spaced timesteps.
            num_meshes_per_prompt (`int`, *optional*, defaults to 1):
                The number of meshes to generate per input image.
            guidance_scale (`float`, *optional*, defaults to 7.5):
                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                Higher guidance scale encourages generation that closely matches the input image.
            generator (`int`, *optional*):
                A seed to make the generation deterministic.
            latents (`torch.FloatTensor`, *optional*):
                Pre-generated noisy latents to use as inputs for mesh generation.
            force_remove_background (`bool`, *optional*, defaults to `False`):
                Whether to force remove the background from the input image before processing.
            background_color (`List[int]`, *optional*, defaults to `[255, 255, 255]`):
                RGB color values for the background if it needs to be removed or modified.
            foreground_ratio (`float`, *optional*, defaults to 0.95):
                Ratio of the image to consider as foreground when processing.
            surface_extractor_type (`str`, *optional*, defaults to "mc"):
                Type of surface extraction method to use ("mc" for Marching Cubes or other available methods).
            bounds (`float`, *optional*, defaults to 1.05):
                Bounding box size for the generated mesh.
            mc_level (`float`, *optional*, defaults to 0.0):
                Iso-surface level value for Marching Cubes extraction.
            octree_resolution (`int`, *optional*, defaults to 256):
                Resolution of the octree used for mesh generation.
            output_type (`str`, *optional*, defaults to "trimesh"):
                Type of output mesh format ("trimesh" or other supported formats).
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a `MeshPipelineOutput` instead of a plain tuple.

        Returns:
            [`MeshPipelineOutput`] or `tuple`:
                If `return_dict` is `True`, [`MeshPipelineOutput`] is returned, otherwise a `tuple` is returned where the
                first element is a list of generated meshes and the second element is a list of corresponding metadata.
        """
        # 0. Check inputs. Raise error if not correct
        self.check_inputs(
            image=image,
        )
        device = self._execution_device
        self._guidance_scale = guidance_scale

        # 1. Define call parameters
        if isinstance(image, torch.Tensor):
            batch_size = image.shape[0]
        elif isinstance(image, PIL.Image.Image) or isinstance(image, str):
            batch_size = 1

        # 2. Preprocess input image
        if isinstance(image, torch.Tensor):
            assert image.ndim == 3  # H, W, 3
            image_pil = TF.to_pil_image(image)
        elif isinstance(image, PIL.Image.Image):
            image_pil = image
        elif isinstance(image, str):
            if image.startswith("http"):
                import requests

                image_pil = PIL.Image.open(requests.get(image, stream=True).raw)
            else:
                image_pil = PIL.Image.open(image)
        image_pil = preprocess_image(image_pil, force=force_remove_background, background_color=background_color, foreground_ratio=foreground_ratio)  # remove the background images

        # 3. Encode condition
        image_embeds, negative_image_embeds = self.encode_image(
            image_pil, device, num_meshes_per_prompt
        )
        if self.do_classifier_free_guidance and image_embeds is not None:
            image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0)
        # 3.1 Encode label condition
        label_embeds = None
        if self.transformer.cfg.use_label_condition:
            if label is not None:
                label_embeds, negative_label_embeds = self.encode_label(
                    label, device, num_meshes_per_prompt
                )
                if self.do_classifier_free_guidance:
                    label_embeds = torch.cat(
                        [negative_label_embeds, label_embeds], dim=0
                    )
            else:
                uncond_label_embeds = self.label_encoder.empty_label_embeds.repeat(
                    num_meshes_per_prompt, 1, 1
                ).to(image_embeds)
                if self.do_classifier_free_guidance:
                    label_embeds = torch.cat(
                        [uncond_label_embeds, uncond_label_embeds], dim=0
                    )
        # 3.3 Encode caption condition
        caption_embeds = None
        if self.transformer.cfg.use_caption_condition:
            if caption is not None:
                caption_embeds, negative_caption_embeds = self.encode_caption(
                    caption, device, num_meshes_per_prompt
                )
                if self.do_classifier_free_guidance:
                    caption_embeds = torch.cat(
                        [negative_caption_embeds, caption_embeds], dim=0
                    )
            else:
                uncond_caption_embeds = self.caption_encoder.empty_text_embeds.repeat(
                    num_meshes_per_prompt, 1, 1
                ).to(image_embeds)
                if self.do_classifier_free_guidance:
                    caption_embeds = torch.cat(
                        [uncond_caption_embeds, uncond_caption_embeds], dim=0
                    )

        # 4. Prepare timesteps
        timesteps, num_inference_steps = retrieve_timesteps(
            self.scheduler, num_inference_steps, device, timesteps
        )
        num_warmup_steps = max(
            len(timesteps) - num_inference_steps * self.scheduler.order, 0
        )
        self._num_timesteps = len(timesteps)

        # 5. Prepare latent variables
        num_latents = self.vae.cfg.num_latents
        num_channels_latents = self.transformer.cfg.input_channels
        latents = self.prepare_latents(
            batch_size * num_meshes_per_prompt,
            num_latents,
            num_channels_latents,
            image_embeds.dtype,
            device,
            generator,
            latents,
        )

        # 6. Denoising loop
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                # expand the latents if we are doing classifier free guidance
                latent_model_input = (
                    torch.cat([latents] * 2)
                    if self.do_classifier_free_guidance
                    else latents
                )
                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                timestep = t.expand(latent_model_input.shape[0])

                noise_pred = self.transformer(
                    latent_model_input,
                    timestep,
                    visual_condition=image_embeds,
                    label_condition=label_embeds,
                    caption_condition=caption_embeds,
                    return_dict=False,
                )[0]

                # perform guidance
                if self.do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_image = noise_pred.chunk(2)
                    noise_pred = noise_pred_uncond + self.guidance_scale * (
                        noise_pred_image - noise_pred_uncond
                    )

                if (i <= zero_steps) and use_zero_init:
                    noise_pred = noise_pred * 0.0

                # compute the previous noisy sample x_t -> x_t-1
                latents_dtype = latents.dtype
                latents = self.scheduler.step(
                    noise_pred, t, latents, return_dict=False
                )[0]

                if latents.dtype != latents_dtype:
                    if torch.backends.mps.is_available():
                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
                        latents = latents.to(latents_dtype)

                if i == len(timesteps) - 1 or (
                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
                ):
                    progress_bar.update()

        # 4. Post-processing
        if not output_type == "latent":
            if latents.dtype == torch.bfloat16:
                self.vae.to(torch.float16)
                latents = latents.to(torch.float16)
            mesh = self.vae.extract_geometry(
                self.vae.decode(latents),
                surface_extractor_type=surface_extractor_type,
                bounds=bounds,
                mc_level=mc_level,
                octree_resolution=octree_resolution,
                enable_pbar=False,
            )
            if output_type != "raw":
                mesh_list = []
                for i, cur_mesh in enumerate(mesh):
                    print(f"Generating mesh {i+1}/{num_meshes_per_prompt}")
                    if output_type == "trimesh":
                        import trimesh

                        cur_mesh = trimesh.Trimesh(
                            vertices=cur_mesh.verts.cpu().numpy(),
                            faces=cur_mesh.faces.cpu().numpy(),
                        )
                        cur_mesh.fix_normals()
                        cur_mesh.face_normals
                        cur_mesh.vertex_normals
                        cur_mesh.visual = trimesh.visual.TextureVisuals(
                            material=trimesh.visual.material.PBRMaterial(
                                baseColorFactor=(255, 255, 255),
                                main_color=(255, 255, 255),
                                metallicFactor=0.05,
                                roughnessFactor=1.0,
                            )
                        )
                        if do_remove_floater:
                            cur_mesh = remove_floater(cur_mesh)
                        if do_remove_degenerate_face:
                            cur_mesh = remove_degenerate_face(cur_mesh)
                        if do_reduce_face and max_facenum > 0:
                            cur_mesh = reduce_face(cur_mesh, max_facenum)
                        if do_shade_smooth:
                            cur_mesh = cur_mesh.smooth_shaded
                        mesh_list.append(cur_mesh)
                    elif output_type == "np":
                        if do_remove_floater:
                            print(
                                'remove floater is NOT used when output_type is "np". '
                            )
                        if do_remove_degenerate_face:
                            print(
                                'remove degenerate face is NOT used when output_type is "np". '
                            )
                        if do_reduce_face:
                            print(
                                'reduce floater is NOT used when output_type is "np". '
                            )
                        if do_shade_smooth:
                            print('shade smooth is NOT used when output_type is "np". ')
                        mesh_list.append(
                            [
                                cur_mesh[0].verts.cpu().numpy(),
                                cur_mesh[0].faces.cpu().numpy(),
                            ]
                        )
                mesh = mesh_list
            else:
                if do_remove_floater:
                    print('remove floater is NOT used when output_type is "raw". ')
                if do_remove_degenerate_face:
                    print(
                        'remove degenerate face is NOT used when output_type is "raw". '
                    )
                if do_reduce_face:
                    print('reduce floater is NOT used when output_type is "raw". ')

        else:
            mesh = latents

        if not return_dict:
            return tuple(image_pil), tuple(mesh)
        return Step1X3DGeometryPipelineOutput(image=image_pil, mesh=mesh)