Update app.py

app.py

@@ -1,36 +1,37 @@
-import os
+import io
+import base64
 import torch
-import logging
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, send_file
 from diffusers import DiffusionPipeline
 from PIL import Image
-from io import BytesIO
-import base64
+import logging
 
-# Logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Flask app
 app = Flask(__name__)
 
-# Load Zero123Plus pipeline (for CPU)
-logger.info("Loading Zero123Plus pipeline...")
+# 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,
-        variant=None,  # avoid fp16 issues
+        torch_dtype=torch.float32,  # CPU needs float32
     )
     pipe.to("cpu")
-    logger.info("Pipeline loaded successfully.")
+    logger.info("=== Application Startup at CPU mode =====")
 except Exception as e:
     logger.error(f"Error loading model: {e}")
     pipe = None
 
-@app.route("/", methods=["GET"])
-def index():
-    return "Zero123Plus API is running."
+def pil_to_base64(image):
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+@app.route("/")
+def home():
+    return "Zero123Plus CPU API is running!"
 
 @app.route("/generate", methods=["POST"])
 def generate():
@@ -44,25 +45,19 @@ def generate():
         if not image_data:
             return jsonify({"error": "No image provided"}), 400
 
-        # Decode base64 to PIL image
-        image = Image.open(BytesIO(base64.b64decode(image_data.split(",")[-1]))).convert("RGB")
+        if image_data.startswith("data:image"):
+            image_data = image_data.split(",")[1]
 
-        # Run inference
-        logger.info("Generating 3D views...")
-        output = pipe(image)
-        generated_image = output.images[0]
+        image = Image.open(io.BytesIO(base64.b64decode(image_data))).convert("RGB")
 
-        # Convert output to base64
-        buffered = BytesIO()
-        generated_image.save(buffered, format="PNG")
-        img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        result = pipe(image)
+        output_image = result.images[0]
+
+        return jsonify({"image": f"data:image/png;base64,{pil_to_base64(output_image)}"})
 
-        return jsonify({"image": f"data:image/png;base64,{img_base64}"})
-    
     except Exception as e:
-        logger.error(f"Generation failed: {e}")
+        logger.error(f"Error generating image: {e}")
         return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":
-    logger.info("=== Application Startup at CPU mode =====")
     app.run(host="0.0.0.0", port=7860)