import gradio as gr
import spaces

import os
import shutil
import json
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
import time
# If psutil is available in the environment, we can use it for memory info
try:
    import psutil
    PSUTIL_AVAILABLE = True
except ImportError:
    PSUTIL_AVAILABLE = False

# --- Environment Variables --- 
os.environ['TOKENIZERS_PARALLELISM'] = 'true'
os.environ['SPCONV_ALGO'] = 'native'
# ---------------------------

from typing import *

# Add JSON encoder for NumPy arrays
class NumpyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return json.JSONEncoder.default(self, obj)


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)
    
    
def end_session(req: gr.Request):
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    # Use shutil.rmtree with ignore_errors=True for robustness
    shutil.rmtree(user_dir, ignore_errors=True)


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'],
    )
    # Ensure tensors are created on the correct device ('cuda')
    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda', dtype=torch.float32)
    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda', dtype=torch.float32)
    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda', dtype=torch.float32)
    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda', dtype=torch.float32)
    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda', dtype=torch.float32)
    
    mesh = edict(
        vertices=torch.tensor(state['mesh']['vertices'], device='cuda', dtype=torch.float32),
        faces=torch.tensor(state['mesh']['faces'], device='cuda', dtype=torch.int64), # Faces are usually integers
    )
    
    return gs, mesh


def get_seed(randomize_seed: bool, seed: int) -> int:
    """
    Get the random seed.
    """
    return np.random.randint(0, MAX_SEED) if randomize_seed else 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,
) -> dict: # MODIFIED: Now returns only the state dict
    """
    Convert a text prompt to a 3D model state object.
    Args:
        prompt (str): The text prompt.
        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.
    Returns:
        dict: The JSON-serializable state object containing the generated 3D model info.
    """
    # Ensure user directory exists (redundant if start_session is always called, but safe)
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True) 
    
    print(f"[{req.session_hash}] Running text_to_3d for prompt: {prompt}") # Add logging
    
    outputs = pipeline.run(
        prompt,
        seed=seed,
        formats=["gaussian", "mesh"],
        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,
        },
    )

    # REMOVED: Video rendering logic moved to render_preview_video
    
    # Create the state object and ensure it's JSON serializable for API calls
    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
    # Convert to serializable format
    serializable_state = json.loads(json.dumps(state, cls=NumpyEncoder))
    
    print(f"[{req.session_hash}] text_to_3d completed. Returning state.") # Modified log message
    
    torch.cuda.empty_cache()
        
    # --- REVERTED DEBUGGING --- 
    # Remove the temporary simple dictionary return
    # print("[DEBUG] Returning simple dict for API test.")
    # return {"status": "test_success", "received_prompt": prompt}
    # --- END REVERTED DEBUGGING ---
    
    # Original return line (restored):
    return serializable_state # MODIFIED: Return only state

# --- NEW FUNCTION ---
@spaces.GPU
def render_preview_video(state: dict, req: gr.Request) -> str:
    """
    Renders a preview video from the provided state object.
    Args:
        state (dict): The state object containing Gaussian and mesh data.
        req (gr.Request): Gradio request object for session hash.
    Returns:
        str: The path to the rendered video file.
    """
    if not state:
        print(f"[{req.session_hash}] render_preview_video called with empty state. Returning None.")
        return None 

    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True) # Ensure directory exists
    
    print(f"[{req.session_hash}] Unpacking state for video rendering.") 
    # Only unpack gs, as mesh causes type errors with render_utils after unpacking
    gs, _ = unpack_state(state) # We still need the mesh for GLB, but not for this video preview
    
    print(f"[{req.session_hash}] Rendering video (Gaussian only)...") 
    # Render ONLY the Gaussian splats, as rendering the unpacked mesh fails
    video = render_utils.render_video(gs, num_frames=120)['color']
    # REMOVED: video_geo = render_utils.render_video(mesh, num_frames=120)['normal']
    # REMOVED: video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
    
    video_path = os.path.join(user_dir, 'preview_sample.mp4') 
    print(f"[{req.session_hash}] Saving video to {video_path}") 
    # Save only the Gaussian render
    imageio.mimsave(video_path, video, fps=15)
    
    torch.cuda.empty_cache()
    return video_path
