Update app.py

app.py

@@ -2,7 +2,7 @@ import io
 import base64
 import torch
 from flask import Flask, request, jsonify, send_file
-from diffusers import DiffusionPipeline
+from diffusers import Zero123PlusPipeline
 from PIL import Image
 import logging
 
@@ -14,14 +14,15 @@ app = Flask(__name__)
 # Load the model once at startup (on CPU)
 try:
     logger.info("Loading Zero123Plus pipeline...")
-    pipe = DiffusionPipeline.from_pretrained(
+    pipe = Zero123PlusPipeline.from_pretrained(
         "sudo-ai/zero123plus-v1.2",
         torch_dtype=torch.float32,  # CPU needs float32
+        cache_dir="/tmp/hf_home",
     )
     pipe.to("cpu")
     logger.info("=== Application Startup at CPU mode =====")
 except Exception as e:
-    logger.error(f"Error loading model: {e}")
+    logger.error(f"Error loading model: {e}", exc_info=True)
     pipe = None
 
 def pil_to_base64(image):
@@ -50,13 +51,15 @@ def generate():
 
         image = Image.open(io.BytesIO(base64.b64decode(image_data))).convert("RGB")
 
+        logger.info("Processing image with pipeline...")
         result = pipe(image)
         output_image = result.images[0]
+        logger.info("Image processed successfully")
 
         return jsonify({"image": f"data:image/png;base64,{pil_to_base64(output_image)}"})
 
     except Exception as e:
-        logger.error(f"Error generating image: {e}")
+        logger.error(f"Error generating image: {e}", exc_info=True)
         return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":