app.py

# Version: 1.1.3 - API State Fix + DEBUG (Video Disabled) + unload() Fix (2025-05-04)
# Changes:
# - FIXED TypeError in demo.unload() by removing incorrect 'inputs'/'outputs' arguments.
# - ENSURED `import spaces` is present for the @spaces.GPU decorator.
# - TEMPORARY DEBUGGING STEP: Commented out video rendering in `text_to_3d`
#   and return None for video_path to isolate the "Session not found" error.
# - Modified `text_to_3d` to explicitly return the serializable `state_dict` from `pack_state`.
# - Modified `extract_glb`/`extract_gaussian` to accept `state_dict: dict`.
# - Kept Gradio UI bindings using `output_buf`.
# - Added minor safety checks and logging.

import gradio as gr
import spaces  # <<<--- ENSURE THIS IMPORT IS PRESENT

import os
import shutil
os.environ['TOKENIZERS_PARALLELISM'] = 'true'
os.environ['SPCONV_ALGO'] = 'native' # Direct set as per original

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
# Use standard /tmp directory which is usually available in container environments
TMP_DIR = '/tmp/gradio_sessions'
print(f"Using temporary directory: {TMP_DIR}")
# Ensure the base temp directory exists
try:
    os.makedirs(TMP_DIR, exist_ok=True)
except OSError as e:
    print(f"Warning: Could not create base temp directory {TMP_DIR}: {e}", file=sys.stderr)
    # Potentially fall back or exit if temp dir is critical
    TMP_DIR = '.' # Fallback to current directory (less ideal)
    print(f"Warning: Falling back to use current directory for temp files: {os.path.abspath(TMP_DIR)}")


def start_session(req: gr.Request):
    """Creates a temporary directory for the user session."""
    user_dir = None # Initialize
    try:
        session_hash = req.session_hash
        if not session_hash:
             session_hash = f"no_session_{np.random.randint(10000, 99999)}"
             print(f"Warning: No session_hash in request, using temporary ID: {session_hash}")

        # Ensure TMP_DIR exists before joining path
        if not os.path.exists(TMP_DIR):
             os.makedirs(TMP_DIR, exist_ok=True)

        user_dir = os.path.join(TMP_DIR, str(session_hash))
        os.makedirs(user_dir, exist_ok=True)
        print(f"Started session, ensured directory exists: {user_dir}")
    except Exception as e:
        print(f"Error in start_session creating directory '{user_dir}': {e}", file=sys.stderr)
        traceback.print_exc()


def end_session(req: gr.Request):
    """Removes the temporary directory for the user session."""
    user_dir = None # Initialize
    try:
        session_hash = req.session_hash
        if not session_hash:
            print("Warning: No session_hash in end_session request, cannot clean up.")
            return

        user_dir = os.path.join(TMP_DIR, str(session_hash))
        if os.path.exists(user_dir) and os.path.isdir(user_dir): # Extra check if it's a directory
            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 not found or not a directory: {user_dir}")
    except Exception as e:
        print(f"Error in end_session cleaning directory '{user_dir}': {e}", file=sys.stderr)


def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
    """Packs Gaussian and Mesh data into a serializable dictionary."""
    print("[pack_state] Packing state to dictionary...")
    try:
        packed_data = {
            'gaussian': {
                **{k: v for k, v in gs.init_params.items()},
                '_xyz': gs._xyz.detach().cpu().numpy(),
                '_features_dc': gs._features_dc.detach().cpu().numpy(),
                '_scaling': gs._scaling.detach().cpu().numpy(),
                '_rotation': gs._rotation.detach().cpu().numpy(),
                '_opacity': gs._opacity.detach().cpu().numpy(),
            },
            'mesh': {
                'vertices': mesh.vertices.detach().cpu().numpy(),
                'faces': mesh.faces.detach().cpu().numpy(),
            },
        }
        print(f"[pack_state] Dictionary created. Keys: {list(packed_data.keys())}, Gaussian points: {len(packed_data['gaussian']['_xyz'])}, Mesh vertices: {len(packed_data['mesh']['vertices'])}")
        return packed_data
    except Exception as e:
        print(f"Error during pack_state: {e}", file=sys.stderr)
        traceback.print_exc()
        raise


