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.
    """
    # Validate image
    if image is None:
        raise ValueError("No image provided.")
    if image.mode != "RGBA":
        image = image.convert("RGBA")
    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]:
    """
    Convert an image to a 3D model.

    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]: The state dictionary and the path to the generated video.
    """
    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)
    torch.cuda.empty_cache()
    return state, video_path

# Existing GLB Extraction Function
@spaces.GPU
def extract_glb(
    state: dict,
    mesh_simplify: float,
    texture_size: int,
    req: gr.Request,
) -> Tuple[str, str]:
    """
    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[str, str]: The path to the extracted GLB file.
    """
    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)
    torch.cuda.empty_cache()
    return glb_path, glb_path

# **Addition: High-Quality GLB Extraction Function**
@spaces.GPU
def extract_glb_high_quality(
    state: dict,
    req: gr.Request,
) -> Tuple[str, str]:
    """
    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[str, str]: The path to the high-quality GLB file.
    """
    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)
    torch.cuda.empty_cache()
    return glb_path, glb_path

# Gradio Interface Definition
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" 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.
    * If you find the generated 3D asset satisfactory, click "Extract GLB" to extract the GLB file and download it.
    * **New:** Click "Download High Quality GLB" to download the GLB file without any polygon reduction and with maximum texture quality.
    """)
    
    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)
            
            # Existing Extract GLB Button
            extract_glb_btn = gr.Button("Extract GLB", interactive=False)
            
            # **Addition: Download High Quality GLB Button**
            extract_glb_high_quality_btn = gr.Button("Download High Quality GLB", interactive=False)

        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)
            # Existing Download GLB Button
            download_glb = gr.DownloadButton(
                label="Download GLB",
                # Removed 'file_count' to prevent runtime error
            )
            
            # **Addition: Download High Quality GLB DownloadButton**
            download_high_quality_glb = gr.DownloadButton(
                label="Download High Quality GLB",
                # Removed 'file_count' to prevent runtime error
            )
    
    # State Variables
    output_buf = gr.State()
    glb_path_state = gr.State()  # For standard GLB
    glb_high_quality_path_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_prompt.upload(
        preprocess_image,
        inputs=[image_prompt],
        outputs=[image_prompt],
    )

    generate_btn.click(
        get_seed,
        inputs=[randomize_seed, seed],
        outputs=[seed],
        concurrency_limit=1  # Set concurrency limit for Generate
    ).then(
        image_to_3d,
        inputs=[image_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, gr.Request()],
        outputs=[output_buf, video_output],
        concurrency_limit=1  # Set concurrency limit for image_to_3d
    ).then(
        # Enable the Extract GLB and Download High Quality GLB buttons after generation
        lambda: (True, True),
        outputs=[extract_glb_btn, extract_glb_high_quality_btn],
    )

    video_output.clear(
        lambda: (False, False),
        outputs=[extract_glb_btn, extract_glb_high_quality_btn],
    )

    extract_glb_btn.click(
        extract_glb,
        inputs=[output_buf, mesh_simplify, texture_size, gr.Request()],
        outputs=[model_output, glb_path_state],
        concurrency_limit=1  # Set concurrency limit for extract_glb
    ).then(
        lambda glb_path: glb_path if glb_path else "",
        inputs=[glb_path_state],
        outputs=[download_glb],
    )

    # **Addition: High-Quality GLB Extraction and Download**
    extract_glb_high_quality_btn.click(
        extract_glb_high_quality,
        inputs=[output_buf, gr.Request()],
        outputs=[model_output, glb_high_quality_path_state],
        concurrency_limit=1  # Set concurrency limit for extract_glb_high_quality
    ).then(
        lambda glb_path: glb_path if glb_path else "",
        inputs=[glb_high_quality_path_state],
        outputs=[download_high_quality_glb],
    )

    model_output.clear(
        lambda: (gr.File.update(value=None), gr.File.update(value=None)),
        outputs=[download_glb, download_high_quality_glb],
    )

# 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 Exception as e:
        print(f"Preloading rembg failed: {e}")
    
    # Configure Gradio's queue without deprecated parameters
    demo.queue().launch()