chat gpot

app.py

@@ -1,22 +1,17 @@
-# Version: 1.1.0 - API State Fix (2025-05-04)
-# Changes:
-# - 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.
+# Version: 1.1.1 - Targeted signature and state fixes
+# 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
+#  - Bumped version, added comments at change sites
 
 import gradio as gr
 import spaces
-
 import os
 import shutil
 os.environ['TOKENIZERS_PARALLELISM'] = 'true'
-# 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
+os.environ['SPCONV_ALGO'] = 'native'
 
 from typing import *
 import torch
@@ -26,105 +21,80 @@ 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):
-    """Creates a temporary directory for the user session."""
     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 = 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)
+        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...")
+    """Packs Gaussian and Mesh data into a JSON-serializable dictionary."""
     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(),
+            **{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(),
-            'faces': mesh.faces.detach().cpu().numpy(),
+            'vertices': mesh.vertices.detach().cpu().numpy().tolist(),
+            'faces': mesh.faces.detach().cpu().numpy().tolist(),
         },
     }
-    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...")
-    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
+    print("[unpack_state] Unpacking state from dictionary... ")
     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
+        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'),
     )
-    # 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
+    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(mesh_data['vertices'], device=device, dtype=torch.float32),
-        faces=torch.tensor(mesh_data['faces'], device=device, dtype=torch.int64), # Faces are typically long/int64
+        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),
     )
-    print(f"[unpack_state] Mesh unpacked. Vertices: {mesh.vertices.shape[0]}, Faces: {mesh.faces.shape[0]}")
-
     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) # Ensure it's a standard int
-
+    return int(new_seed)
 
 @spaces.GPU
 def text_to_3d(
@@ -135,331 +105,137 @@ def text_to_3d(
     slat_guidance_strength: float,
     slat_sampling_steps: int,
     req: gr.Request,
-) -> Tuple[dict, str]: # Return type changed for clarity
-    """
-    Generates a 3D model (Gaussian and Mesh) from text and returns a
-    serializable state dictionary and a video preview path.
-    """
-    print(f"[text_to_3d] Received prompt: '{prompt}', Seed: {seed}")
+) -> 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)
-    print(f"[text_to_3d] User directory: {user_dir}")
-
-    # --- Generation Pipeline ---
-    try:
-        print("[text_to_3d] Running Trellis pipeline...")
-        outputs = pipeline.run(
-            prompt,
-            seed=seed,
-            formats=["gaussian", "mesh"], # Ensure both are generated
-            sparse_structure_sampler_params={
-                "steps": int(ss_sampling_steps), # Ensure steps are int
-                "cfg_strength": float(ss_guidance_strength),
-            },
-            slat_sampler_params={
-                "steps": int(slat_sampling_steps), # Ensure steps are int
-                "cfg_strength": float(slat_guidance_strength),
-            },
-        )
-        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 --- 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] Video rendering/saving error: {e}", file=sys.stderr)
-        traceback.print_exc()
-        # 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] Cleared CUDA cache.")
-
-    # --- Return Serializable Dictionary and Video Path --- VITAL CHANGE for API
-    print("[text_to_3d] Returning state dictionary and video path.")
+    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) # Increased duration slightly
+@spaces.GPU(duration=120)
 def extract_glb(
-    state_dict: dict, # <-- VITAL CHANGE: Accept the dictionary directly
+    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}")
-    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.")
-
+    gs, mesh = unpack_state(state_dict)
     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}")
-
-    # --- 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 ---
-    try:
-        print("[extract_glb] Converting to GLB...")
-        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] GLB conversion/export error: {e}", file=sys.stderr)
-        traceback.print_exc()
-        raise gr.Error(f"Failed to extract GLB: {e}")
-
-    # --- Cleanup and Return ---
-    if torch.cuda.is_available():
-        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.")
+    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, # <-- VITAL CHANGE: Accept the dictionary directly
+    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.")
-    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.")
-
+    gs, _ = unpack_state(state_dict)
     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 --- 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 ---
-    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] PLY saving error: {e}", file=sys.stderr)
-        traceback.print_exc()
-        raise gr.Error(f"Failed to extract Gaussian PLY: {e}")
-
-    # --- Cleanup and Return ---
-    if torch.cuda.is_available():
-        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.")
+    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 ---
-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.
     """)
 
-    # --- 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.
+    # 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.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: 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)
-
+                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): # 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): # 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...")
-
-    # Handle session start/end
-demo.load(start_session)
-demo.unload(end_session)
-
-    # --- Generate Button Click Flow ---
-    # 1. Get Seed -> 2. Run text_to_3d -> 3. Enable extraction buttons
+            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],
-        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], # output_buf receives state_dict
-        api_name="text_to_3d"
+        outputs=[output_buf, video_output],
     ).then(
-        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)
-        ),
-        outputs=[extract_glb_btn, extract_gs_btn, download_glb, download_gs], # Update interactivity
+        lambda: (extract_glb_btn.update(interactive=True), extract_gs_btn.update(interactive=True)),
+        outputs=[extract_glb_btn, extract_gs_btn],
     )
 
-    # --- 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], # Pass the state_dict via output_buf
-        outputs=[model_output, download_glb], # Returns path to both
-        api_name="extract_glb"
+        inputs=[output_buf, mesh_simplify, texture_size],
+        outputs=[model_output, download_glb],
     ).then(
-        lambda: gr.DownloadButton(interactive=True), # Enable download button
+        lambda: download_glb.update(interactive=True),
         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], # Pass the state_dict via output_buf
-        outputs=[model_output, download_gs], # Returns path to both
-        api_name="extract_gaussian"
+        inputs=[output_buf],
+        outputs=[model_output, download_gs],
     ).then(
-        lambda: gr.DownloadButton(interactive=True), # Enable download button
+        lambda: download_gs.update(interactive=True),
         outputs=[download_gs],
     )
 
-    # --- Clear Download Button Interactivity when model preview is cleared ---
-    # This might be redundant if generate disables them, but adds safety
+    # Clear callbacks
     model_output.clear(
-        lambda: (gr.DownloadButton(interactive=False), gr.DownloadButton(interactive=False)),
-        outputs=[download_glb, download_gs]
+        lambda: (download_glb.update(interactive=False), download_gs.update(interactive=False)),
+        outputs=[download_glb, download_gs],
     )
-    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)
-        ),
+    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],
     )
 
-    print("Gradio interface setup complete.")
-
-
-# --- Launch the Gradio app ---
 if __name__ == "__main__":
-    print("Loading Trellis pipeline...")
-    try:
-        # Ensure model/variant matches requirements, use revision if needed
-        pipeline = TrellisTextTo3DPipeline.from_pretrained(
-            "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 (will be very slow).")
-            print("✅ Trellis pipeline loaded successfully to CPU.")
-    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)
-
-    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
+    pipeline = TrellisTextTo3DPipeline.from_pretrained(
+        "JeffreyXiang/TRELLIS-text-xlarge",
+        torch_dtype=torch.float16
     )
-    print("Gradio demo launched.")
+    if torch.cuda.is_available(): pipeline = pipeline.to("cuda")
+    demo.queue().launch(debug=True)