Update app.py

app.py

@@ -1,63 +1,49 @@
-import io
-import base64
+import os
 import torch
 from flask import Flask, request, jsonify, send_file
-from diffusers import DiffusionPipeline
+from pipeline import Zero123PlusPipeline  # from your local pipeline.py
 from PIL import Image
-import logging
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+from io import BytesIO
 
 app = Flask(__name__)
 
 # Load the model once at startup (on CPU)
-try:
-    logger.info("Loading Zero123Plus pipeline...")
-    pipe = DiffusionPipeline.from_pretrained(
-        "sudo-ai/zero123plus-v1.2",
-        torch_dtype=torch.float32,  # CPU needs float32
-    )
-    pipe.to("cpu")
-    logger.info("=== Application Startup at CPU mode =====")
-except Exception as e:
-    logger.error(f"Error loading model: {e}")
-    pipe = None
-
-def pil_to_base64(image):
-    buffer = io.BytesIO()
-    image.save(buffer, format="PNG")
-    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+print("Loading Zero123Plus pipeline on CPU...")
+pipe = Zero123PlusPipeline.from_pretrained(
+    "sudo-ai/zero123plus-v1.2",
+    torch_dtype=torch.float32,
+)
+pipe.to("cpu")
+pipe.enable_model_cpu_offload()
+print("Model loaded.")
 
 @app.route("/")
 def home():
-    return "Zero123Plus CPU API is running!"
+    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 pipe is None:
-        return jsonify({"error": "Model not loaded"}), 500
-
-    try:
-        data = request.get_json()
-        image_data = data.get("image")
-
-        if not image_data:
-            return jsonify({"error": "No image provided"}), 400
-
-        if image_data.startswith("data:image"):
-            image_data = image_data.split(",")[1]
-
-        image = Image.open(io.BytesIO(base64.b64decode(image_data))).convert("RGB")
+    if "image" not in request.files:
+        return jsonify({"error": "No image uploaded"}), 400
 
-        result = pipe(image)
-        output_image = result.images[0]
+    file = request.files["image"]
+    image = Image.open(file.stream).convert("RGB")
 
-        return jsonify({"image": f"data:image/png;base64,{pil_to_base64(output_image)}"})
+    print("Generating 3D view...")
+    result = pipe(image, num_inference_steps=50, guidance_scale=3.0)
 
-    except Exception as e:
-        logger.error(f"Error generating image: {e}")
-        return jsonify({"error": str(e)}), 500
+    output = result.images[0]
+    img_io = BytesIO()
+    output.save(img_io, "PNG")
+    img_io.seek(0)
+    return send_file(img_io, mimetype="image/png")
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)