Update app.py

app.py

@@ -1,21 +1,42 @@
 import os
+import logging
 from flask import Flask, request, jsonify, send_file
 from diffusers import DiffusionPipeline
 import torch
 from PIL import Image
 import io
 
+# Configure logging to stdout instead of files
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler()]
+)
+logger = logging.getLogger(__name__)
+
 # Set Hugging Face cache directory to a writable path
-os.environ['HF_HOME'] = '/home/user/.cache/huggingface'
+os.environ['HF_HOME'] = '/tmp/.cache/huggingface'
+
+# Create cache directory if it doesn't exist
+os.makedirs('/tmp/.cache/huggingface', exist_ok=True)
 
 app = Flask(__name__)
 
-# Load the model on CPU
-pipe = DiffusionPipeline.from_pretrained(
-    "sudo-ai/zero123plus-v1.2",
-    torch_dtype=torch.float32
-)
-pipe.to("cpu")
+# Initialize the model
+@app.before_first_request
+def load_model():
+    global pipe
+    try:
+        logger.info("Loading Zero123Plus model...")
+        pipe = DiffusionPipeline.from_pretrained(
+            "sudo-ai/zero123plus-v1.2",
+            torch_dtype=torch.float32
+        )
+        pipe.to("cpu")
+        logger.info("Model loaded successfully")
+    except Exception as e:
+        logger.error(f"Error loading model: {str(e)}")
+        raise
 
 @app.route("/", methods=["GET"])
 def index():
@@ -23,22 +44,34 @@ def index():
 
 @app.route("/generate", methods=["POST"])
 def generate():
-    if 'image' not in request.files:
-        return jsonify({"error": "No image uploaded"}), 400
+    try:
+        if 'image' not in request.files:
+            logger.warning("No image uploaded")
+            return jsonify({"error": "No image uploaded"}), 400
 
-    image_file = request.files['image']
-    input_image = Image.open(image_file).convert("RGB")
+        image_file = request.files['image']
+        input_image = Image.open(image_file).convert("RGB")
+        logger.info(f"Received image of size {input_image.size}")
 
-    # Generate new views
-    result = pipe(image=input_image, num_inference_steps=25)
-    output_image = result.images[0]
+        # Get optional parameters with defaults
+        num_steps = int(request.form.get('num_inference_steps', 25))
+        
+        logger.info(f"Starting image generation with {num_steps} steps")
+        # Generate new views
+        result = pipe(image=input_image, num_inference_steps=num_steps)
+        output_image = result.images[0]
+        logger.info(f"Generated image of size {output_image.size}")
 
-    # Save to a BytesIO object
-    img_io = io.BytesIO()
-    output_image.save(img_io, 'PNG')
-    img_io.seek(0)
+        # Save to a BytesIO object
+        img_io = io.BytesIO()
+        output_image.save(img_io, 'PNG')
+        img_io.seek(0)
 
-    return send_file(img_io, mimetype='image/png')
+        return send_file(img_io, mimetype='image/png')
+    
+    except Exception as e:
+        logger.error(f"Error during generation: {str(e)}")
+        return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)