revert back to 1.1

app.py

@@ -1,20 +1,22 @@
-# Version: 1.1.0 - API State Fix + DEBUG (Video Disabled - FINAL CORRECTED BASELINE) (2025-05-04)
+# Version: 1.1.0 - API State Fix (2025-05-04)
 # Changes:
-# - Based *EXACTLY* on user-provided Version 1.1.0 code that fixed state dict & loaded pipeline.
-# - TEMPORARY DEBUGGING STEP: Commented out video rendering/saving in `text_to_3d`
-#   and return None for video_path to isolate the "Session not found" error.
-# - No other changes were made. Removed previous erroneous additions.
+# - Modified `text_to_3d` to explicitly return the serializable `state_dict` from `pack_state`
+#   as the first return value. This ensures the dictionary is available via the API.
+# - Modified `extract_glb` and `extract_gaussian` to accept `state_dict: dict` as their first argument
+#   instead of relying on the implicit `gr.State` object type when called via API.
+# - Kept Gradio UI bindings (`outputs=[output_buf, ...]`, `inputs=[output_buf, ...]`)
+#   so the UI continues to function by passing the dictionary through output_buf.
+# - Added minor safety checks and logging.
 
 import gradio as gr
-# NOTE: Ensuring 'spaces' is imported if decorators are used (was missing in user provided snippet but needed)
-# If @spaces.GPU decorators are not used, this import is not needed.
-# Assuming they ARE used based on previous context:
 import spaces
 
 import os
 import shutil
 os.environ['TOKENIZERS_PARALLELISM'] = 'true'
-os.environ['SPCONV_ALGO'] = 'native' # Direct set as per original user code
+# Fix potential SpConv issue if needed, try 'hash' or 'native'
+# os.environ.setdefault('SPCONV_ALGO', 'native') # Use setdefault to avoid overwriting if already set
+os.environ['SPCONV_ALGO'] = 'native' # Direct set as per original
 
 from typing import *
 import torch
@@ -28,128 +30,102 @@ from trellis.utils import render_utils, postprocessing_utils
 import traceback
 import sys
 
+
 MAX_SEED = np.iinfo(np.int32).max
-# Using path relative to file as in original user provided code
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
-# Ensure base directory exists
-try:
-    os.makedirs(TMP_DIR, exist_ok=True)
-    print(f"Using temporary directory: {TMP_DIR}")
-except OSError as e:
-     print(f"Warning: Could not create base temp directory {TMP_DIR}: {e}", file=sys.stderr)
-     TMP_DIR = '.' # Fallback
-     print(f"Warning: Falling back to use current directory for temp files: {os.path.abspath(TMP_DIR)}")
+os.makedirs(TMP_DIR, exist_ok=True)
+
 
 def start_session(req: gr.Request):
     """Creates a temporary directory for the user session."""
-    user_dir = None
-    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}")
-        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)
+    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):
     """Removes the temporary directory for the user session."""
-    user_dir = None
-    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):
-            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)
+    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 serializable dictionary."""
+    # Ensure tensors are on CPU and converted to numpy before returning the dict
     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
+    packed_data = {
+        'gaussian': {
+            # Spread init_params first to ensure correct types
+            **{k: v for k, v in gs.init_params.items()}, # Ensure init_params are included
+            '_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
 
 
 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]}")
+    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.")
+
+    # Ensure the device is correctly set when unpacking
+    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']
+
+    # Recreate Gaussian object using parameters stored during packing
+    gs = Gaussian(
+        aabb=gauss_data.get('aabb'), # Use .get for safety
+        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'),
+    )
+    # Load tensors, ensuring they are created on the correct device
+    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]}")
+
+    # Recreate mesh object using edict for compatibility if needed elsewhere
+    mesh = edict(
+        vertices=torch.tensor(mesh_data['vertices'], device=device, dtype=torch.float32),
+        faces=torch.tensor(mesh_data['faces'], device=device, dtype=torch.int64), # Faces are typically long/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
+    return gs, mesh
 
 
 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)
+    return int(new_seed) # Ensure it's a standard int
 
 
-# Decorator requires 'import spaces' at the top
 @spaces.GPU
 def text_to_3d(
     prompt: str,
@@ -159,84 +135,79 @@ def text_to_3d(
     slat_guidance_strength: float,
     slat_sampling_steps: int,
     req: gr.Request,
-) -> Tuple[dict, Optional[str]]: # Return Optional[str] for video path
+) -> Tuple[dict, str]: # Return type changed for clarity
     """
     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 <<<
