Update app.py

app.py

@@ -7,20 +7,20 @@ from PIL import Image
 from io import BytesIO
 import base64
 
-# Setup logging
+# Logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Initialize Flask app
+# Flask app
 app = Flask(__name__)
 
-# Load the pipeline (CPU, no fp16 variant)
+# Load Zero123Plus pipeline (for CPU)
 logger.info("Loading Zero123Plus pipeline...")
 try:
     pipe = DiffusionPipeline.from_pretrained(
         "sudo-ai/zero123plus-v1.2",
         torch_dtype=torch.float32,
-        variant=None,
+        variant=None,  # avoid fp16 issues
     )
     pipe.to("cpu")
     logger.info("Pipeline loaded successfully.")
@@ -38,28 +38,31 @@ def generate():
         return jsonify({"error": "Model not loaded"}), 500
 
     try:
-        image_data = request.json.get("image")
+        data = request.get_json()
+        image_data = data.get("image")
+
         if not image_data:
             return jsonify({"error": "No image provided"}), 400
 
-        # Decode the base64 image
+        # Decode base64 to PIL image
         image = Image.open(BytesIO(base64.b64decode(image_data.split(",")[-1]))).convert("RGB")
 
-        # Run the pipeline
+        # Run inference
         logger.info("Generating 3D views...")
         output = pipe(image)
-        images = output.images  # List of generated views
+        generated_image = output.images[0]
 
-        # Return the first image as base64
+        # Convert output to base64
         buffered = BytesIO()
-        images[0].save(buffered, format="PNG")
-        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        return jsonify({"image": f"data:image/png;base64,{img_str}"})
+        generated_image.save(buffered, format="PNG")
+        img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
 
+        return jsonify({"image": f"data:image/png;base64,{img_base64}"})
+    
     except Exception as e:
-        logger.error(f"Inference error: {e}")
+        logger.error(f"Generation failed: {e}")
         return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":
-    logger.info("=== Application Startup at {} =====".format(torch.cuda.utilization().get('time') if torch.cuda.is_available() else "CPU mode"))
+    logger.info("=== Application Startup at CPU mode =====")
     app.run(host="0.0.0.0", port=7860)