Update app.py

app.py

@@ -15,6 +15,7 @@ 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 additional dependencies
 subprocess.run("pip install spandrel==0.4.1 --no-deps", shell=True, check=True)
@@ -58,7 +59,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")
+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:
@@ -73,10 +74,88 @@ 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"}
+HEADER = """
+# 🌌 PolyGenixAI: Craft 3D Worlds with Cosmic Precision
+## Unleash Infinite Creativity with AI-Powered 3D Generation by AnvilInteractive Solutions
+<p style="font-size: 1.1em; color: #A78BFA;">By <a href="https://www.anvilinteractive.com/" style="color: #A78BFA; text-decoration: none; font-weight: bold;">AnvilInteractive Solutions</a></p>
+## 🚀 Launch Your Creation:
+1. **Upload an Image** (clear, single-object images shine brightest)
+2. **Choose a Style Filter** to infuse your unique vision
+3. Click **Generate 3D Model** to sculpt your mesh
+4. Click **Apply Texture** to bring your model to life
+5. **Download GLB** to share your masterpiece
+<p style="font-size: 0.9em; margin-top: 10px; color: #D1D5DB;">Powered by cutting-edge AI and multi-view technology from AnvilInteractive Solutions. Join our <a href="https://www.anvilinteractive.com/community" style="color: #A78BFA; text-decoration: none;">PolyGenixAI Community</a> to connect with creators and spark inspiration.</p>
+<style>
+  @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');
+  body {
+    background-color: #1A1A1A !important;
+    font-family: 'Inter', sans-serif !important;
+    color: #D1D5DB !important;
+  }
+  .gr-panel {
+    background-color: #2D2D2D !important;
+    border: 1px solid #7C3AED !important;
+    border-radius: 12px !important;
+    padding: 20px !important;
+    box-shadow: 0 4px 10px rgba(124, 58, 237, 0.2) !important;
+  }
+  .gr-button-primary {
+    background: linear-gradient(45deg, #7C3AED, #A78BFA) !important;
+    color: white !important;
+    border: none !important;
+    border-radius: 8px !important;
+    padding: 12px 24px !important;
+    font-weight: 600 !important;
+    transition: transform 0.2s, box-shadow 0.2s !important;
+  }
+  .gr-button-primary:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 4px 12px rgba(124, 58, 237, 0.5) !important;
+  }
+  .gr-button-secondary {
+    background-color: #4B4B4B !important;
+    color: #D1D5DB !important;
+    border: 1px solid #A78BFA !important;
+    border-radius: 8px !important;
+    padding: 10px 20px !important;
+    transition: transform 0.2s !important;
+  }
+  .gr-button-secondary:hover {
+    transform: translateY(-1px) !important;
+    background-color: #6B6B6B !important;
+  }
+  .gr-accordion {
+    background-color: #2D2D2D !important;
+    border-radius: 8px !important;
+    border: 1px solid #7C3AED !important;
+  }
+  .gr-tab {
+    background-color: #2D2D2D !important;
+    color: #A78BFA !important;
+    border: 1px solid #7C3AED !important;
+    border-radius: 8px !important;
+    margin: 5px !important;
+  }
+  .gr-tab:hover, .gr-tab-selected {
+    background: linear-gradient(45deg, #7C3AED, #A78BFA) !important;
+    color: white !important;
+  }
+  .gr-slider input[type=range]::-webkit-slider-thumb {
+    background-color: #7C3AED !important;
+    border: 2px solid #A78BFA !important;
+  }
+  .gr-dropdown {
+    background-color: #2D2D2D !important;
+    color: #D1D5DB !important;
+    border: 1px solid #A78BFA !important;
+    border-radius: 8px !important;
+  }
+  h1, h3 {
+    color: #A78BFA !important;
+    text-shadow: 0 0 10px rgba(124, 58, 237, 0.5) !important;
+  }
+</style>
+"""
 
 # triposg
 from image_process import prepare_image
@@ -105,8 +184,9 @@ mv_adapter_pipe = prepare_pipeline(
     dtype=torch.float16,
 )
 birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "ZhengPeng7/BiRefNet", trust_remote_code=True
