Update app.py

app.py

@@ -1,136 +1,56 @@
 import os
-import logging
-import pip
-from flask import Flask, request, jsonify, send_file
 import torch
+from flask import Flask, request, jsonify, send_file
+from diffusers import DiffusionPipeline
 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__)
-
-# Try to update the required packages
-try:
-    logger.info("Updating huggingface_hub and diffusers...")
-    pip.main(['install', '--upgrade', 'huggingface_hub', '--quiet'])
-    pip.main(['install', '--upgrade', 'diffusers', '--quiet'])
-except Exception as e:
-    logger.warning(f"Failed to update libraries: {str(e)}")
-
-# Set Hugging Face cache directory to a writable path
-os.environ['HF_HOME'] = '/tmp/hf_home'
-os.environ['XDG_CACHE_HOME'] = '/tmp/cache'
-
-# Create cache directories if they don't exist
-os.makedirs('/tmp/hf_home', exist_ok=True)
-os.makedirs('/tmp/cache', exist_ok=True)
-os.makedirs('/tmp/diffusers_cache', exist_ok=True)
-
-# Global variable for the model
-pipe = None
+from io import BytesIO
 
-# Initialize the model at startup
-def load_model():
-    global pipe
-    try:
-        logger.info("Loading Zero123Plus model...")
-        # Import here to ensure the environment variables are set before import
-        from diffusers import AutoPipelineForImage2Image
-        from huggingface_hub import snapshot_download
-        
-        try:
-            # First try to download the model files
-            model_path = snapshot_download(
-                "sudo-ai/zero123plus-v1.2",
-                cache_dir="/tmp/diffusers_cache",
-                local_files_only=False
-            )
-            
-            # Then load from local path
-            pipe = AutoPipelineForImage2Image.from_pretrained(
-                model_path,
-                torch_dtype=torch.float32,
-                safety_checker=None,
-                low_cpu_mem_usage=True
-            )
-        except Exception as download_error:
-            logger.warning(f"Failed to download using snapshot_download: {str(download_error)}")
-            
-            # Fallback to direct loading with local_files_only=False
-            pipe = AutoPipelineForImage2Image.from_pretrained(
-                "sudo-ai/zero123plus-v1.2",
-                torch_dtype=torch.float32,
-                cache_dir="/tmp/diffusers_cache",
-                safety_checker=None,
-                low_cpu_mem_usage=True,
-                local_files_only=False
-            )
-        
-        pipe.to("cpu")
-        logger.info("Model loaded successfully")
-        return True
-    except Exception as e:
-        logger.error(f"Error loading model: {str(e)}")
-        return False
+# Optional: logs
+import logging
+logging.basicConfig(level=logging.INFO)
 
-# Initialize Flask app
+# Flask app setup
 app = Flask(__name__)
 
-# Load the model immediately
-load_model()
+# Load model once at startup
+logging.info("Loading Zero123Plus pipeline...")
 
-@app.route("/", methods=["GET"])
-def index():
-    # Check if logs parameter is present
-    if request.args.get('logs') == 'container':
-        return jsonify({"message": "Zero123Plus API is running.", "status": "logs viewed"})
-    return jsonify({"message": "Zero123Plus API is running."})
+MODEL_ID = "sudo-ai/zero123plus-v1.2"  # Or your preferred model
+try:
+    pipe = DiffusionPipeline.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float16,
+        variant="fp16"
+    ).to("cuda" if torch.cuda.is_available() else "cpu")
+except Exception as e:
+    logging.error(f"Error loading model: {e}")
+    pipe = None
+
+@app.route("/")
+def health_check():
+    return jsonify({"status": "Zero123 API is running!"})
 
 @app.route("/generate", methods=["POST"])
-def generate():
-    try:
-        global pipe
-        # Check if model is loaded
-        if pipe is None:
-            success = load_model()
-            if not success:
-                return jsonify({"error": "Failed to initialize model"}), 500
-        
-        if 'image' not in request.files:
-            logger.warning("No image uploaded")
-            return jsonify({"error": "No image uploaded"}), 400
+def generate_image():
+    if pipe is None:
+        return jsonify({"error": "Model not loaded properly"}), 500
 
-        image_file = request.files['image']
-        input_image = Image.open(image_file).convert("RGB")
-        logger.info(f"Received image of size {input_image.size}")
+    data = request.files.get("image")
+    if not data:
+        return jsonify({"error": "No image provided"}), 400
 
-        # 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
-        )
-        
+    try:
+        input_image = Image.open(data).convert("RGB")
+        result = pipe(image=input_image, num_inference_steps=30)
         output_image = result.images[0]
-        logger.info(f"Generated image of size {output_image.size}")
 
-        # Save to a BytesIO object
-        img_io = io.BytesIO()
+        # Return as image file
+        img_io = BytesIO()
         output_image.save(img_io, 'PNG')
         img_io.seek(0)
-
         return send_file(img_io, mimetype='image/png')
-    
     except Exception as e:
-        logger.error(f"Error during generation: {str(e)}")
+        logging.error(f"Generation error: {e}")
         return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":