Update app.py

app.py

@@ -16,10 +16,11 @@ from fastapi.security import APIKeyHeader
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 import uvicorn
+import asyncio
 
 # 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)
+subprocess.run("pip install fastapi==0.116.1 uvicorn==0.31.0", shell=True, check=True)
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.float16
@@ -59,7 +60,7 @@ 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
+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:
@@ -74,6 +75,37 @@ class GenerateRequest(BaseModel):
     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"}
+
+# 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)
+
+# Gradio-related functions (unchanged)
 HEADER = """
 # 🌌 PolyGenixAI: Craft 3D Worlds with Cosmic Precision
 ## Unleash Infinite Creativity with AI-Powered 3D Generation by AnvilInteractive Solutions
@@ -184,9 +216,8 @@ mv_adapter_pipe = prepare_pipeline(
     dtype=torch.float16,
 )
 birefnet = AutoModelForImageSegmentation.from_pretrained(
-        "ZhengPeng7/BiRefNet", trust_remote_code=True
-    )
-birefnet.to(DEVICE)
+    "ZhengPeng7/BiRefNet", trust_remote_code=True
+).to(DEVICE)
 transform_image = transforms.Compose(
     [
         transforms.Resize((1024, 1024)),
@@ -199,7 +230,7 @@ remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, DEVICE)
 if not os.path.exists("checkpoints/RealESRGAN_x2plus.pth"):
     hf_hub_download("dtarnow/UPscaler", filename="RealESRGAN_x2plus.pth", local_dir="checkpoints")
 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)
+    hf_hub_download("Sanster/models", filename="big-lama.pt", local_dir="checkpoints")
 
 def start_session(req: gr.Request):
     save_dir = os.path.join(TMP_DIR, str(req.session_hash))
@@ -253,7 +284,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,
@@ -337,31 +367,6 @@ 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)
-
 @spaces.GPU()
 @torch.no_grad()
 def run_segmentation(image: str):
@@ -406,7 +411,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,
@@ -491,14 +495,9 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
 
     return textured_glb_path
 
-# Gradio interface (unchanged)
+# Gradio interface
 with gr.Blocks(title="PolyGenixAI", css="body { background-color: #1A1A1A; } .gr-panel { background-color: #2D2D2D; }") as demo:
     gr.Markdown(HEADER)
-
-    @app.get("/api/test")
-    async def test_endpoint():
-        return {"message": "FastAPI is running"}
-
     with gr.Tabs(elem_classes="gr-tab"):
         with gr.Tab("Create 3D Model"):
             with gr.Row():
@@ -551,12 +550,10 @@ with gr.Blocks(title="PolyGenixAI", css="body { background-color: #1A1A1A; } .gr
                         )
                         gen_button = gr.Button("Generate 3D Model", variant="primary", elem_classes="gr-button-primary")
                         gen_texture_button = gr.Button("Apply Texture", variant="secondary", interactive=False, elem_classes="gr-button-secondary")
-
                 with gr.Column(scale=1):
                     model_output = gr.Model3D(label="3D Model Preview", interactive=False, height=400, elem_classes="gr-panel")
                     textured_model_output = gr.Model3D(label="Textured 3D Model", interactive=False, height=400, elem_classes="gr-panel")
                     download_button = gr.Button("Download GLB", variant="secondary", elem_classes="gr-button-secondary")
-
         with gr.Tab("Cosmic Gallery"):
             gr.Markdown("### Discover Stellar Creations")
             gr.Examples(
@@ -570,7 +567,6 @@ with gr.Blocks(title="PolyGenixAI", css="body { background-color: #1A1A1A; } .gr
                 cache_examples=True,
             )
             gr.Markdown("Connect with creators in our <a href='https://www.anvilinteractive.com/community' style='color: #A78BFA; text-decoration: none;'>PolyGenixAI Cosmic Community</a>!")
-
     gen_button.click(
         run_segmentation,
         inputs=[image_prompts],
@@ -591,16 +587,25 @@ with gr.Blocks(title="PolyGenixAI", css="body { background-color: #1A1A1A; } .gr
         ],
         outputs=[model_output]
     ).then(lambda: gr.Button(interactive=True), outputs=[gen_texture_button])
-
     gen_texture_button.click(
         run_texture,
         inputs=[image_prompts, model_output, seed],
         outputs=[textured_model_output]
     )
-
     demo.load(start_session)
     demo.unload(end_session)
 
-# Run both Gradio and FastAPI
+# Run FastAPI and Gradio concurrently
+async def run_servers():
+    # Start FastAPI with Uvicorn
+    config = uvicorn.Config(app=app, host="0.0.0.0", port=8000)
+    server = uvicorn.Server(config)
+    fastapi_task = asyncio.create_task(server.serve())
+    
+    # Start Gradio
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+    
+    await fastapi_task
+
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    asyncio.run(run_servers())