Update app.py

app.py

@@ -5,80 +5,102 @@ import cv2
 import os
 import tempfile
 import time
-from gradio_client import Client, file
+from gradio_client import Client, handle_file
 
-# --- CONFIGURATION ---
-# We switched to "hysts/Shap-E" (OpenAI's model).
-# It is extremely reliable and rarely crashes compared to TripoSR/InstantMesh.
-REMOTE_MODEL_ID = "hysts/Shap-E"
+# --- CONFIGURATION: PRIORITY LIST ---
+# We will try these models in order until one works.
+# 1. Official TripoSR (Best speed)
+# 2. Community Mirror (Backup)
+# 3. OpenAI Shap-E (Old reliable)
+MODELS = [
+    {"id": "stabilityai/TripoSR", "api": "/generate", "type": "tripo"},
+    {"id": "virattt/TripoSR", "api": "/generate", "type": "tripo"},
+    {"id": "hysts/Shap-E", "api": "/image-to-3d", "type": "shape"}
+]
 
 def photo_to_sketch(image):
-    """Instant local sketch – always returns PIL Image"""
-    print("-> Starting Sketch Generation...")
-    if image is None:
-        return None
-
+    """Instant local sketch"""
+    if image is None: return None
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image.astype('uint8'))
-
     gray = image.convert("L")
     img_array = np.array(gray)
-
     blurred = cv2.GaussianBlur(img_array, (5, 5), 0)
     edges = cv2.Canny(blurred, 60, 150)
-
     sketch_np = 255 - edges
-    sketch_pil = Image.fromarray(sketch_np)
-    
-    return sketch_pil.convert("RGB")
+    return Image.fromarray(sketch_np).convert("RGB")
 
 def generate_3d_avatar(sketch_image, height, weight, muscle, gender, breast):
-    """Generate 3D model using OpenAI Shap-E (Reliable Fallback)"""
-    print(f"-> Starting 3D Generation using {REMOTE_MODEL_ID}...")
+    """Try multiple remote models until one succeeds"""
+    print(f"-> Starting 3D Generation Process...")
     
     if sketch_image is None:
         raise gr.Error("Please upload an image first!")
 
-    # 1. Save Sketch to a temporary file
+    # Save temp file
     if isinstance(sketch_image, np.ndarray):
         sketch_image = Image.fromarray(sketch_image.astype('uint8'))
-        
     temp_dir = tempfile.gettempdir()
     sketch_path = os.path.join(temp_dir, f"sketch_{int(time.time())}.png")
     sketch_image.save(sketch_path)
-    
-    # 2. Connect to Shap-E
-    try:
-        print(f"-> Connecting to {REMOTE_MODEL_ID}...")
-        client = Client(REMOTE_MODEL_ID)
-        
-        # 3. Send request (Shap-E API parameters)
-        print("-> Sending request to Shap-E...")
-        
-        # Shap-E expects: [image, prompt_text, seed, guidance_scale, num_inference_steps]
-        # We leave prompt empty to force Image-to-3D mode
-        result = client.predict(
-            file(sketch_path), # Input Image
-            "",                # Text Prompt (Empty for Img-to-3D)
-            0,                 # Seed
-            15,                # Guidance Scale
-            64,                # Steps (64 is fast, 128 is better)
-            api_name="/image-to-3d"
-        )
-        
-        # Result is just the file path string
-        print(f"-> Success! Received model: {result}")
-        return result, result
+    print(f"-> Saved input to {sketch_path}")
+
+    last_error = ""
 
-    except Exception as e:
-        print(f"-> Connection Error: {e}")
-        # If even Shap-E fails, we raise a clear error
-        raise gr.Error(f"All 3D models are currently busy. Error: {e}")
+    # --- RETRY LOOP ---
+    for model in MODELS:
+        try:
+            model_id = model["id"]
+            print(f"------------------------------------------")
+            print(f"-> Attempting Connection to: {model_id}...")
+            
+            client = Client(model_id)
+            
+            if model["type"] == "tripo":
+                # TripoSR Parameters
+                print("-> Sending request (TripoSR format)...")
+                result = client.predict(
+                    handle_file(sketch_path), # New handle_file method
+                    False,                    # Do not remove background
+                    0.85,                     # Foreground ratio
+                    api_name=model["api"]
+                )
+                
+            elif model["type"] == "shape":
+                # Shap-E Parameters (Corrected inputs)
+                print("-> Sending request (Shap-E format)...")
+                result = client.predict(
+                    handle_file(sketch_path), # Input Image
+                    "high quality 3d model",  # Prompt (Fixes empty string error)
+                    0,                        # Seed
+                    15,                       # Guidance Scale
+                    64,                       # Steps
+                    api_name=model["api"]
+                )
+
+            # If we get here, it worked!
+            print(f"-> SUCCESS! Model generated by {model_id}")
+            
+            # Handle different return types (tuple vs string)
+            if isinstance(result, (list, tuple)):
+                final_glb = result[0]
+            else:
+                final_glb = result
+                
+            return final_glb, final_glb
+
+        except Exception as e:
+            print(f"-> FAILED: {model_id} | Error: {e}")
+            last_error = str(e)
+            continue # Try next model
+
+    # If loop finishes without success
+    raise gr.Error(f"All 3 backup models failed. Last error: {last_error}")
 
 # =============== UI ===============
 with gr.Blocks(title="SketchToLife") as demo:
-    gr.Markdown("# SketchToLife – Photo → Sketch → 3D Avatar")
-    gr.Markdown(f"**Status:** Connected to `OpenAI Shap-E` (High Stability Backup).")
+    gr.Markdown("# SketchToLife – Robust 3D Generator")
+    gr.Markdown("**Status:** Using Multi-Model Fallback (TripoSR → Mirror → Shap-E)")
 
     with gr.Row():
         with gr.Column():
@@ -87,8 +109,7 @@ with gr.Blocks(title="SketchToLife") as demo:
             out_sketch = gr.Image(label="Your Sketch", height=420, type="pil")
 
         with gr.Column():
-            gr.Markdown("### Customize 3D Body")
-            # Placeholders to prevent argument errors
+            gr.Markdown("### Customize Body")
             h = gr.Dropdown(["short", "average", "tall", "giant"], value="average", label="Height")
             w = gr.Dropdown(["slim", "average", "curvy", "heavy"], value="average", label="Weight")
             m = gr.Dropdown(["slim", "fit", "muscular", "bodybuilder"], value="fit", label="Muscle")
@@ -98,12 +119,10 @@ with gr.Blocks(title="SketchToLife") as demo:
             btn2 = gr.Button("Generate 3D Model", variant="primary", size="lg")
 
     with gr.Row():
-        view3d = gr.Model3D(label="Your 3D Model – Rotate & Zoom!", height=520, interactive=True)
+        view3d = gr.Model3D(label="3D Result", height=520, interactive=True)
         download = gr.File(label="Download .GLB")
 
     btn1.click(photo_to_sketch, inputs=inp, outputs=out_sketch)
-    
-    # 6 inputs -> 6 arguments
     btn2.click(generate_3d_avatar, inputs=[out_sketch, h, w, m, g, b], outputs=[view3d, download])
 
 if __name__ == "__main__":