app.py

import gradio as gr
import spaces
from gradio_litmodel3d import LitModel3D

import os
import shutil
os.environ['SPCONV_ALGO'] = 'native'
from typing import *
import torch
import numpy as np
import imageio
import uuid
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

# Constants
MAX_SEED = np.iinfo(np.int32).max
TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
os.makedirs(TMP_DIR, exist_ok=True)

# Session Management Functions
def start_session(req: gr.Request):
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    print(f'Creating user directory: {user_dir}')
    os.makedirs(user_dir, exist_ok=True)

def end_session(req: gr.Request):
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    print(f'Removing user directory: {user_dir}')
    shutil.rmtree(user_dir)

# Image Preprocessing Function
def preprocess_image(image: Image.Image) -> Image.Image:
    """
    Preprocess the input image.

    Args:
        image (Image.Image): The input image.

    Returns:
        Image.Image: The preprocessed image.
    """
    processed_image = pipeline.preprocess_image(image)
    return processed_image

# State Packing and Unpacking Functions
def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> 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(),
        },
        'trial_id': trial_id,
    }

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, state['trial_id']

# Seed Management Function
def get_seed(randomize_seed: bool, seed: int) -> int:
    """
    Get the random seed.

    Args:
        randomize_seed (bool): Whether to randomize the seed.
        seed (int): The provided seed value.

    Returns:
        int: The final seed to use.
    """
    return np.random.randint(0, MAX_SEED) if randomize_seed else seed

# Core 3D Generation Function
@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,
    req: gr.Request,
) -> Tuple[dict, str, str]:
    """
    Convert an image to a 3D model and generate a GLB file.

    Args:
        image (Image.Image): The input image.
        seed (int): The random seed.
        ss_guidance_strength (float): The guidance strength for sparse structure generation.
        ss_sampling_steps (int): The number of sampling steps for sparse structure generation.
        slat_guidance_strength (float): The guidance strength for structured latent generation.
        slat_sampling_steps (int): The number of sampling steps for structured latent generation.
        req (gr.Request): Gradio request object.

    Returns:
        Tuple[dict, str, str]: The state dictionary, path to the generated video, and path to the standard GLB file.
    """
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    outputs = pipeline.run(
        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,
        },
    )
    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
    trial_id = uuid.uuid4()
    video_path = os.path.join(user_dir, f"{trial_id}.mp4")
    imageio.mimsave(video_path, video, fps=15)
    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], trial_id)
    
    # Generate standard GLB file with default simplification and texture size
    glb = postprocessing_utils.to_glb(
        outputs['gaussian'][0],
        outputs['mesh'][0],
        simplify=0.95,          # Default simplification
        texture_size=1024,      # Default texture size
        verbose=False
    )
    glb_path = os.path.join(user_dir, f"{trial_id}.glb")
    glb.export(glb_path)
    
    torch.cuda.empty_cache()
    return state, video_path, glb_path

# Existing GLB Extraction Function
@spaces.GPU
def extract_glb(
    state: dict,
    mesh_simplify: float,
    texture_size: int,
    req: gr.Request,
) -> Tuple[dict, bytes]:
    """
    Extract a GLB file from the 3D model.

    Args:
        state (dict): The state of the generated 3D model.
        mesh_simplify (float): The mesh simplification factor.
        texture_size (int): The texture resolution.
        req (gr.Request): Gradio request object.

    Returns:
        Tuple[dict, bytes]: The model state and the GLB file bytes.
    """
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    gs, mesh, trial_id = unpack_state(state)
    glb = postprocessing_utils.to_glb(
        gs,
        mesh,
        simplify=mesh_simplify,
        texture_size=texture_size,
        verbose=False
    )
    glb_path = os.path.join(user_dir, f"{trial_id}.glb")
    glb.export(glb_path)
    # Read the GLB file as bytes
    with open(glb_path, "rb") as f:
        glb_bytes = f.read()
    torch.cuda.empty_cache()
    return state, glb_bytes

# New High-Quality GLB Extraction Function
@spaces.GPU
def extract_glb_high_quality(
    state: dict,
    req: gr.Request,
) -> Tuple[dict, bytes]:
    """
    Extract a high-quality GLB file from the 3D model without polygon reduction.

    Args:
        state (dict): The state of the generated 3D model.
        req (gr.Request): Gradio request object.

    Returns:
        Tuple[dict, bytes]: The model state and the high-quality GLB file bytes.
    """
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    gs, mesh, trial_id = unpack_state(state)
    # Set simplify to 0.0 to disable polygon reduction
    # Set texture_size to 2048 for maximum texture quality
    glb = postprocessing_utils.to_glb(
        gs,
        mesh,
        simplify=0.0,
        texture_size=2048,
        verbose=False
    )
    glb_path = os.path.join(user_dir, f"{trial_id}_high_quality.glb")
    glb.export(glb_path)
    # Read the GLB file as bytes
    with open(glb_path, "rb") as f:
        glb_bytes = f.read()
    torch.cuda.empty_cache()
    return state, glb_bytes