def unpack_state(state_dict: dict) -> Tuple[Gaussian, edict]:
    """Unpacks Gaussian and Mesh data from a dictionary."""
    print("[unpack_state] Unpacking state from dictionary...")
    try:
        if not isinstance(state_dict, dict) or 'gaussian' not in state_dict or 'mesh' not in state_dict:
            raise ValueError("Invalid state_dict structure passed to unpack_state.")

        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print(f"[unpack_state] Using device: {device}")

        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(gauss_data['_xyz'], device=device, dtype=torch.float32)
        gs._features_dc = torch.tensor(gauss_data['_features_dc'], device=device, dtype=torch.float32)
        gs._scaling = torch.tensor(gauss_data['_scaling'], device=device, dtype=torch.float32)
        gs._rotation = torch.tensor(gauss_data['_rotation'], device=device, dtype=torch.float32)
        gs._opacity = torch.tensor(gauss_data['_opacity'], device=device, dtype=torch.float32)
        print(f"[unpack_state] Gaussian unpacked. Points: {gs.get_xyz.shape[0]}")

        mesh = edict(
            vertices=torch.tensor(mesh_data['vertices'], device=device, dtype=torch.float32),
            faces=torch.tensor(mesh_data['faces'], device=device, dtype=torch.int64),
        )
        print(f"[unpack_state] Mesh unpacked. Vertices: {mesh.vertices.shape[0]}, Faces: {mesh.faces.shape[0]}")

        return gs, mesh
    except Exception as e:
        print(f"Error during unpack_state: {e}", file=sys.stderr)
        traceback.print_exc()
        raise


def get_seed(randomize_seed: bool, seed: int) -> int:
    """Gets a seed value, randomizing if requested."""
    new_seed = np.random.randint(0, MAX_SEED) if randomize_seed else seed
    print(f"[get_seed] Randomize: {randomize_seed}, Input Seed: {seed}, Output Seed: {new_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, Optional[str]]:
    """
    Generates a 3D model (Gaussian and Mesh) from text and returns a
    serializable state dictionary and potentially a video preview path.
    >>> TEMPORARILY DISABLED VIDEO RENDERING FOR DEBUGGING <<<
    """
    print(f"[text_to_3d - DEBUG MODE] Received prompt: '{prompt}', Seed: {seed}")
    user_dir = None # Initialize
    state_dict = None # Initialize
    try:
        session_hash = req.session_hash
        if not session_hash:
            session_hash = f"no_session_{np.random.randint(10000, 99999)}"
            print(f"Warning: No session_hash in text_to_3d request, using temporary ID: {session_hash}")

        # Ensure user directory exists
        user_dir = os.path.join(TMP_DIR, str(session_hash))
        os.makedirs(user_dir, exist_ok=True)
        print(f"[text_to_3d - DEBUG MODE] User directory: {user_dir}")

        # --- Generation Pipeline ---
        print("[text_to_3d - DEBUG MODE] Running Trellis pipeline...")
        outputs = pipeline.run(
            prompt=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),
            },
        )
        print("[text_to_3d - DEBUG MODE] Pipeline run completed.")

        # --- Create Serializable State Dictionary ---
        state_dict = pack_state(outputs['gaussian'][0], outputs['mesh'][0])

    except Exception as e:
        print(f"❌ [text_to_3d - DEBUG MODE] Error during generation or packing: {e}", file=sys.stderr)
        traceback.print_exc()
        # Raise a Gradio error to send failure message back to client if possible
        raise gr.Error(f"Core generation failed: {e}")

    # --- Render Video Preview (TEMPORARILY DISABLED FOR DEBUGGING) ---
    video_path = None
    print("[text_to_3d - DEBUG MODE] Skipping video rendering.")
    # --- Original Video Code Block (Keep commented) ---
    # ... (video code commented out) ...

    # --- Cleanup and Return ---
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        print("[text_to_3d - DEBUG MODE] Cleared CUDA cache.")

    # --- Return Serializable Dictionary and None Video Path ---
    print("[text_to_3d - DEBUG MODE] Returning state dictionary and None video path.")
    if state_dict is None:
         # This case should ideally be caught by the exception handling above
         print("Error: state_dict is None before return, generation likely failed.", file=sys.stderr)
         raise gr.Error("State dictionary creation failed.")
    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]:
    """
    Extracts a GLB file from the provided 3D model state dictionary.
    """
    print(f"[extract_glb] Received request. Simplify: {mesh_simplify}, Texture Size: {texture_size}")
    user_dir = None # Initialize
    glb_path = None # Initialize
    try:
        session_hash = req.session_hash
        if not session_hash:
            session_hash = f"no_session_{np.random.randint(10000, 99999)}"
            print(f"Warning: No session_hash in extract_glb request, using temporary ID: {session_hash}")

        if not isinstance(state_dict, dict):
            print("❌ [extract_glb] Error: Invalid state_dict received (not a dictionary).")
            raise gr.Error("Invalid state data received. Please generate the model first.")

        user_dir = os.path.join(TMP_DIR, str(session_hash))
        os.makedirs(user_dir, exist_ok=True)
        print(f"[extract_glb] User directory: {user_dir}")

        # --- Unpack state from the dictionary ---
        gs, mesh = unpack_state(state_dict)

        # --- Postprocessing and Export ---
        print("[extract_glb] Converting to GLB...")
        simplify_factor = float(mesh_simplify)
        tex_size = int(texture_size)
        glb = postprocessing_utils.to_glb(gs, mesh, simplify=simplify_factor, texture_size=tex_size, verbose=True)
        glb_path = os.path.join(user_dir, 'sample.glb')
        print(f"[extract_glb] Exporting GLB to: {glb_path}")
        glb.export(glb_path)
        print("[extract_glb] GLB exported successfully.")

    except Exception as e:
        print(f"❌ [extract_glb] Error during GLB extraction: {e}", file=sys.stderr)
        traceback.print_exc()
        raise gr.Error(f"Failed to extract GLB: {e}") # Propagate error

    # --- Cleanup and Return ---
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        print("[extract_glb] Cleared CUDA cache.")

    print("[extract_glb] Returning GLB path.")
    if glb_path is None:
         print("Error: glb_path is None before return, extraction likely failed.", file=sys.stderr)
         raise gr.Error("GLB path generation failed.")
    return glb_path, glb_path


