app.py

import gradio as gr
import spaces
from gradio_litmodel3d import LitModel3D

import threading, time, shutil, os
os.environ['SPCONV_ALGO'] = 'native'
from typing import *
import torch
import numpy as np
import imageio
from easydict import EasyDict as edict
from PIL import Image
from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.representations import Gaussian, MeshExtractResult
from trellis.utils import render_utils, postprocessing_utils


MAX_SEED = np.iinfo(np.int32).max
# Custom save directory for GLB files
CUSTOM_SAVE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'outputs')
os.makedirs(CUSTOM_SAVE_DIR, exist_ok=True)



def start_session(req: gr.Request):
    user_dir = os.path.join(CUSTOM_SAVE_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)
    
def end_session(req: gr.Request):
    """
    Safe session cleanup for Hugging Face Spaces.
    - Runs asynchronously (won’t block Gradio)
    - Only touches your custom /outputs/ folder
    - Waits a few seconds to avoid race conditions
    """
    def delayed_cleanup():
        time.sleep(5)  # allow Gradio to finish file reads
        user_dir = os.path.join(CUSTOM_SAVE_DIR, str(req.session_hash))
        try:
            # safety check: never delete anything outside outputs/
            if not user_dir.startswith(CUSTOM_SAVE_DIR):
                print(f"[Session Cleanup Warning] Unsafe path skipped: {user_dir}")
                return
    
            if os.path.exists(user_dir):
                shutil.rmtree(user_dir, ignore_errors=True)
                print(f"[Session Cleanup] Deleted {user_dir}")
            else:
                print(f"[Session Cleanup] Directory already gone: {user_dir}")
    
        except Exception as e:
            print(f"[Session Cleanup Warning] {e}")
    
    # run cleanup in background so Gradio can unload cleanly
    threading.Thread(target=delayed_cleanup, daemon=True).start()



def preprocess_image(image: Image.Image) -> Image.Image:
    """
    Preprocess the input image for 3D generation.
    
    This function is called when a user uploads an image or selects an example.
    It applies background removal and other preprocessing steps necessary for
    optimal 3D model generation.

    Args:
        image (Image.Image): The input image from the user

    Returns:
        Image.Image: The preprocessed image ready for 3D generation
    """
    processed_image = pipeline.preprocess_image(image)
    return processed_image


def preprocess_images(images: List[Tuple[Image.Image, str]]) -> List[Image.Image]:
    """
    Preprocess a list of input images for multi-image 3D generation.
    
    This function is called when users upload multiple images in the gallery.
    It processes each image to prepare them for the multi-image 3D generation pipeline.
    
    Args:
        images (List[Tuple[Image.Image, str]]): The input images from the gallery
        
    Returns:
        List[Image.Image]: The preprocessed images ready for 3D generation
    """
    images = [image[0] for image in images]
    processed_images = [pipeline.preprocess_image(image) for image in images]
    return processed_images


def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
    return {
        'gaussian': {
            **gs.init_params,
            '_xyz': gs._xyz.cpu().numpy(),
            '_features_dc': gs._features_dc.cpu().numpy(),
            '_scaling': gs._scaling.cpu().numpy(),
            '_rotation': gs._rotation.cpu().numpy(),
            '_opacity': gs._opacity.cpu().numpy(),
        },
        'mesh': {
            'vertices': mesh.vertices.cpu().numpy(),
            'faces': mesh.faces.cpu().numpy(),
        },
    }
    
    
def unpack_state(state: dict) -> 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) -> int:
    """
    Get the random seed for generation.
    
    This function is called by the generate button to determine whether to use
    a random seed or the user-specified seed value.
    
    Args:
        randomize_seed (bool): Whether to generate a random seed
        seed (int): The user-specified seed value
        
    Returns:
        int: The seed to use for generation
    """
    return np.random.randint(0, MAX_SEED) if randomize_seed else seed


