added logging

app.py

@@ -4,9 +4,6 @@ import spaces
 import os
 import shutil
 import json
-os.environ['TOKENIZERS_PARALLELISM'] = 'true'
-os.environ['SPCONV_ALGO'] = 'native'
-from typing import *
 import torch
 import numpy as np
 import imageio
@@ -14,11 +11,22 @@ 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):
@@ -259,24 +267,23 @@ def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
     return gaussian_path, gaussian_path
 
 
-# --- NEW COMBINED API FUNCTION ---
-@spaces.GPU(duration=120) # Allow more time for combined generation + extraction
+# --- NEW COMBINED API FUNCTION (with HEAVY logging) ---
+@spaces.GPU(duration=120)
 def generate_and_extract_glb(
-    # Inputs mirror text_to_3d and extract_glb settings
     prompt: str,
     seed: int,
     ss_guidance_strength: float,
     ss_sampling_steps: int,
     slat_guidance_strength: float,
     slat_sampling_steps: int,
-    mesh_simplify: float, # Added from extract_glb
-    texture_size: int,    # Added from extract_glb
-    req: gr.Request, # Keep req for potential session info if needed
-) -> Optional[str]: # MODIFIED: Explicitly show it can return None
+    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.
@@ -286,37 +293,65 @@ def generate_and_extract_glb(
         slat_sampling_steps (int): Structured latent steps.
         mesh_simplify (float): Mesh simplification factor for GLB.
         texture_size (int): Texture resolution for GLB.
-        req (gr.Request): Gradio request object.
-        
+        req (Optional[gr.Request]): Gradio request object.
     Returns:
-        Optional[str]: The absolute path to the generated GLB file within the Space's filesystem,
-                       or None if any step fails.
+        Optional[str]: Path to the generated GLB file or None on failure.
     """
-    session_hash = "API_CALL" # Use a generic identifier for API calls if req is None or lacks session
+    # --- 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
+        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: ERROR creating directory {user_dir}: {e}")
-        return None # Cannot proceed without directory
+        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: ===== generate_and_extract_glb START ====")
-    print(f"[{session_hash}] API: Prompt: '{prompt}', Seed: {seed}, Simplify: {mesh_simplify}, Texture: {texture_size}")
+    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 (adapted from text_to_3d) ---
+    # --- Step 1: Generate 3D Model --- 
+    print(f"\n[{session_hash}] API: --- Starting Step 1: Generation Pipeline --- ")
     try:
-        print(f"[{session_hash}] API: Step 1 - Running generation pipeline...")
-        t_start_gen = time.time() # Add timing
+        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"], # Need both for GLB extraction
+            formats=["gaussian", "mesh"],
             sparse_structure_sampler_params={
                 "steps": ss_sampling_steps,
                 "cfg_strength": ss_guidance_strength,
@@ -326,87 +361,167 @@ def generate_and_extract_glb(
                 "cfg_strength": slat_guidance_strength,
             },
         )
+        # --- End pipeline call ---
         t_end_gen = time.time()
-        print(f"[{session_hash}] API: Step 1 - Generation pipeline completed in {t_end_gen - t_start_gen:.2f}s.")
+        print(f"[{session_hash}] API: Step 1 - pipeline.run() completed in {t_end_gen - t_start_gen:.2f}s.")
         
-        # Validate outputs immediately
-        if not outputs or 'gaussian' not in outputs or not outputs['gaussian'] or 'mesh' not in outputs or not outputs['mesh']:
-            print(f"[{session_hash}] API: ERROR - Pipeline output is missing expected keys or values.")
-            return None
+        # === 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 or mesh_output is None:
-             print(f"[{session_hash}] API: ERROR - Pipeline returned None for gs_output or mesh_output.")
-             return None
+        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.")
              
-        print(f"[{session_hash}] API: Step 1 - Outputs obtained (gs type: {type(gs_output)}, mesh type: {type(mesh_output)}).")
-        
-    except Exception as e:
-        print(f"[{session_hash}] API: ERROR during generation pipeline step: {e}")
-        # Print detailed traceback
-        traceback.print_exc()
-        # Clean up CUDA memory before returning
-        try: 
-            torch.cuda.empty_cache()
-            print(f"[{session_hash}] API: CUDA cache cleared after generation error.")
-        except Exception as cache_e:
-            print(f"[{session_hash}] API: Error clearing CUDA cache after generation error: {cache_e}")
-        return None # Return None on failure
-        
-    # --- Step 2: Extract GLB (adapted from extract_glb) ---
-    glb_path = None # Initialize glb_path
-    try:
-        print(f"[{session_hash}] API: Step 2 - Extracting GLB (simplify={mesh_simplify}, texture={texture_size})...")
-        # Check if inputs from previous step are valid
-        if gs_output is None or mesh_output is None:
-             print(f"[{session_hash}] API: ERROR - Cannot proceed with GLB extraction, gs_output or mesh_output is None.")
-             return None
+        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.")
              
-        t_start_glb = time.time()
-        # Directly use the outputs from the pipeline
-        glb = postprocessing_utils.to_glb(gs_output, mesh_output, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
-        t_end_glb = time.time()
-        print(f"[{session_hash}] API: Step 2 - GLB object created in {t_end_glb - t_start_glb:.2f}s.")
+        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 glb is None:
-            print(f"[{session_hash}] API: ERROR - postprocessing_utils.to_glb returned None.")
-            return None
+        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")
             
-        glb_path = os.path.join(user_dir, f'api_generated_{session_hash}.glb') # Use unique name
-        print(f"[{session_hash}] API: Step 2 - Saving GLB to {glb_path}...")
-        t_start_save = time.time()
-        glb.export(glb_path)
-        t_end_save = time.time()
-        print(f"[{session_hash}] API: Step 2 - GLB saved in {t_end_save - t_start_save:.2f}s.")
-        print(f"[{session_hash}] API: Step 2 - GLB extraction completed successfully.")
-        
-    except Exception as e:
-        print(f"[{session_hash}] API: ERROR during GLB extraction step: {e}")
-        # Print detailed traceback
+    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()
-        # Clean up CUDA memory before returning
+        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: CUDA cache cleared after extraction error.")
-        except Exception as cache_e:
-            print(f"[{session_hash}] API: Error clearing CUDA cache after extraction error: {cache_e}")
-        return None # Return None on failure
+            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 ---
+    # --- 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:
-        print(f"[{session_hash}] API: Error clearing final CUDA cache: {cache_e}")
+    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: ===== generate_and_extract_glb END (Success) ===== Returning GLB path: {glb_path}")
-        return glb_path # Return only the path to the generated GLB
+        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: ===== generate_and_extract_glb END (Failure) ===== GLB path not generated or does not exist.")
-        return None # Ensure None is returned if glb_path wasn't set or file doesn't exist
+        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 ---
 
 
@@ -535,40 +650,21 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         outputs=[download_glb, download_gs], # Disable both download buttons
     )
     
-    # --- NEW API ENDPOINT DEFINITION ---
-    # Define the combined function as an API endpoint. 
-    # This is *separate* from the UI button clicks.
-    # It directly calls the combined function.
-    demo.load(
-        None, # No function needed on load for this endpoint
-        inputs=[
-            text_prompt, # Map inputs from API request data based on order
-            seed,
-            ss_guidance_strength,
-            ss_sampling_steps,
-            slat_guidance_strength,
-            slat_sampling_steps,
-            mesh_simplify,
-            texture_size
-        ],
-        outputs=None, # Output is handled by the function return for the API
-        api_name="generate_and_extract_glb" # Assign the specific API name
-    )
-
 # --- Launch the Gradio app ---
 if __name__ == "__main__":
-    print("Loading Trellis pipeline...")
-    # Consider adding error handling for pipeline loading
+    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 successfully.")
+        print("Pipeline loaded and moved to CUDA successfully.")
     except Exception as e:
-        print(f"Error loading pipeline: {e}")
-        # Optionally exit or provide a fallback UI
+        print(f"FATAL ERROR initializing pipeline: {e}")
+        traceback.print_exc()
+        # Optionally exit if pipeline loading fails
         sys.exit(1) 
         
-    print("Launching Gradio demo...")
-    # Enable queue for handling multiple users/requests
-    # Set share=True if you need a public link (requires login for private spaces)
+    print("Launching Gradio demo with queue enabled...")
     demo.queue().launch()