@spaces.GPU
def extract_gaussian(
    state_dict: dict,
    req: gr.Request
) -> Tuple[str, str]:
    """
    Extracts a PLY (Gaussian) file from the provided 3D model state dictionary.
    """
    print("[extract_gaussian] Received request.")
    user_dir = None # Initialize
    gaussian_path = None # Initialize
    try:
        session_hash = req.session_hash
        if not session_hash:
            session_hash = f"no_session_{np.random.randint(10000, 99999)}"
            print(f"Warning: No session_hash in extract_gaussian request, using temporary ID: {session_hash}")

        if not isinstance(state_dict, dict):
            print("❌ [extract_gaussian] Error: Invalid state_dict received (not a dictionary).")
            raise gr.Error("Invalid state data received. Please generate the model first.")

        user_dir = os.path.join(TMP_DIR, str(session_hash))
        os.makedirs(user_dir, exist_ok=True)
        print(f"[extract_gaussian] User directory: {user_dir}")

        # --- Unpack state from the dictionary ---
        gs, _ = unpack_state(state_dict)

        # --- Export PLY ---
        gaussian_path = os.path.join(user_dir, 'sample.ply')
        print(f"[extract_gaussian] Saving PLY to: {gaussian_path}")
        gs.save_ply(gaussian_path)
        print("[extract_gaussian] PLY saved successfully.")

    except Exception as e:
        print(f"❌ [extract_gaussian] Error during Gaussian extraction: {e}", file=sys.stderr)
        traceback.print_exc()
        raise gr.Error(f"Failed to extract Gaussian PLY: {e}") # Propagate error

    # --- Cleanup and Return ---
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        print("[extract_gaussian] Cleared CUDA cache.")

    print("[extract_gaussian] Returning PLY path.")
    if gaussian_path is None:
         print("Error: gaussian_path is None before return, extraction likely failed.", file=sys.stderr)
         raise gr.Error("Gaussian PLY path generation failed.")
    return gaussian_path, gaussian_path