@spaces.GPU(duration=120)
def generate_and_extract_glb(
    image: Image.Image,
    multiimages: List[Tuple[Image.Image, str]],
    seed: int,
    ss_guidance_strength: float,
    ss_sampling_steps: int,
    slat_guidance_strength: float,
    slat_sampling_steps: int,
    multiimage_algo: Literal["multidiffusion", "stochastic"],
    mesh_simplify: float,
    texture_size: int,
    req: gr.Request,
) -> str:
    """
    Generate the 3D model and save only the GLB file to a custom folder.
    The GLB filename matches the input image filename (or first image if multi-image).
    """
    # Derive base filename from the first image in the batch
    if image:
        base_filename = getattr(image, "filename", None) or "generated_model"
    else:
        # Use the first uploaded image filename if available
        base_filename = getattr(multiimages[0][0], "filename", None) or "generated_model"

    # Strip directories and extension to get a clean name
    base_name = os.path.splitext(os.path.basename(base_filename))[0]
    glb_filename = f"{base_name}.glb"
    glb_path = os.path.join(CUSTOM_SAVE_DIR, glb_filename)

    # Generate 3D model
    if image:
        outputs = pipeline.run(
            image,
            seed=seed,
            formats=["gaussian", "mesh"],
            preprocess_image=True,
            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,
            },
        )
    else:
        outputs = pipeline.run_multi_image(
            [img[0] for img in multiimages],
            seed=seed,
            formats=["gaussian", "mesh"],
            preprocess_image=True,
            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,
            },
            mode=multiimage_algo,
        )

    # Export GLB file to custom folder
    gs = outputs['gaussian'][0]
    mesh = outputs['mesh'][0]
    glb = postprocessing_utils.to_glb(
        gs,
        mesh,
        simplify=mesh_simplify,
        texture_size=texture_size,
        verbose=False
    )
    glb.export(glb_path)

    torch.cuda.empty_cache()
    print(f"[Saved GLB] {glb_path}")
    return glb_path



@spaces.GPU
def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
    """
    Extract a Gaussian splatting file from the generated 3D model.
    
    This function is called when the user clicks "Extract Gaussian" button.
    It converts the 3D model state into a .ply file format containing
    Gaussian splatting data for advanced 3D applications.

    Args:
        state (dict): The state of the generated 3D model containing Gaussian data
        req (gr.Request): Gradio request object for session management

    Returns:
        Tuple[str, str]: Paths to the extracted Gaussian file (for display and download)
    """
    user_dir = os.path.join(CUSTOM_SAVE_DIR, str(req.session_hash))
    gs, _ = unpack_state(state)
    gaussian_path = os.path.join(user_dir, 'sample.ply')
    gs.save_ply(gaussian_path)
    torch.cuda.empty_cache()
    return gaussian_path, gaussian_path


def prepare_multi_example() -> List[Image.Image]:
    multi_case = list(set([i.split('_')[0] for i in os.listdir("assets/example_multi_image")]))
    images = []
    for case in multi_case:
        _images = []
        for i in range(1, 4):
            img = Image.open(f'assets/example_multi_image/{case}_{i}.png')
            W, H = img.size
            img = img.resize((int(W / H * 512), 512))
            _images.append(np.array(img))
        images.append(Image.fromarray(np.concatenate(_images, axis=1)))
    return images


def split_image(image: Image.Image) -> List[Image.Image]:
    """
    Split a multi-view image into separate view images.
    
    This function is called when users select multi-image examples that contain
    multiple views in a single concatenated image. It automatically splits them
    based on alpha channel boundaries and preprocesses each view.
    
    Args:
        image (Image.Image): A concatenated image containing multiple views
        
    Returns:
        List[Image.Image]: List of individual preprocessed view images
    """
    image = np.array(image)
    alpha = image[..., 3]
    alpha = np.any(alpha>0, axis=0)
    start_pos = np.where(~alpha[:-1] & alpha[1:])[0].tolist()
    end_pos = np.where(alpha[:-1] & ~alpha[1:])[0].tolist()
    images = []
    for s, e in zip(start_pos, end_pos):
        images.append(Image.fromarray(image[:, s:e+1]))
    return [preprocess_image(image) for image in images]


