Update app.py

app.py

@@ -11,9 +11,14 @@ from torchvision import transforms
 from huggingface_hub import hf_hub_download, snapshot_download
 import subprocess
 import shutil
+from fastapi import FastAPI, HTTPException, Depends, File, UploadFile
+from fastapi.security import APIKeyHeader
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
 
 # Install additional dependencies
 subprocess.run("pip install spandrel==0.4.1 --no-deps", shell=True, check=True)
+subprocess.run("pip install fastapi uvicorn", shell=True, check=True)
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.float16
@@ -45,6 +50,34 @@ sys.path.append(os.path.join(TRIPOSG_CODE_DIR, "scripts"))
 sys.path.append(MV_ADAPTER_CODE_DIR)
 sys.path.append(os.path.join(MV_ADAPTER_CODE_DIR, "scripts"))
 
+# Initialize FastAPI app
+app = FastAPI()
+
+# Mount static files for serving generated models
+app.mount("/files", StaticFiles(directory=TMP_DIR), name="files")
+
+# API key authentication
+api_key_header = APIKeyHeader(name="X-API-Key")
+VALID_API_KEY = os.getenv("POLYGENIX_API_KEY", "your-secret-api-key")
+
+async def verify_api_key(api_key: str = Depends(api_key_header)):
+    if api_key != VALID_API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+    return api_key
+
+# API request model
+class GenerateRequest(BaseModel):
+    seed: int = 0
+    num_inference_steps: int = 50
+    guidance_scale: float = 7.5
+    simplify: bool = True
+    target_face_num: int = DEFAULT_FACE_NUMBER
+
+# Test endpoint
+@app.get("/api/test")
+async def test_endpoint():
+    return {"message": "FastAPI is running"}
+
 # triposg
 from image_process import prepare_image
 from briarmbg import BriaRMBG
@@ -88,27 +121,13 @@ if not os.path.exists("checkpoints/RealESRGAN_x2plus.pth"):
 if not os.path.exists("checkpoints/big-lama.pt"):
     subprocess.run("wget -P checkpoints/ https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt", shell=True, check=True)
 
-def start_session(req: gr.Request):
-    save_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    os.makedirs(save_dir, exist_ok=True)
-    print("start session, mkdir", save_dir)
-
-def end_session(req: gr.Request):
-    save_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    shutil.rmtree(save_dir)
-
 def get_random_hex():
     random_bytes = os.urandom(8)
     random_hex = random_bytes.hex()
     return random_hex
 
-def get_random_seed(randomize_seed, seed):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
-
 @spaces.GPU(duration=180)
-def run_full(image: str, req: gr.Request):
+def run_full(image: str, req=None):
     seed = 0
     num_inference_steps = 50
     guidance_scale = 7.5
@@ -223,6 +242,45 @@ def run_full(image: str, req: gr.Request):
 
     return image_seg, mesh_path, textured_glb_path
 
+# FastAPI endpoint for generating 3D models
+@app.post("/api/generate")
+async def generate_3d_model(request: GenerateRequest, image: UploadFile = File(...), api_key: str = Depends(verify_api_key)):
+    try:
+        # Save uploaded image to temporary directory
+        session_hash = get_random_hex()
+        save_dir = os.path.join(TMP_DIR, session_hash)
+        os.makedirs(save_dir, exist_ok=True)
+        image_path = os.path.join(save_dir, f"input_{get_random_hex()}.png")
+        with open(image_path, "wb") as f:
+            f.write(await image.read())
+
+        # Run the full pipeline
+        image_seg, mesh_path, textured_glb_path = run_full(image_path, req=None)
+
+        # Return the file URL for the textured GLB
+        file_url = f"/files/{session_hash}/{os.path.basename(textured_glb_path)}"
+        return {"file_url": file_url}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+    finally:
+        # Clean up temporary directory
+        if os.path.exists(save_dir):
+            shutil.rmtree(save_dir)
+
+def start_session(req: gr.Request):
+    save_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    os.makedirs(save_dir, exist_ok=True)
+    print("start session, mkdir", save_dir)
+
+def end_session(req: gr.Request):
+    save_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    shutil.rmtree(save_dir)
+
+def get_random_seed(randomize_seed, seed):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
 @spaces.GPU()
 @torch.no_grad()
 def run_segmentation(image: str):
@@ -327,7 +385,6 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
 
     torch.cuda.empty_cache()
 
-    save_dir = os.path.join(TMP_DIR, str(req.session_hash))
     mv_image_path = os.path.join(save_dir, f"polygenixai_mv_{get_random_hex()}.png")
     make_image_grid(images, rows=1).save(mv_image_path)
 
@@ -534,5 +591,9 @@ with gr.Blocks(title="PolyGenixAI", css="body { background-color: #1A1A1A; } .gr
     demo.load(start_session)
     demo.unload(end_session)
 
+# Mount Gradio to FastAPI
+app = gr.mount_gradio_app(app, demo, path="/")
+
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)