# Version: 1.1.3 - API State Fix + DEBUG (Video Disabled) + unload() Fix (2025-05-04) # Changes: # - FIXED TypeError in demo.unload() by removing incorrect 'inputs'/'outputs' arguments. # - ENSURED `import spaces` is present for the @spaces.GPU decorator. # - TEMPORARY DEBUGGING STEP: Commented out video rendering in `text_to_3d` # and return None for video_path to isolate the "Session not found" error. # - Modified `text_to_3d` to explicitly return the serializable `state_dict` from `pack_state`. # - Modified `extract_glb`/`extract_gaussian` to accept `state_dict: dict`. # - Kept Gradio UI bindings using `output_buf`. # - Added minor safety checks and logging. import gradio as gr import spaces # <<<--- ENSURE THIS IMPORT IS PRESENT import os import shutil os.environ['TOKENIZERS_PARALLELISM'] = 'true' os.environ['SPCONV_ALGO'] = 'native' # Direct set as per original from typing import * import torch import numpy as np import imageio from easydict import EasyDict as edict from trellis.pipelines import TrellisTextTo3DPipeline from trellis.representations import Gaussian, MeshExtractResult from trellis.utils import render_utils, postprocessing_utils import traceback import sys MAX_SEED = np.iinfo(np.int32).max # Use standard /tmp directory which is usually available in container environments TMP_DIR = '/tmp/gradio_sessions' print(f"Using temporary directory: {TMP_DIR}") # Ensure the base temp directory exists try: os.makedirs(TMP_DIR, exist_ok=True) except OSError as e: print(f"Warning: Could not create base temp directory {TMP_DIR}: {e}", file=sys.stderr) # Potentially fall back or exit if temp dir is critical TMP_DIR = '.' # Fallback to current directory (less ideal) print(f"Warning: Falling back to use current directory for temp files: {os.path.abspath(TMP_DIR)}") def start_session(req: gr.Request): """Creates a temporary directory for the user session.""" user_dir = None # Initialize try: session_hash = req.session_hash if not session_hash: session_hash = f"no_session_{np.random.randint(10000, 99999)}" print(f"Warning: No session_hash in request, using temporary ID: {session_hash}") # Ensure TMP_DIR exists before joining path if not os.path.exists(TMP_DIR): os.makedirs(TMP_DIR, exist_ok=True) user_dir = os.path.join(TMP_DIR, str(session_hash)) os.makedirs(user_dir, exist_ok=True) print(f"Started session, ensured directory exists: {user_dir}") except Exception as e: print(f"Error in start_session creating directory '{user_dir}': {e}", file=sys.stderr) traceback.print_exc() def end_session(req: gr.Request): """Removes the temporary directory for the user session.""" user_dir = None # Initialize try: session_hash = req.session_hash if not session_hash: print("Warning: No session_hash in end_session request, cannot clean up.") return user_dir = os.path.join(TMP_DIR, str(session_hash)) if os.path.exists(user_dir) and os.path.isdir(user_dir): # Extra check if it's a directory try: shutil.rmtree(user_dir) print(f"Ended session, removed directory: {user_dir}") except OSError as e: print(f"Error removing tmp directory {user_dir}: {e.strerror}", file=sys.stderr) else: print(f"Ended session, directory not found or not a directory: {user_dir}") except Exception as e: print(f"Error in end_session cleaning directory '{user_dir}': {e}", file=sys.stderr) def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict: """Packs Gaussian and Mesh data into a serializable dictionary.""" print("[pack_state] Packing state to dictionary...") try: packed_data = { 'gaussian': { **{k: v for k, v in gs.init_params.items()}, '_xyz': gs._xyz.detach().cpu().numpy(), '_features_dc': gs._features_dc.detach().cpu().numpy(), '_scaling': gs._scaling.detach().cpu().numpy(), '_rotation': gs._rotation.detach().cpu().numpy(), '_opacity': gs._opacity.detach().cpu().numpy(), }, 'mesh': { 'vertices': mesh.vertices.detach().cpu().numpy(), 'faces': mesh.faces.detach().cpu().numpy(), }, } print(f"[pack_state] Dictionary created. Keys: {list(packed_data.keys())}, Gaussian points: {len(packed_data['gaussian']['_xyz'])}, Mesh vertices: {len(packed_data['mesh']['vertices'])}") return packed_data except Exception as e: print(f"Error during pack_state: {e}", file=sys.stderr) traceback.print_exc() raise def unpack_state(state_dict: dict) -> Tuple[Gaussian, edict]: """Unpacks Gaussian and Mesh data from a dictionary.""" print("[unpack_state] Unpacking state from dictionary...") try: if not isinstance(state_dict, dict) or 'gaussian' not in state_dict or 'mesh' not in state_dict: raise ValueError("Invalid state_dict structure passed to unpack_state.") device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f"[unpack_state] Using device: {device}") gauss_data = state_dict['gaussian'] mesh_data = state_dict['mesh'] gs = Gaussian( aabb=gauss_data.get('aabb'), sh_degree=gauss_data.get('sh_degree'), mininum_kernel_size=gauss_data.get('mininum_kernel_size'), scaling_bias=gauss_data.get('scaling_bias'), opacity_bias=gauss_data.get('opacity_bias'), scaling_activation=gauss_data.get('scaling_activation'), ) gs._xyz = torch.tensor(gauss_data['_xyz'], device=device, dtype=torch.float32) gs._features_dc = torch.tensor(gauss_data['_features_dc'], device=device, dtype=torch.float32) gs._scaling = torch.tensor(gauss_data['_scaling'], device=device, dtype=torch.float32) gs._rotation = torch.tensor(gauss_data['_rotation'], device=device, dtype=torch.float32) gs._opacity = torch.tensor(gauss_data['_opacity'], device=device, dtype=torch.float32) print(f"[unpack_state] Gaussian unpacked. Points: {gs.get_xyz.shape[0]}") mesh = edict( vertices=torch.tensor(mesh_data['vertices'], device=device, dtype=torch.float32), faces=torch.tensor(mesh_data['faces'], device=device, dtype=torch.int64), ) print(f"[unpack_state] Mesh unpacked. Vertices: {mesh.vertices.shape[0]}, Faces: {mesh.faces.shape[0]}") return gs, mesh except Exception as e: print(f"Error during unpack_state: {e}", file=sys.stderr) traceback.print_exc() raise def get_seed(randomize_seed: bool, seed: int) -> int: """Gets a seed value, randomizing if requested.""" new_seed = np.random.randint(0, MAX_SEED) if randomize_seed else seed print(f"[get_seed] Randomize: {randomize_seed}, Input Seed: {seed}, Output Seed: {new_seed}") return int(new_seed) @spaces.GPU def text_to_3d( prompt: str, seed: int, ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int, req: gr.Request, ) -> Tuple[dict, Optional[str]]: """ Generates a 3D model (Gaussian and Mesh) from text and returns a serializable state dictionary and potentially a video preview path. >>> TEMPORARILY DISABLED VIDEO RENDERING FOR DEBUGGING <<< """ print(f"[text_to_3d - DEBUG MODE] Received prompt: '{prompt}', Seed: {seed}") user_dir = None # Initialize state_dict = None # Initialize try: session_hash = req.session_hash if not session_hash: session_hash = f"no_session_{np.random.randint(10000, 99999)}" print(f"Warning: No session_hash in text_to_3d request, using temporary ID: {session_hash}") # Ensure user directory exists user_dir = os.path.join(TMP_DIR, str(session_hash)) os.makedirs(user_dir, exist_ok=True) print(f"[text_to_3d - DEBUG MODE] User directory: {user_dir}") # --- Generation Pipeline --- print("[text_to_3d - DEBUG MODE] Running Trellis pipeline...") outputs = pipeline.run( prompt=prompt, seed=seed, formats=["gaussian", "mesh"], sparse_structure_sampler_params={ "steps": int(ss_sampling_steps), "cfg_strength": float(ss_guidance_strength), }, slat_sampler_params={ "steps": int(slat_sampling_steps), "cfg_strength": float(slat_guidance_strength), }, ) print("[text_to_3d - DEBUG MODE] Pipeline run completed.") # --- Create Serializable State Dictionary --- state_dict = pack_state(outputs['gaussian'][0], outputs['mesh'][0]) except Exception as e: print(f"❌ [text_to_3d - DEBUG MODE] Error during generation or packing: {e}", file=sys.stderr) traceback.print_exc() # Raise a Gradio error to send failure message back to client if possible raise gr.Error(f"Core generation failed: {e}") # --- Render Video Preview (TEMPORARILY DISABLED FOR DEBUGGING) --- video_path = None print("[text_to_3d - DEBUG MODE] Skipping video rendering.") # --- Original Video Code Block (Keep commented) --- # ... (video code commented out) ... # --- Cleanup and Return --- if torch.cuda.is_available(): torch.cuda.empty_cache() print("[text_to_3d - DEBUG MODE] Cleared CUDA cache.") # --- Return Serializable Dictionary and None Video Path --- print("[text_to_3d - DEBUG MODE] Returning state dictionary and None video path.") if state_dict is None: # This case should ideally be caught by the exception handling above print("Error: state_dict is None before return, generation likely failed.", file=sys.stderr) raise gr.Error("State dictionary creation failed.") return state_dict, video_path @spaces.GPU(duration=120) def extract_glb( state_dict: dict, mesh_simplify: float, texture_size: int, req: gr.Request, ) -> Tuple[str, str]: """ Extracts a GLB file from the provided 3D model state dictionary. """ print(f"[extract_glb] Received request. Simplify: {mesh_simplify}, Texture Size: {texture_size}") user_dir = None # Initialize glb_path = None # Initialize try: session_hash = req.session_hash if not session_hash: session_hash = f"no_session_{np.random.randint(10000, 99999)}" print(f"Warning: No session_hash in extract_glb request, using temporary ID: {session_hash}") if not isinstance(state_dict, dict): print("❌ [extract_glb] Error: Invalid state_dict received (not a dictionary).") raise gr.Error("Invalid state data received. Please generate the model first.") user_dir = os.path.join(TMP_DIR, str(session_hash)) os.makedirs(user_dir, exist_ok=True) print(f"[extract_glb] User directory: {user_dir}") # --- Unpack state from the dictionary --- gs, mesh = unpack_state(state_dict) # --- Postprocessing and Export --- print("[extract_glb] Converting to GLB...") simplify_factor = float(mesh_simplify) tex_size = int(texture_size) glb = postprocessing_utils.to_glb(gs, mesh, simplify=simplify_factor, texture_size=tex_size, verbose=True) glb_path = os.path.join(user_dir, 'sample.glb') print(f"[extract_glb] Exporting GLB to: {glb_path}") glb.export(glb_path) print("[extract_glb] GLB exported successfully.") except Exception as e: print(f"❌ [extract_glb] Error during GLB extraction: {e}", file=sys.stderr) traceback.print_exc() raise gr.Error(f"Failed to extract GLB: {e}") # Propagate error # --- Cleanup and Return --- if torch.cuda.is_available(): torch.cuda.empty_cache() print("[extract_glb] Cleared CUDA cache.") print("[extract_glb] Returning GLB path.") if glb_path is None: print("Error: glb_path is None before return, extraction likely failed.", file=sys.stderr) raise gr.Error("GLB path generation failed.") return glb_path, glb_path @spaces.GPU def extract_gaussian( state_dict: dict, req: gr.Request ) -> Tuple[str, str]: """ Extracts a PLY (Gaussian) file from the provided 3D model state dictionary. """ print("[extract_gaussian] Received request.") user_dir = None # Initialize gaussian_path = None # Initialize try: session_hash = req.session_hash if not session_hash: session_hash = f"no_session_{np.random.randint(10000, 99999)}" print(f"Warning: No session_hash in extract_gaussian request, using temporary ID: {session_hash}") if not isinstance(state_dict, dict): print("❌ [extract_gaussian] Error: Invalid state_dict received (not a dictionary).") raise gr.Error("Invalid state data received. Please generate the model first.") user_dir = os.path.join(TMP_DIR, str(session_hash)) os.makedirs(user_dir, exist_ok=True) print(f"[extract_gaussian] User directory: {user_dir}") # --- Unpack state from the dictionary --- gs, _ = unpack_state(state_dict) # --- Export PLY --- gaussian_path = os.path.join(user_dir, 'sample.ply') print(f"[extract_gaussian] Saving PLY to: {gaussian_path}") gs.save_ply(gaussian_path) print("[extract_gaussian] PLY saved successfully.") except Exception as e: print(f"❌ [extract_gaussian] Error during Gaussian extraction: {e}", file=sys.stderr) traceback.print_exc() raise gr.Error(f"Failed to extract Gaussian PLY: {e}") # Propagate error # --- Cleanup and Return --- if torch.cuda.is_available(): torch.cuda.empty_cache() print("[extract_gaussian] Cleared CUDA cache.") print("[extract_gaussian] Returning PLY path.") if gaussian_path is None: print("Error: gaussian_path is None before return, extraction likely failed.", file=sys.stderr) raise gr.Error("Gaussian PLY path generation failed.") return gaussian_path, gaussian_path # --- Gradio UI Definition --- print("Setting up Gradio Blocks interface...") with gr.Blocks(delete_cache=(600, 600), title="TRELLIS Text-to-3D") as demo: gr.Markdown(""" # Text to 3D Asset with [TRELLIS](https://trellis3d.github.io/) * Type a text prompt and click "Generate" to create a 3D asset preview. * Adjust extraction settings if desired. * Click "Extract GLB" or "Extract Gaussian" to get the downloadable 3D file. *(Note: Video preview is temporarily disabled for debugging)* """) # --- State Buffer --- output_buf = gr.State() with gr.Row(): with gr.Column(scale=1): # Input column text_prompt = gr.Textbox(label="Text Prompt", lines=5, placeholder="e.g., a cute red dragon") with gr.Accordion(label="Generation Settings", open=False): seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1) randomize_seed = gr.Checkbox(label="Randomize Seed", value=True) gr.Markdown("--- \n **Stage 1: Sparse Structure Generation**") with gr.Row(): ss_guidance_strength = gr.Slider(0.0, 15.0, label="Guidance Strength", value=7.5, step=0.1) ss_sampling_steps = gr.Slider(10, 50, label="Sampling Steps", value=25, step=1) gr.Markdown("--- \n **Stage 2: Structured Latent Generation**") with gr.Row(): slat_guidance_strength = gr.Slider(0.0, 15.0, label="Guidance Strength", value=7.5, step=0.1) slat_sampling_steps = gr.Slider(10, 50, label="Sampling Steps", value=25, step=1) generate_btn = gr.Button("Generate 3D Preview", variant="primary") with gr.Accordion(label="GLB Extraction Settings", open=True): mesh_simplify = gr.Slider(0.9, 0.99, label="Simplify Factor", value=0.95, step=0.01, info="Higher value = less simplification (more polys)") texture_size = gr.Slider(512, 2048, label="Texture Size (pixels)", value=1024, step=512, info="Size of the generated texture map") with gr.Row(): extract_glb_btn = gr.Button("Extract GLB", interactive=False) extract_gs_btn = gr.Button("Extract Gaussian (PLY)", interactive=False) gr.Markdown(""" *NOTE: Gaussian file (.ply) can be very large (~50MB+) and may take time to process/download.* """) with gr.Column(scale=1): # Output column # Video component remains for layout but won't show anything in this debug version video_output = gr.Video(label="Generated 3D Preview (DISABLED FOR DEBUG)", autoplay=False, loop=False, value=None, height=350) model_output = gr.Model3D(label="Extracted Model Preview", height=350, clear_color=[0.95, 0.95, 0.95, 1.0]) with gr.Row(): download_glb = gr.DownloadButton(label="Download GLB", interactive=False) download_gs = gr.DownloadButton(label="Download Gaussian (PLY)", interactive=False) # --- Event Handlers --- print("Defining Gradio event handlers...") # Handle session start/end # demo.load() is valid with inputs=None, outputs=None (though default) demo.load(start_session, inputs=None, outputs=None) # >>> FIX: demo.unload() does NOT take inputs/outputs arguments <<< demo.unload(end_session) # Removed inputs/outputs kwargs # --- Generate Button Click Flow --- generate_event = generate_btn.click( get_seed, inputs=[randomize_seed, seed], outputs=[seed], api_name="get_seed" ).then( text_to_3d, inputs=[text_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps], outputs=[output_buf, video_output], # state_dict -> output_buf, None -> video_output api_name="text_to_3d" ).then( lambda: ( gr.Button(interactive=True), gr.Button(interactive=True), gr.DownloadButton(interactive=False), gr.DownloadButton(interactive=False) ), inputs=None, outputs=[extract_glb_btn, extract_gs_btn, download_glb, download_gs], ) # --- Extract GLB Button Click Flow --- extract_glb_event = extract_glb_btn.click( extract_glb, inputs=[output_buf, mesh_simplify, texture_size], outputs=[model_output, download_glb], api_name="extract_glb" ).then( lambda: gr.DownloadButton(interactive=True), inputs=None, outputs=[download_glb], ) # --- Extract Gaussian Button Click Flow --- extract_gs_event = extract_gs_btn.click( extract_gaussian, inputs=[output_buf], outputs=[model_output, download_gs], api_name="extract_gaussian" ).then( lambda: gr.DownloadButton(interactive=True), inputs=None, outputs=[download_gs], ) # --- Clear Download Button Interactivity when model preview is cleared --- model_output.clear( lambda: (gr.DownloadButton(interactive=False), gr.DownloadButton(interactive=False)), inputs=None, outputs=[download_glb, download_gs] ) video_output.clear( lambda: ( gr.Button(interactive=False), gr.Button(interactive=False), gr.DownloadButton(interactive=False), gr.DownloadButton(interactive=False) ), inputs=None, outputs=[extract_glb_btn, extract_gs_btn, download_glb, download_gs], ) print("Gradio interface setup complete.") # --- Launch the Gradio app --- if __name__ == "__main__": print("Loading Trellis pipeline...") pipeline_loaded = False pipeline = None # Initialize try: pipeline = TrellisTextTo3DPipeline.from_pretrained( "JeffreyXiang/TRELLIS-text-xlarge", torch_dtype=torch.float16 # Use float16 if GPU supports it ) if torch.cuda.is_available(): pipeline = pipeline.to("cuda") print("✅ Trellis pipeline loaded successfully to GPU.") else: print("⚠️ WARNING: CUDA not available, running on CPU (will be very slow).") print("✅ Trellis pipeline loaded successfully to CPU.") pipeline_loaded = True except Exception as e: print(f"❌ Failed to load Trellis pipeline: {e}", file=sys.stderr) traceback.print_exc() print("❌ Exiting due to pipeline load failure.") sys.exit(1) if pipeline_loaded: print("Launching Gradio demo...") # Consider increasing queue timeout if tasks are long demo.queue( # default_concurrency_limit=2, # Limit concurrency if resource issues suspected # status_update_rate='auto' ).launch( # server_name="0.0.0.0", # Allows access from local network # share=False, # Set True for public link (careful with resources) debug=True, # Enable Gradio/FastAPI debug logs # prevent_thread_lock=True # Might help sometimes ) print("Gradio demo launched.") else: print("Gradio demo not launched due to pipeline loading failure.")