Update app.py

app.py

@@ -13,11 +13,33 @@ import subprocess
 import shutil
 import base64
 import logging
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse
+from fastapi.middleware.wsgi import WSGIMiddleware
+import uvicorn
+
+# Job status dictionary
+generation_status = {}
+
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
+# FastAPI app to wrap both Gradio and API routes
+api_app = FastAPI()
+
+@api_app.get("/generation_status/{job_id}")
+async def check_generation_status(job_id: str):
+    if job_id in generation_status:
+        return JSONResponse(content=generation_status[job_id])
+    else:
+        return JSONResponse(content={"status": "not_found"}, status_code=404)
+
+
+# Mount the Gradio Blocks demo interface as WSGI
+api_app.mount("/", WSGIMiddleware(demo))
+
 # Install additional dependencies
 try:
     subprocess.run("pip install spandrel==0.4.1 --no-deps", shell=True, check=True)
@@ -384,6 +406,14 @@ def image_to_3d(
         mesh_path = os.path.join(save_dir, f"polygenixai_{get_random_hex()}.glb")
         mesh.export(mesh_path)
         logger.info(f"Saved mesh to {mesh_path}")
+        
+        # Track job status
+        job_id = req.session_hash
+        generation_status[job_id] = {
+            "status": "ready",
+            "mesh_path": mesh_path
+        }
+
 
         torch.cuda.empty_cache()
         return mesh_path
@@ -391,6 +421,8 @@ def image_to_3d(
         logger.error(f"Error in image_to_3d: {str(e)}")
         raise
 
+
+
 @spaces.GPU(duration=5)
 @torch.no_grad()
 def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):