Update app.py

app.py

@@ -1,49 +1,55 @@
 import os
-import torch
-from flask import Flask, request, jsonify, send_file
-from pipeline import Zero123PlusPipeline  # from your local pipeline.py
+import sys
+from flask import Flask, request, jsonify
 from PIL import Image
-from io import BytesIO
+import torch
+
+# Add the current directory to sys.path to allow local import
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+from pipeline import Zero123PlusPipeline
 
 app = Flask(__name__)
 
-# Load the model once at startup (on CPU)
-print("Loading Zero123Plus pipeline on CPU...")
+# Load the pipeline once when the app starts
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Running on device: {device}")
+
 pipe = Zero123PlusPipeline.from_pretrained(
     "sudo-ai/zero123plus-v1.2",
-    torch_dtype=torch.float32,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
 )
-pipe.to("cpu")
-pipe.enable_model_cpu_offload()
-print("Model loaded.")
+pipe = pipe.to(device)
 
 @app.route("/")
-def home():
-    return '''
-        <h1>Zero123Plus Image to 3D Generator</h1>
-        <form action="/generate" method="post" enctype="multipart/form-data">
-            <p>Upload a single-view image:</p>
-            <input type="file" name="image"><br><br>
-            <input type="submit" value="Generate 3D View">
-        </form>
-    '''
-
-@app.route("/generate", methods=["POST"])
-def generate():
-    if "image" not in request.files:
+def index():
+    return "Zero123Plus API is running!"
+
+@app.route("/predict", methods=["POST"])
+def predict():
+    if 'image' not in request.files:
         return jsonify({"error": "No image uploaded"}), 400
 
-    file = request.files["image"]
-    image = Image.open(file.stream).convert("RGB")
+    image = request.files["image"]
+    try:
+        input_image = Image.open(image).convert("RGB")
+
+        result = pipe(input_image, num_inference_steps=75, num_images_per_prompt=4)
+
+        images = result.images  # List of PIL Images
+        output_dir = "outputs"
+        os.makedirs(output_dir, exist_ok=True)
+        saved_paths = []
+
+        for i, img in enumerate(images):
+            path = os.path.join(output_dir, f"output_{i}.png")
+            img.save(path)
+            saved_paths.append(path)
 
-    print("Generating 3D view...")
-    result = pipe(image, num_inference_steps=50, guidance_scale=3.0)
+        return jsonify({"outputs": saved_paths})
 
-    output = result.images[0]
-    img_io = BytesIO()
-    output.save(img_io, "PNG")
-    img_io.seek(0)
-    return send_file(img_io, mimetype="image/png")
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)