-).to(DEVICE)
+        "ZhengPeng7/BiRefNet", trust_remote_code=True
+    )
+birefnet.to(DEVICE)
 transform_image = transforms.Compose(
     [
         transforms.Resize((1024, 1024)),
@@ -121,13 +201,27 @@ 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=None):
+def run_full(image: str, req: gr.Request):
     seed = 0
     num_inference_steps = 50
     guidance_scale = 7.5
@@ -159,6 +253,7 @@ def run_full(image: str, req=None):
     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,
@@ -267,20 +362,6 @@ async def generate_3d_model(request: GenerateRequest, image: UploadFile = File(.
         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):
@@ -325,6 +406,7 @@ 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,
@@ -385,6 +467,7 @@ 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)
 
@@ -408,92 +491,14 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
 
     return textured_glb_path
 
-HEADER = """
-# 🌌 PolyGenixAI: Craft 3D Worlds with Cosmic Precision
-## Unleash Infinite Creativity with AI-Powered 3D Generation by AnvilInteractive Solutions
-<p style="font-size: 1.1em; color: #A78BFA;">By <a href="https://www.anvilinteractive.com/" style="color: #A78BFA; text-decoration: none; font-weight: bold;">AnvilInteractive Solutions</a></p>
-## 🚀 Launch Your Creation:
-1. **Upload an Image** (clear, single-object images shine brightest)
-2. **Choose a Style Filter** to infuse your unique vision
-3. Click **Generate 3D Model** to sculpt your mesh
-4. Click **Apply Texture** to bring your model to life
-5. **Download GLB** to share your masterpiece
-<p style="font-size: 0.9em; margin-top: 10px; color: #D1D5DB;">Powered by cutting-edge AI and multi-view technology from AnvilInteractive Solutions. Join our <a href="https://www.anvilinteractive.com/community" style="color: #A78BFA; text-decoration: none;">PolyGenixAI Community</a> to connect with creators and spark inspiration.</p>
-<style>
-  @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');
-  body {
-    background-color: #1A1A1A !important;
-    font-family: 'Inter', sans-serif !important;
-    color: #D1D5DB !important;
-  }
-  .gr-panel {
-    background-color: #2D2D2D !important;
-    border: 1px solid #7C3AED !important;
-    border-radius: 12px !important;
-    padding: 20px !important;
-    box-shadow: 0 4px 10px rgba(124, 58, 237, 0.2) !important;
-  }
-  .gr-button-primary {
-    background: linear-gradient(45deg, #7C3AED, #A78BFA) !important;
-    color: white !important;
-    border: none !important;
-    border-radius: 8px !important;
-    padding: 12px 24px !important;
-    font-weight: 600 !important;
-    transition: transform 0.2s, box-shadow 0.2s !important;
-  }
-  .gr-button-primary:hover {
-    transform: translateY(-2px) !important;
-    box-shadow: 0 4px 12px rgba(124, 58, 237, 0.5) !important;
-  }
-  .gr-button-secondary {
-    background-color: #4B4B4B !important;
-    color: #D1D5DB !important;
-    border: 1px solid #A78BFA !important;
-    border-radius: 8px !important;
-    padding: 10px 20px !important;
-    transition: transform 0.2s !important;
-  }
-  .gr-button-secondary:hover {
-    transform: translateY(-1px) !important;
-    background-color: #6B6B6B !important;
-  }
-  .gr-accordion {
-    background-color: #2D2D2D !important;
-    border-radius: 8px !important;
-    border: 1px solid #7C3AED !important;
-  }
-  .gr-tab {
-    background-color: #2D2D2D !important;
-    color: #A78BFA !important;
-    border: 1px solid #7C3AED !important;
-    border-radius: 8px !important;
-    margin: 5px !important;
-  }
-  .gr-tab:hover, .gr-tab-selected {
-    background: linear-gradient(45deg, #7C3AED, #A78BFA) !important;
-    color: white !important;
-  }
-  .gr-slider input[type=range]::-webkit-slider-thumb {
-    background-color: #7C3AED !important;
-    border: 2px solid #A78BFA !important;
-  }
-  .gr-dropdown {
-    background-color: #2D2D2D !important;
-    color: #D1D5DB !important;
-    border: 1px solid #A78BFA !important;
-    border-radius: 8px !important;
-  }
-  h1, h3 {
-    color: #A78BFA !important;
-    text-shadow: 0 0 10px rgba(124, 58, 237, 0.5) !important;
-  }
-</style>
-"""
-
-# Gradio interface
+# Gradio interface (unchanged)
 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():
@@ -546,10 +551,12 @@ 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(
@@ -563,6 +570,7 @@ 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],
@@ -583,17 +591,16 @@ 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)
 
-# Mount Gradio to FastAPI
-app = gr.mount_gradio_app(app, demo, path="/")
-
+# Run both Gradio and FastAPI
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)