custom api endpoints

app.py

@@ -17,11 +17,24 @@ from trellis.utils import render_utils, postprocessing_utils
 import traceback
 import sys
 
+# --- FastAPI / Threading Imports ---
+import threading
+import uvicorn
+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')
 os.makedirs(TMP_DIR, exist_ok=True)
 
+# --- Global Pipeline Variable ---
+pipeline = None
+
+# --- Logging Setup ---
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 def start_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
@@ -201,6 +214,80 @@ def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
     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")
+@spaces.GPU(duration=300) # Allow longer duration for API calls
+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("""
     ## Text to 3D Asset with [TRELLIS](https://trellis3d.github.io/)
@@ -293,8 +380,30 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     )
     
 
-# Launch the Gradio app
+# --- 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
     pipeline = TrellisTextTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-text-xlarge")
     pipeline.cuda()
+    logger.info("Trellis Pipeline Initialized.")
+
+    # Start the background API server
+    start_api_thread()
+
+    # Launch the Gradio interface (blocking call)
+    logger.info("Launching Gradio Demo...")
     demo.launch()