# Gradio Interface Definition
with gr.Blocks(delete_cache=(600, 600)) as demo:
    gr.Markdown("""
    ## Image to 3D Asset with [TRELLIS](https://trellis3d.github.io/)
    * **Generate:** Upload an image and click "Generate" to create a 3D asset. If the image has an alpha channel, it will be used as the mask. Otherwise, the background will be removed automatically.
    * **Extract GLB:** If the generated 3D asset is satisfactory, click "Extract GLB" to extract the GLB file based on your chosen settings and download it.
    * **Download High Quality GLB:** Click this button to download a high-quality GLB file without any polygon reduction and with maximum texture quality.
    * **Status:** View messages and feedback about your actions below.
    """)
    
    with gr.Row():
        with gr.Column():
            # Image Input
            image_prompt = gr.Image(
                label="Image Prompt",
                format="png",
                image_mode="RGBA",
                type="pil",
                height=300
            )
            
            # Generation Settings Accordion
            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,
                        500,
                        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,
                        500,
                        label="Sampling Steps",
                        value=12,
                        step=1
                    )

            # Generate Button
            generate_btn = gr.Button("Generate")
            
            # GLB Extraction Settings Accordion
            with gr.Accordion(label="GLB Extraction Settings", open=False):
                mesh_simplify = gr.Slider(
                    0.0,
                    0.98,
                    label="Simplify",
                    value=0.95,
                    step=0.01
                )
                texture_size = gr.Slider(
                    512,
                    2048,
                    label="Texture Size",
                    value=1024,
                    step=512
                )
            
            # Extract GLB Button
            extract_glb_btn = gr.Button("Extract GLB")
            
            # Download High Quality GLB Button
            download_glb_high_quality_btn = gr.Button("Download High Quality GLB")
        
        with gr.Column():
            # Video Output
            video_output = gr.Video(
                label="Generated 3D Asset",
                autoplay=True,
                loop=True,
                height=300
            )
            # 3D Model Display
            model_output = LitModel3D(
                label="Extracted GLB",
                exposure=20.0,
                height=300
            )
            # Download GLB Buttons
            download_glb = gr.DownloadButton(
                label="Download GLB"
            )
            download_high_quality_glb = gr.DownloadButton(
                label="Download High Quality GLB"
            )
            
            # Status Message
            status_message = gr.Textbox(
                label="Status",
                value="Awaiting your action...",
                interactive=False,
                lines=2
            )

    # State Variables
    output_buf = gr.State()
    glb_bytes_state = gr.State()  # For standard GLB
    glb_high_quality_bytes_state = gr.State()  # For high-quality GLB

    # Example Images
    with gr.Row():
        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,
        )

    # Event Handlers
    demo.load(start_session)
    demo.unload(end_session)
    
    # Image Upload Handler
    image_prompt.upload(
        preprocess_image,
        inputs=[image_prompt],
        outputs=[image_prompt],
    )

    # Generate Button Click Handler
    generate_btn.click(
        get_seed,
        inputs=[randomize_seed, seed],
        outputs=[seed],
    ).then(
        image_to_3d,
        inputs=[
            image_prompt,
            seed,
            ss_guidance_strength,
            ss_sampling_steps,
            slat_guidance_strength,
            slat_sampling_steps
        ],
        outputs=[output_buf, video_output, glb_path := gr.State()]
    ).then(
        lambda state, video, glb: "Generation successful! You can now extract and download GLB files.",
        inputs=[output_buf, video_output, glb_path],
        outputs=[status_message]
    )

    # Extract GLB Button Click Handler
    extract_glb_btn.click(
        extract_glb,
        inputs=[output_buf, mesh_simplify, texture_size],
        outputs=[output_buf, glb_bytes_state],
    ).then(
        # Map the GLB bytes to the DownloadButton with filename
        lambda state, glb_bytes: (glb_bytes, "model.glb"),
        inputs=[output_buf, glb_bytes_state],
        outputs=[download_glb],
    ).then(
        # Update status message
        lambda: "GLB extraction successful! Click 'Download GLB' to save your model.",
        inputs=None,
        outputs=[status_message]
    )

    # Download High Quality GLB Button Click Handler
    download_glb_high_quality_btn.click(
        extract_glb_high_quality,
        inputs=[output_buf],
        outputs=[output_buf, glb_high_quality_bytes_state],
    ).then(
        # Map the high-quality GLB bytes to the DownloadButton with filename
        lambda state, glb_bytes: (glb_bytes, "model_high_quality.glb"),
        inputs=[output_buf, glb_high_quality_bytes_state],
        outputs=[download_high_quality_glb],
    ).then(
        # Update status message
        lambda: "High-quality GLB extraction successful! Click 'Download High Quality GLB' to save your model.",
        inputs=None,
        outputs=[status_message]
    )

    # Handle Clearing of Video Output
    video_output.clear(
        lambda: "Video output cleared. Please generate a new 3D asset.",
        outputs=[status_message],
    )

    # Handle Clearing of Model Output
    model_output.clear(
        lambda: "Model output cleared. Please extract and download GLB files again.",
        outputs=[status_message],
    )

# Launch the Gradio app
if __name__ == "__main__":
    # Initialize the pipeline
    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()