with gr.Blocks(delete_cache=(600, 600)) as demo:
    gr.Markdown("""
    ## Image to 3D Asset with [TRELLIS](https://trellis3d.github.io/)
    * Upload an image and click "Generate & Extract GLB" to create a 3D asset and automatically extract the GLB file.
    * If you want the Gaussian file as well, click "Extract Gaussian" after generation.
    * If the image has alpha channel, it will be used as the mask. Otherwise, we use `rembg` to remove the background.
    
    ✨New: 1) Experimental multi-image support. 2) Gaussian file extraction.
    """)
    
    with gr.Row():
        with gr.Column():
            with gr.Tabs() as input_tabs:
                with gr.Tab(label="Single Image", id=0) as single_image_input_tab:
                    image_prompt = gr.Image(label="Image Prompt", format="png", image_mode="RGBA", type="pil", height=300)
                with gr.Tab(label="Multiple Images", id=1) as multiimage_input_tab:
                    multiimage_prompt = gr.Gallery(label="Image Prompt", format="png", type="pil", height=300, columns=3)
                    gr.Markdown("""
                        Input different views of the object in separate images. 
                        
                        *NOTE: this is an experimental algorithm without training a specialized model. It may not produce the best results for all images, especially those having different poses or inconsistent details.*
                    """)
            
            with gr.Accordion(label="Generation Settings", open=False):
                seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                gr.Markdown("Stage 1: Sparse 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("Stage 2: Structured Latent 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)
                multiimage_algo = gr.Radio(["stochastic", "multidiffusion"], label="Multi-image Algorithm", value="stochastic")
            
            with gr.Accordion(label="GLB Extraction Settings", open=False):
                mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
                texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)

            generate_btn = gr.Button("Generate & Extract GLB", variant="primary")
            extract_gs_btn = gr.Button("Extract Gaussian", interactive=False)
            gr.Markdown("""
                        *NOTE: Gaussian file can be very large (~50MB), it will take a while to display and download.*
                        """)

        with gr.Column():
            video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
            model_output = LitModel3D(label="Extracted GLB/Gaussian", exposure=10.0, height=300)
            
            with gr.Row():
                download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
                download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)  
    
    is_multiimage = gr.State(False)
    output_buf = gr.State()

    # Example images at the bottom of the page
    with gr.Row() as single_image_example:
        examples = gr.Examples(
            examples=[
                f'assets/example_image/{image}'
                for image in os.listdir("assets/example_image")
            ],
            inputs=[image_prompt],
            fn=preprocess_image,
            outputs=[image_prompt],
            run_on_click=True,
            examples_per_page=64,
        )
    with gr.Row(visible=False) as multiimage_example:
        examples_multi = gr.Examples(
            examples=prepare_multi_example(),
            inputs=[image_prompt],
            fn=split_image,
            outputs=[multiimage_prompt],
            run_on_click=True,
            examples_per_page=8,
        )

    # Handlers
    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_multiimage, single_image_example, multiimage_example]
    )
    multiimage_input_tab.select(
        lambda: tuple([True, gr.Row.update(visible=False), gr.Row.update(visible=True)]),
        outputs=[is_multiimage, single_image_example, multiimage_example]
    )
    
    image_prompt.upload(
        preprocess_image,
        inputs=[image_prompt],
        outputs=[image_prompt],
    )
    multiimage_prompt.upload(
        preprocess_images,
        inputs=[multiimage_prompt],
        outputs=[multiimage_prompt],
    )

    generate_btn.click(
        get_seed,
        inputs=[randomize_seed, seed],
        outputs=[seed],
    ).then(
        generate_and_extract_glb,
        inputs=[image_prompt, multiimage_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, multiimage_algo, mesh_simplify, texture_size],
        outputs=[download_glb],
    ).then(
        lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]),
        outputs=[extract_gs_btn, download_glb],
    )

    video_output.clear(
        lambda: tuple([gr.Button(interactive=False), gr.Button(interactive=False), gr.Button(interactive=False)]),
        outputs=[extract_gs_btn, download_glb, download_gs],
    )
    
    extract_gs_btn.click(
        extract_gaussian,
        inputs=[output_buf],
        outputs=[model_output, download_gs],
    ).then(
        lambda: gr.Button(interactive=True),
        outputs=[download_gs],
    )

    model_output.clear(
        lambda: tuple([gr.Button(interactive=False), gr.Button(interactive=False)]),
        outputs=[download_glb, download_gs],
    )
    

# Launch the Gradio app
if __name__ == "__main__":
    pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
    pipeline.cuda()
    try:
        pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
    except:
        pass
    demo.launch(
        mcp_server=True,
        show_error=True   # ✅ Enable verbose backend error reporting
    )