# --- Gradio UI Definition ---
print("Setting up Gradio Blocks interface...")
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/)
    * Type a text prompt and click "Generate" to create a 3D asset preview.
    * Adjust extraction settings if desired.
    * Click "Extract GLB" or "Extract Gaussian" to get the downloadable 3D file.
    *(Note: Video preview is temporarily disabled for debugging)*
    """)

    # --- State Buffer ---
    output_buf = gr.State()

    with gr.Row():
        with gr.Column(scale=1): # Input column
            text_prompt = gr.Textbox(label="Text Prompt", lines=5, placeholder="e.g., a cute red dragon")

            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: Sparse Structure Generation**")
                with gr.Row():
                    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: Structured Latent Generation**")
                with gr.Row():
                    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, info="Higher value = less simplification (more polys)")
                texture_size = gr.Slider(512, 2048, label="Texture Size (pixels)", value=1024, step=512, info="Size of the generated texture map")

            with gr.Row():
                extract_glb_btn = gr.Button("Extract GLB", interactive=False)
                extract_gs_btn = gr.Button("Extract Gaussian (PLY)", interactive=False)
            gr.Markdown("""
                        *NOTE: Gaussian file (.ply) can be very large (~50MB+) and may take time to process/download.*
                        """)

        with gr.Column(scale=1): # Output column
            # Video component remains for layout but won't show anything in this debug version
            video_output = gr.Video(label="Generated 3D Preview (DISABLED FOR DEBUG)", autoplay=False, loop=False, value=None, height=350)
            model_output = gr.Model3D(label="Extracted Model Preview", height=350, clear_color=[0.95, 0.95, 0.95, 1.0])

            with gr.Row():
                download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
                download_gs = gr.DownloadButton(label="Download Gaussian (PLY)", interactive=False)

    # --- Event Handlers ---
    print("Defining Gradio event handlers...")

    # Handle session start/end
    # demo.load() is valid with inputs=None, outputs=None (though default)
    demo.load(start_session, inputs=None, outputs=None)
    # >>> FIX: demo.unload() does NOT take inputs/outputs arguments <<<
    demo.unload(end_session) # Removed inputs/outputs kwargs

    # --- Generate Button Click Flow ---
    generate_event = generate_btn.click(
        get_seed,
        inputs=[randomize_seed, seed],
        outputs=[seed],
        api_name="get_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], # state_dict -> output_buf, None -> video_output
        api_name="text_to_3d"
    ).then(
        lambda: (
            gr.Button(interactive=True),
            gr.Button(interactive=True),
            gr.DownloadButton(interactive=False),
            gr.DownloadButton(interactive=False)
        ),
        inputs=None,
        outputs=[extract_glb_btn, extract_gs_btn, download_glb, download_gs],
    )

    # --- Extract GLB Button Click Flow ---
    extract_glb_event = extract_glb_btn.click(
        extract_glb,
        inputs=[output_buf, mesh_simplify, texture_size],
        outputs=[model_output, download_glb],
        api_name="extract_glb"
    ).then(
        lambda: gr.DownloadButton(interactive=True),
        inputs=None,
        outputs=[download_glb],
    )

    # --- Extract Gaussian Button Click Flow ---
    extract_gs_event = extract_gs_btn.click(
        extract_gaussian,
        inputs=[output_buf],
        outputs=[model_output, download_gs],
        api_name="extract_gaussian"
    ).then(
        lambda: gr.DownloadButton(interactive=True),
        inputs=None,
        outputs=[download_gs],
    )

    # --- Clear Download Button Interactivity when model preview is cleared ---
    model_output.clear(
        lambda: (gr.DownloadButton(interactive=False), gr.DownloadButton(interactive=False)),
        inputs=None,
        outputs=[download_glb, download_gs]
    )
    video_output.clear(
         lambda: (
            gr.Button(interactive=False),
            gr.Button(interactive=False),
            gr.DownloadButton(interactive=False),
            gr.DownloadButton(interactive=False)
        ),
        inputs=None,
        outputs=[extract_glb_btn, extract_gs_btn, download_glb, download_gs],
    )

    print("Gradio interface setup complete.")


# --- Launch the Gradio app ---
if __name__ == "__main__":
    print("Loading Trellis pipeline...")
    pipeline_loaded = False
    pipeline = None # Initialize
    try:
        pipeline = TrellisTextTo3DPipeline.from_pretrained(
            "JeffreyXiang/TRELLIS-text-xlarge",
            torch_dtype=torch.float16 # Use float16 if GPU supports it
        )
        if torch.cuda.is_available():
            pipeline = pipeline.to("cuda")
            print("✅ Trellis pipeline loaded successfully to GPU.")
        else:
            print("⚠️ WARNING: CUDA not available, running on CPU (will be very slow).")
            print("✅ Trellis pipeline loaded successfully to CPU.")
        pipeline_loaded = True
    except Exception as e:
        print(f"❌ Failed to load Trellis pipeline: {e}", file=sys.stderr)
        traceback.print_exc()
        print("❌ Exiting due to pipeline load failure.")
        sys.exit(1)

    if pipeline_loaded:
        print("Launching Gradio demo...")
        # Consider increasing queue timeout if tasks are long
        demo.queue(
             # default_concurrency_limit=2, # Limit concurrency if resource issues suspected
             # status_update_rate='auto'
        ).launch(
            # server_name="0.0.0.0", # Allows access from local network
            # share=False, # Set True for public link (careful with resources)
            debug=True, # Enable Gradio/FastAPI debug logs
            # prevent_thread_lock=True # Might help sometimes
        )
        print("Gradio demo launched.")
    else:
         print("Gradio demo not launched due to pipeline loading failure.")