Update app.py

app.py

@@ -11,6 +11,7 @@ from torchvision import transforms
 from huggingface_hub import hf_hub_download, snapshot_download
 import subprocess
 import shutil
+import base64
 
 # Install additional dependencies
 subprocess.run("pip install spandrel==0.4.1 --no-deps", shell=True, check=True)
@@ -204,16 +205,27 @@ def run_full(image: str, seed: int = 0, num_inference_steps: int = 50, guidance_
     return image_seg, mesh_path, textured_glb_path
 
 def gradio_generate(image: str, seed: int = 0, num_inference_steps: int = 50, guidance_scale: float = 7.5, simplify: bool = True, target_face_num: int = DEFAULT_FACE_NUMBER):
-    if not image or not os.path.exists(image):
-        raise ValueError("Invalid or missing image file")
-    
     # Verify API key
     api_key = os.getenv("POLYGENIX_API_KEY", "your-secret-api-key")
     request = gr.Request()
     if not request.headers.get("x-api-key") == api_key:
         raise ValueError("Invalid API key")
 
-    image_seg, mesh_path, textured_glb_path = run_full(image, seed, num_inference_steps, guidance_scale, simplify, target_face_num, req=None)
+    # Handle base64 image
+    if image.startswith("data:image"):
+        # Extract base64 data (e.g., "data:image/jpeg;base64,...")
+        base64_string = image.split(",")[1]
+        image_data = base64.b64decode(base64_string)
+        # Save to temporary file
+        temp_image_path = os.path.join(TMP_DIR, f"input_{get_random_hex()}.png")
+        with open(temp_image_path, "wb") as f:
+            f.write(image_data)
+    else:
+        temp_image_path = image
+        if not os.path.exists(temp_image_path):
+            raise ValueError("Invalid or missing image file")
+
+    image_seg, mesh_path, textured_glb_path = run_full(temp_image_path, seed, num_inference_steps, guidance_scale, simplify, target_face_num, req=None)
     session_hash = os.path.basename(os.path.dirname(textured_glb_path))
     return {"file_url": f"/files/{session_hash}/{os.path.basename(textured_glb_path)}"}
 
@@ -362,7 +374,7 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
 api_interface = gr.Interface(
     fn=gradio_generate,
     inputs=[
-        gr.Image(type="filepath", label="Image"),
+        gr.Image(type="filepath", label="Image"),  # Accepts base64 strings or file paths
         gr.Number(label="Seed", value=0, precision=0),
         gr.Number(label="Inference Steps", value=50, precision=0),
         gr.Number(label="Guidance Scale", value=7.5),