+    serializable state dictionary and a video preview path.
     """
-    print(f"[text_to_3d - DEBUG MODE] Received prompt: '{prompt}', Seed: {seed}")
-    user_dir = None
-    state_dict = None
+    print(f"[text_to_3d] Received prompt: '{prompt}', Seed: {seed}")
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    os.makedirs(user_dir, exist_ok=True)
+    print(f"[text_to_3d] User directory: {user_dir}")
+
+    # --- Generation Pipeline ---
     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}")
-        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...")
+        print("[text_to_3d] Running Trellis pipeline...")
         outputs = pipeline.run(
-            prompt=prompt,
+            prompt,
             seed=seed,
-            formats=["gaussian", "mesh"],
+            formats=["gaussian", "mesh"], # Ensure both are generated
             sparse_structure_sampler_params={
-                "steps": int(ss_sampling_steps),
+                "steps": int(ss_sampling_steps), # Ensure steps are int
                 "cfg_strength": float(ss_guidance_strength),
             },
             slat_sampler_params={
-                "steps": int(slat_sampling_steps),
+                "steps": int(slat_sampling_steps), # Ensure steps are int
                 "cfg_strength": float(slat_guidance_strength),
             },
         )
-        print("[text_to_3d - DEBUG MODE] Pipeline run completed.")
+        print("[text_to_3d] Pipeline run completed.")
+    except Exception as e:
+        print(f"❌ [text_to_3d] Pipeline error: {e}", file=sys.stderr)
+        traceback.print_exc()
+        # Return an empty dict and maybe an error indicator path or None?
+        # For now, re-raise to signal failure clearly upstream.
+        raise gr.Error(f"Trellis pipeline failed: {e}")
 
-        # --- Create Serializable State Dictionary ---
+    # --- Create Serializable State Dictionary --- VITAL CHANGE for API
+    # This dictionary holds the necessary data for later extraction.
+    try:
         state_dict = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
+    except Exception as e:
+        print(f"❌ [text_to_3d] pack_state error: {e}", file=sys.stderr)
+        traceback.print_exc()
+        raise gr.Error(f"Failed to pack state: {e}")
 
+    # --- Render Video Preview ---
+    try:
+        print("[text_to_3d] Rendering video preview...")
+        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']
+        # Ensure video frames are uint8
+        video = [np.concatenate([v.astype(np.uint8), vg.astype(np.uint8)], axis=1) for v, vg in zip(video, video_geo)]
+        video_path = os.path.join(user_dir, 'sample.mp4')
+        imageio.mimsave(video_path, video, fps=15, quality=8) # Added quality setting
+        print(f"[text_to_3d] Video saved to: {video_path}")
     except Exception as e:
-        print(f"❌ [text_to_3d - DEBUG MODE] Error during generation or packing: {e}", file=sys.stderr)
+        print(f"❌ [text_to_3d] Video rendering/saving error: {e}", file=sys.stderr)
         traceback.print_exc()
-        raise gr.Error(f"Core generation failed: {e}")
-
-    # --- Render Video Preview (TEMPORARILY DISABLED FOR DEBUGGING) ---
-    video_path = None # Explicitly set path to None for this debug version
-    print("[text_to_3d - DEBUG MODE] Skipping video rendering.")
-    # --- Start Original Video Code Block (Commented Out) ---
-    # try:
-    #     print("[text_to_3d] Rendering video preview...")
-    #     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([v.astype(np.uint8), vg.astype(np.uint8)], axis=1) for v, vg in zip(video, video_geo)]
-    #     video_path_tmp = os.path.join(user_dir, 'sample.mp4')
-    #     imageio.mimsave(video_path_tmp, video, fps=15, quality=8)
-    #     print(f"[text_to_3d] Video saved to: {video_path_tmp}")
-    #     video_path = video_path_tmp
-    # except Exception as e:
-    #     print(f"❌ [text_to_3d] Video rendering/saving error: {e}", file=sys.stderr)
-    #     traceback.print_exc()
-    #     video_path = None # Indicate video failure
-    # --- End Original Video Code Block ---
+        # Still return state_dict, but maybe signal video error? Return None for path.
+        video_path = None # Indicate video failure
 
     # --- Cleanup and Return ---
+    # Clear CUDA cache if GPU was used
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-        print("[text_to_3d - DEBUG MODE] Cleared CUDA cache.")
+        print("[text_to_3d] Cleared CUDA cache.")
 
-    print("[text_to_3d - DEBUG MODE] Returning state dictionary and None video path.")
-    if state_dict is None:
-         print("Error: state_dict is None before return, generation likely failed.", file=sys.stderr)
-         raise gr.Error("State dictionary creation failed.")
+    # --- Return Serializable Dictionary and Video Path --- VITAL CHANGE for API
+    print("[text_to_3d] Returning state dictionary and video path.")
     return state_dict, video_path
 
 
-# Decorator requires 'import spaces' at the top
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=120) # Increased duration slightly
 def extract_glb(
-    state_dict: dict,
+    state_dict: dict, # <-- VITAL CHANGE: Accept the dictionary directly
     mesh_simplify: float,
     texture_size: int,
     req: gr.Request,
@@ -245,37 +216,32 @@ def extract_glb(
     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
-    glb_path = None
-    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.")
 
-        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(req.session_hash))
+    os.makedirs(user_dir, exist_ok=True)
+    print(f"[extract_glb] User directory: {user_dir}")
 
-        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 ---
+    # --- Unpack state from the dictionary --- VITAL CHANGE for API
+    try:
         gs, mesh = unpack_state(state_dict)
+    except Exception as e:
+        print(f"❌ [extract_glb] unpack_state error: {e}", file=sys.stderr)
+        traceback.print_exc()
+        raise gr.Error(f"Failed to unpack state: {e}")
 
-        # --- Postprocessing and Export ---
+    # --- Postprocessing and Export ---
+    try:
         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 = postprocessing_utils.to_glb(gs, mesh, simplify=float(mesh_simplify), texture_size=int(texture_size), verbose=True) # Verbose for debugging
         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)
+        print(f"❌ [extract_glb] GLB conversion/export error: {e}", file=sys.stderr)
         traceback.print_exc()
         raise gr.Error(f"Failed to extract GLB: {e}")
 
@@ -284,50 +250,44 @@ def extract_glb(
         torch.cuda.empty_cache()
         print("[extract_glb] Cleared CUDA cache.")
 
+    # Return path twice for both Model3D and DownloadButton components
     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
 
 
-# Decorator requires 'import spaces' at the top
 @spaces.GPU
 def extract_gaussian(
-    state_dict: dict,
+    state_dict: dict, # <-- VITAL CHANGE: Accept the dictionary directly
     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
-    gaussian_path = None
-    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.")
+    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}")
+    user_dir = os.path.join(TMP_DIR, str(req.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)
+    # --- Unpack state from the dictionary --- VITAL CHANGE for API
+    try:
+        gs, _ = unpack_state(state_dict) # Only need Gaussian part
+    except Exception as e:
+        print(f"❌ [extract_gaussian] unpack_state error: {e}", file=sys.stderr)
+        traceback.print_exc()
+        raise gr.Error(f"Failed to unpack state: {e}")
 
-        # --- Export PLY ---
+    # --- Export PLY ---
+    try:
         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)
+        print(f"❌ [extract_gaussian] PLY saving error: {e}", file=sys.stderr)
         traceback.print_exc()
         raise gr.Error(f"Failed to extract Gaussian PLY: {e}")
 
@@ -336,10 +296,8 @@ def extract_gaussian(
         torch.cuda.empty_cache()
         print("[extract_gaussian] Cleared CUDA cache.")
 
+    # Return path twice for both Model3D and DownloadButton components
     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
 
 
@@ -351,14 +309,17 @@ with gr.Blocks(delete_cache=(600, 600), title="TRELLIS Text-to-3D") as demo:
     * 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 ---
+    # This hidden component will hold the dictionary returned by text_to_3d,
+    # acting as the state link between generation and extraction for the UI/API.
     output_buf = gr.State()
 
     with gr.Row():
-        with gr.Column(scale=1):
+        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)
@@ -370,83 +331,99 @@ with gr.Blocks(delete_cache=(600, 600), title="TRELLIS Text-to-3D") as demo:
                 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):
+
+            with gr.Accordion(label="GLB Extraction Settings", open=True): # Open by default
+                # Tooltips added for clarity
                 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):
-            # 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.Column(scale=1): # Output column
+            video_output = gr.Video(label="Generated 3D Preview (Geometry | Texture)", autoplay=True, loop=True, height=350) # Slightly larger height
+            model_output = gr.Model3D(label="Extracted Model Preview", height=350, clear_color=[0.95, 0.95, 0.95, 1.0]) # Light background
+
             with gr.Row():
+                # Link download button visibility/interactivity to model_output potentially
                 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...")
-    # Use demo.load as in original user-provided code
-    demo.load(start_session, inputs=None, outputs=None)
-    # Use demo.unload as in original user-provided code (no extra args)
-    demo.unload(end_session) # Corrected: removed inputs/outputs
 
+    # Handle session start/end
+    demo.load(start_session, inputs=None, outputs=None) # Pass None for clarity
+    demo.unload(end_session, inputs=None, outputs=None)
+
+    # --- Generate Button Click Flow ---
+    # 1. Get Seed -> 2. Run text_to_3d -> 3. Enable extraction buttons
     generate_event = generate_btn.click(
         get_seed,
         inputs=[randomize_seed, seed],
         outputs=[seed],
-        api_name="get_seed"
+        api_name="get_seed" # Optional API name
     ).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
+        outputs=[output_buf, video_output], # output_buf receives state_dict
         api_name="text_to_3d"
     ).then(
-        lambda: (
-            gr.Button(interactive=True), gr.Button(interactive=True),
-            gr.DownloadButton(interactive=False), gr.DownloadButton(interactive=False)
+        lambda: ( # Return tuple for multiple outputs
+            gr.Button(interactive=True),
+            gr.Button(interactive=True),
+            gr.DownloadButton(interactive=False), # Ensure download buttons are disabled initially
+            gr.DownloadButton(interactive=False)
         ),
-        inputs=None,
-        outputs=[extract_glb_btn, extract_gs_btn, download_glb, download_gs],
+        outputs=[extract_glb_btn, extract_gs_btn, download_glb, download_gs], # Update interactivity
     )
 
+    # --- Clear video/model outputs if prompt changes (optional, prevents confusion)
+    # text_prompt.change(lambda: (None, None, gr.Button(interactive=False), gr.Button(interactive=False)), outputs=[video_output, model_output, extract_glb_btn, extract_gs_btn])
+
+    # --- Extract GLB Button Click Flow ---
+    # 1. Run extract_glb -> 2. Update Model3D and Download Button
     extract_glb_event = extract_glb_btn.click(
         extract_glb,
-        inputs=[output_buf, mesh_simplify, texture_size],
-        outputs=[model_output, download_glb],
+        inputs=[output_buf, mesh_simplify, texture_size], # Pass the state_dict via output_buf
+        outputs=[model_output, download_glb], # Returns path to both
         api_name="extract_glb"
     ).then(
-        lambda: gr.DownloadButton(interactive=True),
-        inputs=None,
+        lambda: gr.DownloadButton(interactive=True), # Enable download button
         outputs=[download_glb],
     )
 
+    # --- Extract Gaussian Button Click Flow ---
+    # 1. Run extract_gaussian -> 2. Update Model3D and Download Button
     extract_gs_event = extract_gs_btn.click(
         extract_gaussian,
-        inputs=[output_buf],
-        outputs=[model_output, download_gs],
+        inputs=[output_buf], # Pass the state_dict via output_buf
+        outputs=[model_output, download_gs], # Returns path to both
         api_name="extract_gaussian"
     ).then(
-        lambda: gr.DownloadButton(interactive=True),
-        inputs=None,
+        lambda: gr.DownloadButton(interactive=True), # Enable download button
         outputs=[download_gs],
     )
 
+    # --- Clear Download Button Interactivity when model preview is cleared ---
+    # This might be redundant if generate disables them, but adds safety
     model_output.clear(
         lambda: (gr.DownloadButton(interactive=False), gr.DownloadButton(interactive=False)),
-        inputs=None,
         outputs=[download_glb, download_gs]
     )
-    video_output.clear(
+    video_output.clear( # Also disable extraction if video is cleared (e.g., new generation starts)
          lambda: (
-            gr.Button(interactive=False), gr.Button(interactive=False),
-            gr.DownloadButton(interactive=False), gr.DownloadButton(interactive=False)
+            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],
     )
 
@@ -456,30 +433,33 @@ with gr.Blocks(delete_cache=(600, 600), title="TRELLIS Text-to-3D") as demo:
 # --- Launch the Gradio app ---
 if __name__ == "__main__":
     print("Loading Trellis pipeline...")
-    pipeline = None
-    pipeline_loaded = False
     try:
-        # --- Load pipeline WITHOUT torch_dtype (As per original working version) ---
+        # Ensure model/variant matches requirements, use revision if needed
         pipeline = TrellisTextTo3DPipeline.from_pretrained(
-            "JeffreyXiang/TRELLIS-text-xlarge"
+            "JeffreyXiang/TRELLIS-text-xlarge",
+            # revision="main", # Specify if needed
+            torch_dtype=torch.float16 # Use float16 if GPU supports it for less memory
         )
+        # Move to GPU if available
         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.")
+            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()
+        # Exit if pipeline is critical for the app to run
         print("❌ Exiting due to pipeline load failure.")
         sys.exit(1)
 
-    if pipeline_loaded:
-        print("Launching Gradio demo...")
-        # Use queue and debug=True as in original user-provided code
-        demo.queue().launch(debug=True)
-        print("Gradio demo launched.")
-    else:
-         print("Gradio demo not launched.")
+    print("Launching Gradio demo...")
+    # Set share=True if you need a public link (e.g., for testing from outside local network)
+    # Set server_name="0.0.0.0" to allow access from local network IP
+    demo.queue().launch( # Use queue for potentially long-running tasks
+        # server_name="0.0.0.0",
+        # share=False,
+        debug=True # Enable debug mode for more logs
+    )
+    print("Gradio demo launched.")