# --- END NEW FUNCTION ---


@spaces.GPU(duration=90)
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 state.
    Args:
        state (dict): The state of the generated 3D model.
        mesh_simplify (float): The mesh simplification factor.
        texture_size (int): The texture resolution.
    Returns:
        str: The path to the extracted GLB file (for Model3D component).
        str: The path to the extracted GLB file (for DownloadButton).
    """
    if not state:
       print(f"[{req.session_hash}] extract_glb called with empty state. Returning None.")
       return None, None # Return Nones if state is missing
       
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)
    
    print(f"[{req.session_hash}] Unpacking state for GLB extraction.") # Add logging
    gs, mesh = unpack_state(state)
    
    print(f"[{req.session_hash}] Extracting GLB (simplify={mesh_simplify}, texture={texture_size})...") # Add logging
    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
    glb_path = os.path.join(user_dir, 'sample.glb')
    print(f"[{req.session_hash}] Saving GLB to {glb_path}") # Add logging
    glb.export(glb_path)
    
    torch.cuda.empty_cache()
    # Return the same path for both Model3D and DownloadButton components
    return glb_path, glb_path


@spaces.GPU
def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
    """
    Extract a Gaussian PLY file from the 3D model state.
    Args:
        state (dict): The state of the generated 3D model.
    Returns:
        str: The path to the extracted Gaussian file (for Model3D component).
        str: The path to the extracted Gaussian file (for DownloadButton).
    """
    if not state:
       print(f"[{req.session_hash}] extract_gaussian called with empty state. Returning None.")
       return None, None # Return Nones if state is missing
       
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)
    
    print(f"[{req.session_hash}] Unpacking state for Gaussian extraction.") # Add logging
    gs, _ = unpack_state(state)
    
    gaussian_path = os.path.join(user_dir, 'sample.ply')
    print(f"[{req.session_hash}] Saving Gaussian PLY to {gaussian_path}") # Add logging
    gs.save_ply(gaussian_path)
    
    torch.cuda.empty_cache()
    # Return the same path for both Model3D and DownloadButton components
    return gaussian_path, gaussian_path


# --- NEW COMBINED API FUNCTION (with HEAVY logging) ---
@spaces.GPU(duration=120)
def generate_and_extract_glb(
    prompt: str,
    seed: int,
    ss_guidance_strength: float,
    ss_sampling_steps: int,
    slat_guidance_strength: float,
    slat_sampling_steps: int,
    mesh_simplify: float,
    texture_size: int,
    req: Optional[gr.Request] = None, # Make req optional for robustness
) -> Optional[str]:
    """
    Combines 3D model generation and GLB extraction into a single step 
    for API usage, avoiding the need to transfer the state object.
    Includes extensive logging.
    Args:
        prompt (str): Text prompt for generation.
        seed (int): Random seed.
        ss_guidance_strength (float): Sparse structure guidance.
        ss_sampling_steps (int): Sparse structure steps.
        slat_guidance_strength (float): Structured latent guidance.
        slat_sampling_steps (int): Structured latent steps.
        mesh_simplify (float): Mesh simplification factor for GLB.
        texture_size (int): Texture resolution for GLB.
        req (Optional[gr.Request]): Gradio request object.
    Returns:
        Optional[str]: Path to the generated GLB file or None on failure.
    """
    # --- Setup & Initial Logging ---
    pid = os.getpid()
    session_hash = f"API_CALL_{pid}_{int(time.time()*1000)}" # More unique ID for API calls
    if req and hasattr(req, 'session_hash') and req.session_hash:
        session_hash = req.session_hash # Use session hash if available from UI call
        
    print(f"\n[{session_hash}] ========= generate_and_extract_glb INVOKED =========")
    print(f"[{session_hash}] API: PID: {pid}")
    if PSUTIL_AVAILABLE:
        process = psutil.Process(pid)
        mem_info_start = process.memory_info()
        print(f"[{session_hash}] API: Initial Memory: RSS={mem_info_start.rss / (1024**2):.2f} MB, VMS={mem_info_start.vms / (1024**2):.2f} MB")
    else:
        print(f"[{session_hash}] API: psutil not available, cannot log memory usage.")
        
    user_dir = os.path.join(TMP_DIR, str(session_hash))
    try:
        print(f"[{session_hash}] API: Ensuring directory exists: {user_dir}")
        os.makedirs(user_dir, exist_ok=True)
        print(f"[{session_hash}] API: Directory ensured.")
    except Exception as e:
        print(f"[{session_hash}] API: FATAL ERROR creating directory {user_dir}: {e}")
        traceback.print_exc()
        print(f"[{session_hash}] ========= generate_and_extract_glb FAILED (Directory Creation) =========")
        return None
    
    print(f"[{session_hash}] API: Input Params: Prompt='{prompt}', Seed={seed}, Simplify={mesh_simplify}, Texture={texture_size}")
    print(f"[{session_hash}] API: Input Params: SS Steps={ss_sampling_steps}, SS Cfg={ss_guidance_strength}, Slat Steps={slat_sampling_steps}, Slat Cfg={slat_guidance_strength}")
    
    # Check CUDA availability
    cuda_available = torch.cuda.is_available()
    print(f"[{session_hash}] API: torch.cuda.is_available(): {cuda_available}")
    if not cuda_available:
        print(f"[{session_hash}] API: FATAL ERROR - CUDA not available!")
        print(f"[{session_hash}] ========= generate_and_extract_glb FAILED (CUDA Unavailable) =========")
        return None
        
    gs_output = None
    mesh_output = None
    glb_path = None
    
    # --- Step 1: Generate 3D Model --- 
    print(f"\n[{session_hash}] API: --- Starting Step 1: Generation Pipeline --- ")
    try:
        if pipeline is None:
             print(f"[{session_hash}] API: FATAL ERROR - `pipeline` object is None!")
             raise ValueError("Trellis pipeline is not loaded.")
             
        print(f"[{session_hash}] API: Step 1 - Calling pipeline.run()...")
        t_start_gen = time.time()
        # --- The actual pipeline call --- 
        outputs = pipeline.run(
            prompt,
            seed=seed,
            formats=["gaussian", "mesh"],
            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,
            },
        )
        # --- End pipeline call ---
        t_end_gen = time.time()
        print(f"[{session_hash}] API: Step 1 - pipeline.run() completed in {t_end_gen - t_start_gen:.2f}s.")
        
        # === Validate pipeline outputs ===
        print(f"[{session_hash}] API: Step 1 - Validating pipeline outputs...")
        if not outputs:
            print(f"[{session_hash}] API: ERROR - Pipeline output dictionary is None or empty.")
            raise ValueError("Pipeline returned empty output.")
            
        if 'gaussian' not in outputs or not outputs['gaussian']:
            print(f"[{session_hash}] API: ERROR - Pipeline output missing 'gaussian' key or value is empty.")
            raise ValueError("Pipeline output missing Gaussian result.")
            
        if 'mesh' not in outputs or not outputs['mesh']:
            print(f"[{session_hash}] API: ERROR - Pipeline output missing 'mesh' key or value is empty.")
            raise ValueError("Pipeline output missing Mesh result.")
            
        gs_output = outputs['gaussian'][0]
        mesh_output = outputs['mesh'][0]
        
        if gs_output is None:
             print(f"[{session_hash}] API: ERROR - Pipeline returned gs_output as None.")
             raise ValueError("Pipeline returned None for Gaussian output.")
             
        if mesh_output is None:
             print(f"[{session_hash}] API: ERROR - Pipeline returned mesh_output as None.")
             raise ValueError("Pipeline returned None for Mesh output.")
             
        print(f"[{session_hash}] API: Step 1 - Outputs validated successfully.")
        print(f"[{session_hash}] API: Step 1 - gs_output type: {type(gs_output)}")
        # Add more details if useful, e.g., number of Gaussians
        if hasattr(gs_output, '_xyz'):
             print(f"[{session_hash}] API: Step 1 - gs_output num points: {len(gs_output._xyz)}")
             
        print(f"[{session_hash}] API: Step 1 - mesh_output type: {type(mesh_output)}")
        # Add more details if useful, e.g., number of vertices/faces
        if hasattr(mesh_output, 'vertices') and hasattr(mesh_output, 'faces'):
             print(f"[{session_hash}] API: Step 1 - mesh_output verts: {len(mesh_output.vertices)}, faces: {len(mesh_output.faces)}")
        # =================================
        
        if PSUTIL_AVAILABLE:
            mem_info_after_gen = process.memory_info()
            print(f"[{session_hash}] API: Memory After Gen: RSS={mem_info_after_gen.rss / (1024**2):.2f} MB, VMS={mem_info_after_gen.vms / (1024**2):.2f} MB")
            
    except Exception as e_gen:
        print(f"\n[{session_hash}] API: ******** ERROR IN STEP 1: Generation Pipeline ********")
        print(f"[{session_hash}] API: Error Type: {type(e_gen).__name__}, Message: {e_gen}")
        print(f"[{session_hash}] API: Printing traceback...")
        traceback.print_exc()
        print(f"[{session_hash}] API: ********************************************************")
        gs_output = None # Ensure reset on error
        mesh_output = None
        # Fall through to finally block for cleanup
    finally:
        # Attempt cleanup regardless of success/failure in try block
        print(f"[{session_hash}] API: Step 1 - Entering finally block for potential cleanup.")
        try: 
            print(f"[{session_hash}] API: Step 1 - Attempting CUDA cache clear (finally)...")
            torch.cuda.empty_cache()
            print(f"[{session_hash}] API: Step 1 - CUDA cache cleared (finally).")
        except Exception as cache_e_gen:
            print(f"[{session_hash}] API: WARNING - Error clearing CUDA cache in Step 1 finally block: {cache_e_gen}")
        print(f"[{session_hash}] API: --- Finished Step 1: Generation Pipeline (gs valid: {gs_output is not None}, mesh valid: {mesh_output is not None}) --- \n")

    # --- Step 2: Extract GLB --- 
    # Proceed only if Step 1 was successful
    if gs_output is not None and mesh_output is not None:
        print(f"\n[{session_hash}] API: --- Starting Step 2: GLB Extraction --- ")
        try:
            print(f"[{session_hash}] API: Step 2 - Inputs: gs type {type(gs_output)}, mesh type {type(mesh_output)}")
            print(f"[{session_hash}] API: Step 2 - Params: Simplify={mesh_simplify}, Texture Size={texture_size}")
            print(f"[{session_hash}] API: Step 2 - Calling postprocessing_utils.to_glb()...")
            t_start_glb = time.time()
            # --- The actual GLB conversion call --- 
            glb = postprocessing_utils.to_glb(gs_output, mesh_output, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
            # --- End GLB conversion call ---
            t_end_glb = time.time()
            print(f"[{session_hash}] API: Step 2 - postprocessing_utils.to_glb() completed in {t_end_glb - t_start_glb:.2f}s.")
            
            # === Validate GLB output ===
            print(f"[{session_hash}] API: Step 2 - Validating GLB object...")
            if glb is None:
                print(f"[{session_hash}] API: ERROR - postprocessing_utils.to_glb returned None.")
                raise ValueError("GLB conversion returned None.")
            print(f"[{session_hash}] API: Step 2 - GLB object validated successfully (type: {type(glb)})...")
            # ==========================
            
            # === Save GLB ===
            glb_path = os.path.join(user_dir, f'api_generated_{session_hash}_{int(time.time()*1000)}.glb') # More unique name
            print(f"[{session_hash}] API: Step 2 - Saving GLB to path: {glb_path}...")
            t_start_save = time.time()
            # --- The actual GLB export call --- 
            glb.export(glb_path)
            # --- End GLB export call ---
            t_end_save = time.time()
            print(f"[{session_hash}] API: Step 2 - glb.export() completed in {t_end_save - t_start_save:.2f}s.")
            # =================

            # === Verify File Exists ===
            print(f"[{session_hash}] API: Step 2 - Verifying saved file exists at {glb_path}...")
            if not os.path.exists(glb_path):
                 print(f"[{session_hash}] API: ERROR - GLB file was not found after export at {glb_path}.")
                 raise IOError(f"GLB export failed, file not found: {glb_path}")
            print(f"[{session_hash}] API: Step 2 - Saved file verified.")
            # =========================
            
            print(f"[{session_hash}] API: Step 2 - GLB extraction and saving completed successfully.")
            if PSUTIL_AVAILABLE:
                mem_info_after_glb = process.memory_info()
                print(f"[{session_hash}] API: Memory After GLB: RSS={mem_info_after_glb.rss / (1024**2):.2f} MB, VMS={mem_info_after_glb.vms / (1024**2):.2f} MB")
                
        except Exception as e_glb:
            print(f"\n[{session_hash}] API: ******** ERROR IN STEP 2: GLB Extraction ********")
            print(f"[{session_hash}] API: Error Type: {type(e_glb).__name__}, Message: {e_glb}")
            print(f"[{session_hash}] API: Printing traceback...")
            traceback.print_exc()
            print(f"[{session_hash}] API: *****************************************************")
            glb_path = None # Ensure reset on error
            # Fall through to finally block for cleanup
        finally:
            # Attempt cleanup regardless of success/failure in try block
            print(f"[{session_hash}] API: Step 2 - Entering finally block for potential cleanup.")
            # Explicitly delete large objects if possible (might help memory)
            del glb
            print(f"[{session_hash}] API: Step 2 - Deleted intermediate 'glb' object.")
            try: 
                print(f"[{session_hash}] API: Step 2 - Attempting CUDA cache clear (finally)...")
                torch.cuda.empty_cache()
                print(f"[{session_hash}] API: Step 2 - CUDA cache cleared (finally).")
            except Exception as cache_e_glb:
                print(f"[{session_hash}] API: WARNING - Error clearing CUDA cache in Step 2 finally block: {cache_e_glb}")
            print(f"[{session_hash}] API: --- Finished Step 2: GLB Extraction (path valid: {glb_path is not None}) --- \n")
    else:
        print(f"[{session_hash}] API: Skipping Step 2 (GLB Extraction) because Step 1 failed or produced invalid outputs.")
        glb_path = None # Ensure glb_path is None if Step 1 failed

    # --- Final Cleanup and Return --- 
    print(f"[{session_hash}] API: --- Entering Final Cleanup and Return --- ")
    # Final attempt to clear CUDA cache
    try:
        print(f"[{session_hash}] API: Final CUDA cache clear attempt...")
        torch.cuda.empty_cache()
        print(f"[{session_hash}] API: Final CUDA cache cleared.")
    except Exception as cache_e_final:
        print(f"[{session_hash}] API: WARNING - Error clearing final CUDA cache: {cache_e_final}")
        
    # Explicitly delete pipeline outputs if they exist
    del gs_output
    del mesh_output
    print(f"[{session_hash}] API: Deleted intermediate 'gs_output' and 'mesh_output' objects.")
    
    # Final decision based on glb_path status
    if glb_path and os.path.exists(glb_path):
        print(f"[{session_hash}] API: Final Result: SUCCESS. GLB Path: {glb_path}")
        print(f"[{session_hash}] ========= generate_and_extract_glb END (Success) =========")
        return glb_path
    else:
        print(f"[{session_hash}] API: Final Result: FAILURE. GLB Path: {glb_path} (Exists: {os.path.exists(glb_path) if glb_path else 'N/A'})")
        print(f"[{session_hash}] ========= generate_and_extract_glb END (Failure) =========")
        return None
# --- END NEW COMBINED API FUNCTION ---


# State object to hold the generated model info between steps
output_buf = gr.State() 
# Video component placeholder (will be populated by render_preview_video)
# video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300) # Defined later inside the Blocks

with gr.Blocks(delete_cache=(600, 600)) 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.
    * The preview video will appear after generation.
    * If you find the generated 3D asset satisfactory, click "Extract GLB" or "Extract Gaussian" to extract the file and download it.
    """)
    
    with gr.Row():
        with gr.Column():
            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("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=25, 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=7.5, step=0.1)
                    slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=25, step=1)

            generate_btn = gr.Button("Generate")
            
            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)
            
            with gr.Row():
                # Buttons start non-interactive, enabled after generation
                extract_glb_btn = gr.Button("Extract GLB", interactive=False)
                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():
            # Define UI components here
            video_output = gr.Video(label="Generated 3D Asset Preview", autoplay=True, loop=True, height=300)
            model_output = gr.Model3D(label="Extracted GLB/Gaussian", height=300)
            
            with gr.Row():
                 # Buttons start non-interactive, enabled after extraction
                download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
                download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)  
    
    # Define the state buffer here, outside the component definitions but inside the Blocks scope
    output_buf = gr.State()

    # --- Handlers ---
    demo.load(start_session)
    demo.unload(end_session)

    # --- MODIFIED UI CHAIN ---
    # 1. Get Seed
    # 2. Run text_to_3d -> outputs state to output_buf
    # 3. Run render_preview_video (using state from output_buf) -> outputs video to video_output
    # 4. Enable extraction buttons
    generate_btn.click(
        get_seed,
        inputs=[randomize_seed, seed],
        outputs=[seed],
        queue=True # Use queue for potentially long-running steps
    ).then(
        text_to_3d,
        inputs=[text_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
        outputs=[output_buf], # text_to_3d now ONLY outputs state
        api_name="text_to_3d" # Keep API name consistent if needed
    ).then(
        render_preview_video, # NEW step: Render video from state
        inputs=[output_buf],
        outputs=[video_output],
        api_name="render_preview_video" # Assign API name if you want to call this separately
    ).then(
        lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]), # Enable extraction buttons
        outputs=[extract_glb_btn, extract_gs_btn],
    )

    # Clear video and disable extraction buttons if prompt is cleared or generation restarted
    # (Consider adding logic to clear prompt on successful generation if desired)
    text_prompt.change( # Example: Clear video if prompt changes
         lambda: (None, gr.Button(interactive=False), gr.Button(interactive=False)),
         outputs=[video_output, extract_glb_btn, extract_gs_btn]
    )
    video_output.clear( # This might be redundant if text_prompt.change handles it
        lambda: tuple([gr.Button(interactive=False), gr.Button(interactive=False)]),
        outputs=[extract_glb_btn, extract_gs_btn],
    )

    # --- Extraction Handlers ---
    # GLB Extraction: Takes state from output_buf, outputs model and download path
    extract_glb_btn.click(
        extract_glb,
        inputs=[output_buf, mesh_simplify, texture_size],
        outputs=[model_output, download_glb], # Outputs to Model3D and DownloadButton path
        api_name="extract_glb"
    ).then(
        lambda: gr.Button(interactive=True), # Enable download button
        outputs=[download_glb],
    )
    
    # Gaussian Extraction: Takes state from output_buf, outputs model and download path
    extract_gs_btn.click(
        extract_gaussian,
        inputs=[output_buf],
        outputs=[model_output, download_gs], # Outputs to Model3D and DownloadButton path
        api_name="extract_gaussian"
    ).then(
        lambda: gr.Button(interactive=True), # Enable download button
        outputs=[download_gs],
    )

    # Clear model and disable download buttons if video/state is cleared
    model_output.clear(
        lambda: (gr.Button(interactive=False), gr.Button(interactive=False)),
        outputs=[download_glb, download_gs], # Disable both download buttons
    )
    
# --- Launch the Gradio app ---
if __name__ == "__main__":
    print("Initializing pipeline...")
    pipeline = None # Initialize pipeline variable
    try:
        # Load the pipeline
        pipeline = TrellisTextTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-text-xlarge")
        # Move pipeline to CUDA device
        pipeline.cuda()
        print("Pipeline loaded and moved to CUDA successfully.")
    except Exception as e:
        print(f"FATAL ERROR initializing pipeline: {e}")
        traceback.print_exc()
        # Optionally exit if pipeline loading fails
        sys.exit(1) 
        
    print("Launching Gradio demo with queue enabled...")
    demo.queue().launch()