Update app.py

app.py

@@ -11,9 +11,15 @@ 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
+import uvicorn
 
-# install others
+# 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 +51,30 @@ 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")  # Set in Hugging Face Space secrets
+
+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
+
+# Existing Gradio header (unchanged)
 HEADER = """
 # 🌌 PolyGenixAI: Craft 3D Worlds with Cosmic Precision
 ## Unleash Infinite Creativity with AI-Powered 3D Generation by AnvilInteractive Solutions
@@ -128,7 +158,7 @@ HEADER = """
 </style>
 """
 
-# triposg
+# TripoSG and MV-Adapter setup (unchanged)
 from image_process import prepare_image
 from briarmbg import BriaRMBG
 snapshot_download("briaai/RMBG-1.4", local_dir=RMBG_PRETRAINED_MODEL)
@@ -138,7 +168,6 @@ from triposg.pipelines.pipeline_triposg import TripoSGPipeline
 snapshot_download("VAST-AI/TripoSG", local_dir=TRIPOSG_PRETRAINED_MODEL)
 triposg_pipe = TripoSGPipeline.from_pretrained(TRIPOSG_PRETRAINED_MODEL).to(DEVICE, DTYPE)
 
-# mv adapter
 NUM_VIEWS = 6
 from inference_ig2mv_sdxl import prepare_pipeline, preprocess_image, remove_bg
 from mvadapter.utils import get_orthogonal_camera, tensor_to_image, make_image_grid
@@ -224,7 +253,6 @@ def run_full(image: str, req: gr.Request):
     torch.cuda.empty_cache()
 
     height, width = 768, 768
-    # Prepare cameras
     cameras = get_orthogonal_camera(
         elevation_deg=[0, 0, 0, 0, 89.99, -89.99],
         distance=[1.8] * NUM_VIEWS,
@@ -307,7 +335,6 @@ def run_full(image: str, req: gr.Request):
     )
 
     return image_seg, mesh_path, textured_glb_path
-    
 
 @spaces.GPU()
 @torch.no_grad()
@@ -353,7 +380,6 @@ def image_to_3d(
 @torch.no_grad()
 def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
     height, width = 768, 768
-    # Prepare cameras
     cameras = get_orthogonal_camera(
         elevation_deg=[0, 0, 0, 0, 89.99, -89.99],
         distance=[1.8] * NUM_VIEWS,
@@ -421,7 +447,7 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
     from texture import TexturePipeline, ModProcessConfig
     texture_pipe = TexturePipeline(
         upscaler_ckpt_path="checkpoints/RealESRGAN_x2plus.pth",
-        inpaint_ckpt_path="checkpoints/big-lama.pt",
+        inpaint_cckpt_path="checkpoints/big-lama.pt",
         device=DEVICE,
     )
 
@@ -438,7 +464,33 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
 
     return textured_glb_path
 
-
+# FastAPI endpoint for generating 3D models
+@spaces.GPU(duration=180)
+@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)
+
+# Gradio interface (unchanged)
 with gr.Blocks(title="PolyGenixAI", css="body { background-color: #1A1A1A; } .gr-panel { background-color: #2D2D2D; }") as demo:
     gr.Markdown(HEADER)
 
@@ -544,4 +596,6 @@ with gr.Blocks(title="PolyGenixAI", css="body { background-color: #1A1A1A; } .gr
     demo.load(start_session)
     demo.unload(end_session)
 
-demo.launch()
+# Run both Gradio and FastAPI
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)