# trellis_fastAPI_integration.py # Version: 1.0.0 # a.1 Imports and Initial Setup import os import shutil import threading import uvicorn import logging import numpy as np import torch from fastapi import FastAPI, HTTPException from pydantic import BaseModel from easydict import EasyDict as edict # Assuming EasyDict might be needed if state used # Assuming these are available or installed correctly in the environment from trellis.utils import postprocessing_utils # We get the pipeline object passed in, so no direct import needed here # Set up logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # FastAPI app api_app = FastAPI() # --- Temporary Directory --- (Consistent with appTrellis.py) TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp') os.makedirs(TMP_DIR, exist_ok=True) # b.1 Request/Response Models 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 (e.g., guidance, steps) # c.1 API Endpoint for Synchronous Generation @api_app.post("/generate-sync") async def generate_sync_api(request_data: GenerateRequest): """API endpoint to synchronously generate a model and return the GLB path.""" logger.info("API /generate-sync endpoint hit.") # Log when endpoint is called # Access the pipeline object stored in app state pipeline = api_app.state.pipeline if pipeline is None: logger.error("API Error: Pipeline not initialized or passed correctly") 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 ss_sampling_steps = 25 # Example default ss_guidance_strength = 7.5 # Example default slat_sampling_steps = 25 # Example default slat_guidance_strength = 7.5 # Example default logger.info(f"API /generate-sync received prompt: {prompt}") user_dir = None # Define user_dir outside try for cleanup try: # --- Determine a unique temporary directory for this API call --- # Using a simpler random hash name for the API call directory 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...") # Ensure pipeline is run with appropriate parameters outputs = pipeline.run( prompt, seed=seed, formats=["gaussian", "mesh"], sparse_structure_sampler_params={ "steps": ss_sampling_steps, "cfg_strength": ss_guidance_strength, }, slat_sampler_params={ "steps": slat_sampling_steps, "cfg_strength": slat_guidance_strength, }, ) gs = outputs['gaussian'][0] # Get the Gaussian representation mesh = outputs['mesh'][0] # Get the Mesh representation logger.info("API pipeline finished.") torch.cuda.empty_cache() # --- Stage 2: Extract GLB --- logger.info("API extracting GLB...") # Use the postprocessing utility 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 # This path needs to be accessible via the /file= route from outside absolute_glb_path = os.path.abspath(glb_path) logger.info(f"API returning absolute path: {absolute_glb_path}") return {"status": "success", "glb_path": absolute_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 and was created if user_dir and os.path.exists(user_dir): try: shutil.rmtree(user_dir) logger.info(f"API cleaned up failed directory: {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. # d.1 API Server Setup Functions def run_api(): """Run the FastAPI server.""" logger.info("FastAPI Integration: run_api function called.") # Ensure pipeline is available in app state before starting if not hasattr(api_app.state, 'pipeline') or api_app.state.pipeline is None: logger.error("FastAPI Integration: Cannot start API server - Pipeline object not found in app state.") return logger.info("FastAPI Integration: Pipeline object found in state. Attempting to start Uvicorn...") # Run on port 8000 - ensure this doesn't conflict if Gradio also tries this port try: uvicorn.run(api_app, host="0.0.0.0", port=8000) logger.info("FastAPI Integration: Uvicorn server stopped.") # Logged when server exits cleanly except Exception as e: logger.error(f"FastAPI Integration: Uvicorn server failed to run or crashed: {e}", exc_info=True) def start_api_thread(pipeline_object): """Start the API server in a background thread Args: pipeline_object: The initialized TrellisTextTo3DPipeline object """ logger.info("FastAPI Integration: start_api_thread called.") # Store the passed pipeline object in the app's state if pipeline_object is None: logger.error("FastAPI Integration: start_api_thread received a None pipeline_object. Aborting thread start.") return None try: api_app.state.pipeline = pipeline_object logger.info("FastAPI Integration: Pipeline object successfully stored in app state.") except Exception as e: logger.error(f"FastAPI Integration: Failed to store pipeline object in app state: {e}", exc_info=True) return None logger.info("FastAPI Integration: Creating API thread...") api_thread = threading.Thread(target=run_api, daemon=True) logger.info("FastAPI Integration: Attempting to start API thread...") try: api_thread.start() logger.info("FastAPI Integration: API thread started (start() method called).") except Exception as e: logger.error(f"FastAPI Integration: Failed to start API thread: {e}", exc_info=True) return None # Indicate thread failed to start logger.info("Started Trellis FastAPI integration server thread function finished.") # Confirms this function completed return api_thread