app.py

# Version: 1.1.2 - Removed torch_dtype from from_pretrained call
# Applied:
#  - Removed unsupported inputs/outputs kwargs on demo.load/unload
#  - Converted NumPy arrays to lists in pack_state for JSON safety
#  - Fixed indentation in Blocks event-handlers
#  - Verified clear() callbacks use only callback + outputs
#  - Removed `torch_dtype` arg from TrellisTextTo3DPipeline.from_pretrained
#  - Bumped version, added comments at change sites

import gradio as gr
import spaces
import os
import shutil
os.environ['TOKENIZERS_PARALLELISM'] = 'true'
os.environ['SPCONV_ALGO'] = 'native'

from typing import *
import torch
import numpy as np
import imageio
from easydict import EasyDict as edict
from trellis.pipelines import TrellisTextTo3DPipeline
from trellis.representations import Gaussian, MeshExtractResult
from trellis.utils import render_utils, postprocessing_utils
import traceback
import sys

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)


def start_session(req: gr.Request):
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)
    print(f"Started session, created directory: {user_dir}")


def end_session(req: gr.Request):
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    if os.path.exists(user_dir):
        try:
            shutil.rmtree(user_dir)
            print(f"Ended session, removed directory: {user_dir}")
        except OSError as e:
            print(f"Error removing tmp directory {user_dir}: {e.strerror}", file=sys.stderr)
    else:
        print(f"Ended session, directory already removed: {user_dir}")


def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
    """Packs Gaussian and Mesh data into a JSON-serializable dictionary."""
    packed_data = {
        'gaussian': {
            **{k: v for k, v in gs.init_params.items()},
            # FIX: convert arrays to lists for JSON
            '_xyz': gs._xyz.detach().cpu().numpy().tolist(),
            '_features_dc': gs._features_dc.detach().cpu().numpy().tolist(),
            '_scaling': gs._scaling.detach().cpu().numpy().tolist(),
            '_rotation': gs._rotation.detach().cpu().numpy().tolist(),
            '_opacity': gs._opacity.detach().cpu().numpy().tolist(),
        },
        'mesh': {
            'vertices': mesh.vertices.detach().cpu().numpy().tolist(),
            'faces': mesh.faces.detach().cpu().numpy().tolist(),
        },
    }
    return packed_data


def unpack_state(state_dict: dict) -> Tuple[Gaussian, edict]:
    print("[unpack_state] Unpacking state from dictionary... ")
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    gauss_data = state_dict['gaussian']
    mesh_data = state_dict['mesh']
    gs = Gaussian(
        aabb=gauss_data.get('aabb'),
        sh_degree=gauss_data.get('sh_degree'),
        mininum_kernel_size=gauss_data.get('mininum_kernel_size'),
        scaling_bias=gauss_data.get('scaling_bias'),
        opacity_bias=gauss_data.get('opacity_bias'),
        scaling_activation=gauss_data.get('scaling_activation'),
    )
    gs._xyz = torch.tensor(np.array(gauss_data['_xyz']), device=device, dtype=torch.float32)
    gs._features_dc = torch.tensor(np.array(gauss_data['_features_dc']), device=device, dtype=torch.float32)
    gs._scaling = torch.tensor(np.array(gauss_data['_scaling']), device=device, dtype=torch.float32)
    gs._rotation = torch.tensor(np.array(gauss_data['_rotation']), device=device, dtype=torch.float32)
    gs._opacity = torch.tensor(np.array(gauss_data['_opacity']), device=device, dtype=torch.float32)
    mesh = edict(
        vertices=torch.tensor(np.array(mesh_data['vertices']), device=device, dtype=torch.float32),
        faces=torch.tensor(np.array(mesh_data['faces']), device=device, dtype=torch.int64),
    )
    return gs, mesh


def get_seed(randomize_seed: bool, seed: int) -> int:
    new_seed = np.random.randint(0, MAX_SEED) if randomize_seed else seed
    return int(new_seed)

@spaces.GPU
def text_to_3d(
    prompt: str,
    seed: int,
    ss_guidance_strength: float,
    ss_sampling_steps: int,
    slat_guidance_strength: float,
    slat_sampling_steps: int,
    req: gr.Request,
) -> Tuple[dict, str]:
    outputs = pipeline.run(
        prompt,
        seed=seed,
        formats=["gaussian", "mesh"],
        sparse_structure_sampler_params={"steps": int(ss_sampling_steps), "cfg_strength": float(ss_guidance_strength)},
        slat_sampler_params={"steps": int(slat_sampling_steps), "cfg_strength": float(slat_guidance_strength)},
    )
    state_dict = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
    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_combined = [np.concatenate([v.astype(np.uint8), vg.astype(np.uint8)], axis=1) for v, vg in zip(video, video_geo)]
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)
    video_path = os.path.join(user_dir, 'sample.mp4')
    imageio.mimsave(video_path, video_combined, fps=15, quality=8)
    if torch.cuda.is_available(): torch.cuda.empty_cache()
    return state_dict, video_path

