mimic fast api speearet file inegration file

app.py

@@ -17,13 +17,9 @@ from trellis.utils import render_utils, postprocessing_utils
 import traceback
 import sys
 
-# --- FastAPI / Threading Imports ---
-import threading
-import uvicorn
+# --- Import the FastAPI integration module ---
+import trellis_fastAPI_integration
 import logging
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-
 
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
@@ -43,7 +39,11 @@ def start_session(req: gr.Request):
     
 def end_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    shutil.rmtree(user_dir)
+    try:
+        if os.path.exists(user_dir):
+             shutil.rmtree(user_dir)
+    except OSError as e:
+        logger.warning(f"Warning: Could not remove temp session dir {user_dir}: {e}")
 
 
 def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
@@ -63,7 +63,7 @@ def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
     }
     
     
-def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
+def unpack_state(state: dict) -> Tuple[Gaussian, edict]:
     gs = Gaussian(
         aabb=state['gaussian']['aabb'],
         sh_degree=state['gaussian']['sh_degree'],
@@ -119,10 +119,16 @@ def text_to_3d(
         str: The path to the video of the 3D model.
     """
     # --- Determine user_dir robustly ---
-    session_hash_str = str(req.session_hash) if hasattr(req, 'session_hash') and req.session_hash else f"api_call_{np.random.randint(10000)}"
+    session_hash_str = str(req.session_hash) if hasattr(req, 'session_hash') and req.session_hash else f"gradio_call_{np.random.randint(10000)}"
     user_dir = os.path.join(TMP_DIR, session_hash_str)
     os.makedirs(user_dir, exist_ok=True) # Ensure directory exists
 
+    # Use the global pipeline initialized later
+    if pipeline is None:
+        logger.error("Gradio Error: Pipeline not initialized")
+        # Handle error appropriately for Gradio - maybe return None or raise gr.Error?
+        return {}, None
+
     outputs = pipeline.run(
         prompt,
         seed=seed,
@@ -143,9 +149,7 @@ def text_to_3d(
     try:
         imageio.mimsave(video_path, video, fps=15) # Now the directory should exist
     except FileNotFoundError:
-         print(f"ERROR: Directory {user_dir} still not found before mimsave!", file=sys.stderr)
-         # Decide if we should raise or return an error state?
-         # Returning a dummy path might hide the error, so let's raise for now
+         logger.error(f"ERROR: Directory {user_dir} still not found before mimsave!", exc_info=True)
          raise
     state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
     torch.cuda.empty_cache()
@@ -171,7 +175,7 @@ def extract_glb(
         str: The path to the extracted GLB file.
     """
     # --- Determine user_dir robustly ---
-    session_hash_str = str(req.session_hash) if hasattr(req, 'session_hash') and req.session_hash else f"api_call_{np.random.randint(10000)}"
+    session_hash_str = str(req.session_hash) if hasattr(req, 'session_hash') and req.session_hash else f"gradio_call_{np.random.randint(10000)}"
     user_dir = os.path.join(TMP_DIR, session_hash_str)
     os.makedirs(user_dir, exist_ok=True) # Ensure directory exists
 
@@ -181,7 +185,7 @@ def extract_glb(
     try:
         glb.export(glb_path) # Now the directory should exist
     except FileNotFoundError:
-        print(f"ERROR: Directory {user_dir} still not found before glb.export!", file=sys.stderr)
+        logger.error(f"ERROR: Directory {user_dir} still not found before glb.export!", exc_info=True)
         raise
     torch.cuda.empty_cache()
     return glb_path, glb_path
@@ -199,7 +203,7 @@ def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
         str: The path to the extracted Gaussian file.
     """
     # --- Determine user_dir robustly ---
-    session_hash_str = str(req.session_hash) if hasattr(req, 'session_hash') and req.session_hash else f"api_call_{np.random.randint(10000)}"
+    session_hash_str = str(req.session_hash) if hasattr(req, 'session_hash') and req.session_hash else f"gradio_call_{np.random.randint(10000)}"
     user_dir = os.path.join(TMP_DIR, session_hash_str)
     os.makedirs(user_dir, exist_ok=True) # Ensure directory exists
 
@@ -208,84 +212,12 @@ def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
     try:
         gs.save_ply(gaussian_path) # Now the directory should exist
     except FileNotFoundError:
-        print(f"ERROR: Directory {user_dir} still not found before gs.save_ply!", file=sys.stderr)
+        logger.error(f"ERROR: Directory {user_dir} still not found before gs.save_ply!", exc_info=True)
         raise
     torch.cuda.empty_cache()
     return gaussian_path, gaussian_path
 
 
-# --- FastAPI App Setup ---
-api_app = FastAPI()
-
-class GenerateRequest(BaseModel):
-    prompt: str
-    seed: int = 0 # Default seed
-    mesh_simplify: float = 0.95 # Default simplify factor
-    texture_size: int = 1024 # Default texture size
-    # Add other generation parameters if needed
-
-@api_app.post("/api/generate-sync")
-async def generate_sync_api(request_data: GenerateRequest):
-    global pipeline # Access the globally initialized pipeline
-    if pipeline is None:
-        logger.error("API Error: Pipeline not initialized")
-        raise HTTPException(status_code=503, detail="Pipeline not ready")
-
-    prompt = request_data.prompt
-    seed = request_data.seed
-    mesh_simplify = request_data.mesh_simplify
-    texture_size = request_data.texture_size
-    # Extract other params if added to GenerateRequest
-
-    logger.info(f"API /generate-sync received prompt: {prompt}")
-
-    try:
-        # --- Determine a unique temporary directory for this API call ---
-        api_call_hash = f"api_sync_{np.random.randint(100000)}"
-        user_dir = os.path.join(TMP_DIR, api_call_hash)
-        os.makedirs(user_dir, exist_ok=True)
-        logger.info(f"API using temp dir: {user_dir}")
-
-        # --- Stage 1: Run the text-to-3D pipeline ---
-        logger.info("API running pipeline...")
-        # Use default values for parameters not exposed in the simple API for now
-        outputs = pipeline.run(
-            prompt,
-            seed=seed,
-            formats=["gaussian", "mesh"],
-            sparse_structure_sampler_params={"steps": 25, "cfg_strength": 7.5},
-            slat_sampler_params={"steps": 25, "cfg_strength": 7.5},
-        )
-        gs = outputs['gaussian'][0]
-        mesh = outputs['mesh'][0]
-        logger.info("API pipeline finished.")
-        torch.cuda.empty_cache()
-
-        # --- Stage 2: Extract GLB ---
-        logger.info("API extracting GLB...")
-        glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
-        glb_path = os.path.join(user_dir, 'generated_sync.glb')
-        glb.export(glb_path)
-        logger.info(f"API GLB exported to: {glb_path}")
-        torch.cuda.empty_cache()
-
-        # Return the absolute path within the container
-        return {"status": "success", "glb_path": os.path.abspath(glb_path)}
-
-    except Exception as e:
-        logger.error(f"API /generate-sync error: {str(e)}", exc_info=True)
-        # Clean up temp dir on error if it exists
-        if os.path.exists(user_dir):
-            try:
-                shutil.rmtree(user_dir)
-            except Exception as cleanup_e:
-                logger.error(f"API Error cleaning up dir {user_dir}: {cleanup_e}")
-        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
-    # Note: We don't automatically clean up the user_dir on success,
-    # as the file needs to be accessible for download by the calling server.
-    # A separate cleanup mechanism might be needed eventually.
-
-
 # --- Gradio Blocks Definition ---
 with gr.Blocks(delete_cache=(600, 600)) as demo:
     gr.Markdown("""
@@ -379,29 +311,16 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     )
     
 
-# --- Functions to Run FastAPI in Background ---
-def run_api():
-    """Run the FastAPI server."""
-    uvicorn.run(api_app, host="0.0.0.0", port=8000) # Run on port 8000
-
-def start_api_thread():
-    """Start the API server in a background thread."""
-    api_thread = threading.Thread(target=run_api, daemon=True)
-    api_thread.start()
-    logger.info("Started FastAPI server thread on port 8000")
-    return api_thread
-
-
 # Launch the Gradio app and FastAPI server
 if __name__ == "__main__":
     logger.info("Initializing Trellis Pipeline...")
-    # Make pipeline global so API endpoint can access it
+    # Make pipeline global so Gradio functions and API endpoint can access it
     pipeline = TrellisTextTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-text-xlarge")
     pipeline.cuda()
     logger.info("Trellis Pipeline Initialized.")
 
-    # Start the background API server
-    start_api_thread()
+    # Start the background API server using the integration module
+    trellis_fastAPI_integration.start_api_thread(pipeline)
 
     # Launch the Gradio interface (blocking call)
     logger.info("Launching Gradio Demo...")