Update app.py

app.py

@@ -1,54 +1,62 @@
-import os
-import sys
-
-# Add this before importing pipeline
-sys.path.append(os.path.dirname(__file__))
-
-from flask import Flask, request, jsonify
-from PIL import Image
+import io
+import base64
 import torch
-from zero_pipeline import Zero123PlusPipeline  # Now this will work
+from flask import Flask, request, jsonify, send_file
+from diffusers import DiffusionPipeline
+from PIL import Image
+import logging
 
-app = Flask(__name__)
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Load the pipeline once when the app starts
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Running on device: {device}")
+app = Flask(__name__)
 
-pipe = Zero123PlusPipeline.from_pretrained(
-    "sudo-ai/zero123plus-v1.2",
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-)
-pipe = pipe.to(device)
+# 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")
 
 @app.route("/")
-def index():
-    return "Zero123Plus API is running!"
+def home():
+    return "Zero123Plus CPU API is running!"
 
-@app.route("/predict", methods=["POST"])
-def predict():
-    if 'image' not in request.files:
-        return jsonify({"error": "No image uploaded"}), 400
+@app.route("/generate", methods=["POST"])
+def generate():
+    if pipe is None:
+        return jsonify({"error": "Model not loaded"}), 500
 
-    image = request.files["image"]
     try:
-        input_image = Image.open(image).convert("RGB")
+        data = request.get_json()
+        image_data = data.get("image")
+
+        if not image_data:
+            return jsonify({"error": "No image provided"}), 400
 
-        result = pipe(input_image, num_inference_steps=75, num_images_per_prompt=4)
+        if image_data.startswith("data:image"):
+            image_data = image_data.split(",")[1]
 
-        images = result.images  # List of PIL Images
-        output_dir = "outputs"
-        os.makedirs(output_dir, exist_ok=True)
-        saved_paths = []
+        image = Image.open(io.BytesIO(base64.b64decode(image_data))).convert("RGB")
 
-        for i, img in enumerate(images):
-            path = os.path.join(output_dir, f"output_{i}.png")
-            img.save(path)
-            saved_paths.append(path)
+        result = pipe(image)
+        output_image = result.images[0]
 
-        return jsonify({"outputs": saved_paths})
+        return jsonify({"image": f"data:image/png;base64,{pil_to_base64(output_image)}"})
 
     except Exception as e:
+        logger.error(f"Error generating image: {e}")
         return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":