@spaces.GPU(duration=120)
def extract_glb(
    state_dict: dict,
    mesh_simplify: float,
    texture_size: int,
    req: gr.Request,
) -> Tuple[str, str]:
    gs, mesh = unpack_state(state_dict)
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)
    glb = postprocessing_utils.to_glb(gs, mesh, simplify=float(mesh_simplify), texture_size=int(texture_size), verbose=True)
    glb_path = os.path.join(user_dir, 'sample.glb')
    glb.export(glb_path)
    if torch.cuda.is_available(): torch.cuda.empty_cache()
    return glb_path, glb_path

@spaces.GPU
def extract_gaussian(
    state_dict: dict,
    req: gr.Request
) -> Tuple[str, str]:
    gs, _ = unpack_state(state_dict)
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)
    gaussian_path = os.path.join(user_dir, 'sample.ply')
    gs.save_ply(gaussian_path)
    if torch.cuda.is_available(): torch.cuda.empty_cache()
    return gaussian_path, gaussian_path

# --- Gradio UI Definition ---
with gr.Blocks(delete_cache=(600, 600), title="TRELLIS Text-to-3D") as demo:
    gr.Markdown("""
    # Text to 3D Asset with [TRELLIS](https://trellis3d.github.io/)
    """)

    # State buffer
    output_buf = gr.State()

    with gr.Row():
        with gr.Column(scale=1):
            text_prompt = gr.Textbox(label="Text Prompt", lines=5)
            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("---\n**Stage 1**")
                ss_guidance_strength = gr.Slider(0.0, 15.0, label="Guidance Strength", value=7.5, step=0.1)
                ss_sampling_steps = gr.Slider(10, 50, label="Sampling Steps", value=25, step=1)
                gr.Markdown("---\n**Stage 2**")
                slat_guidance_strength = gr.Slider(0.0, 15.0, label="Guidance Strength", value=7.5, step=0.1)
                slat_sampling_steps = gr.Slider(10, 50, label="Sampling Steps", value=25, step=1)
            generate_btn = gr.Button("Generate 3D Preview", variant="primary")
            with gr.Accordion(label="GLB Extraction Settings", open=True):
                mesh_simplify = gr.Slider(0.9, 0.99, label="Simplify Factor", value=0.95, step=0.01)
                texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
            extract_glb_btn = gr.Button("Extract GLB", interactive=False)
            extract_gs_btn = gr.Button("Extract Gaussian (PLY)", interactive=False)
            download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
            download_gs = gr.DownloadButton(label="Download Gaussian (PLY)", interactive=False)
        with gr.Column(scale=1):
            video_output = gr.Video(label="3D Preview", autoplay=True, loop=True)
            model_output = gr.Model3D(label="Extracted Model Preview")

    # --- Event handlers ---
    demo.load(start_session)  # FIX: remove inputs/outputs kwargs
    demo.unload(end_session)  # FIX: remove inputs/outputs kwargs

    # Align indentation to one level under Blocks
    generate_event = generate_btn.click(
        get_seed,
        inputs=[randomize_seed, seed],
        outputs=[seed],
    ).then(
        text_to_3d,
        inputs=[text_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
        outputs=[output_buf, video_output],
    ).then(
        lambda: (extract_glb_btn.update(interactive=True), extract_gs_btn.update(interactive=True)),
        outputs=[extract_glb_btn, extract_gs_btn],
    )

    extract_glb_event = extract_glb_btn.click(
        extract_glb,
        inputs=[output_buf, mesh_simplify, texture_size],
        outputs=[model_output, download_glb],
    ).then(
        lambda: download_glb.update(interactive=True),
        outputs=[download_glb],
    )

    extract_gs_event = extract_gs_btn.click(
        extract_gaussian,
        inputs=[output_buf],
        outputs=[model_output, download_gs],
    ).then(
        lambda: download_gaussian.update(interactive=True),
        outputs=[download_gs],
    )

    # Clear callbacks
    model_output.clear(
        lambda: (download_glb.update(interactive=False), download_gs.update(interactive=False)),
        outputs=[download_glb, download_gs],
    )
    video_output.clear(
        lambda: (extract_glb_btn.update(interactive=False), extract_gs_btn.update(interactive=False), download_glb.update(interactive=False), download_gs.update(interactive=False)),
        outputs=[extract_glb_btn, extract_gs_btn, download_glb, download_gs],
    )

if __name__ == "__main__":
    # Removed torch_dtype argument to match current API
    pipeline = TrellisTextTo3DPipeline.from_pretrained(
        "JeffreyXiang/TRELLIS-text-xlarge"
    )
    if torch.cuda.is_available(): pipeline = pipeline.to("cuda")
    demo.queue().launch(debug=True)