Update app.py

app.py

@@ -1,11 +1,6 @@
 import os
 import logging
 from flask import Flask, request, jsonify, send_file
-from diffusers.pipelines import DiffusionPipeline
-try:
-    from diffusers.pipelines.zero123plus.pipeline_zero123plus import Zero123PlusPipeline
-except ImportError:
-    print("Zero123PlusPipeline not found in current diffusers version")
 import torch
 from PIL import Image
 import io
@@ -19,7 +14,6 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 # Set Hugging Face cache directory to a writable path
-# Make sure to set this BEFORE importing or initializing any models
 os.environ['HF_HOME'] = '/tmp/hf_home'
 os.environ['XDG_CACHE_HOME'] = '/tmp/cache'
 
@@ -36,19 +30,28 @@ def load_model():
     global pipe
     try:
         logger.info("Loading Zero123Plus model...")
-        # Use auto_mapping to let diffusers figure out the correct pipeline class
-        pipe = DiffusionPipeline.from_pretrained(
+        # Import here to ensure the environment variables are set before import
+        from diffusers import AutoPipelineForImage2Image
+        
+        # Use AutoPipelineForImage2Image which should work with any image-to-image model
+        pipe = AutoPipelineForImage2Image.from_pretrained(
             "sudo-ai/zero123plus-v1.2",
             torch_dtype=torch.float32,
             cache_dir="/tmp/diffusers_cache",
-            local_files_only=False,
-            resume_download=True
+            safety_checker=None,
+            low_cpu_mem_usage=True
         )
         pipe.to("cpu")
         logger.info("Model loaded successfully")
+        return True
     except Exception as e:
         logger.error(f"Error loading model: {str(e)}")
-        raise
+        return False
+
+# Don't try to load the model at startup - we'll load it on the first request
+# This prevents the app from crashing if the model can't be loaded immediately
+
+app = Flask(__name__)
 
 # Load the model immediately
 load_model()
@@ -65,10 +68,8 @@ def generate():
         global pipe
         # Check if model is loaded
         if pipe is None:
-            try:
-                load_model()
-            except Exception as e:
-                logger.error(f"Failed to load model: {str(e)}")
+            success = load_model()
+            if not success:
                 return jsonify({"error": "Failed to initialize model"}), 500
         
         if 'image' not in request.files:
@@ -84,7 +85,11 @@ def generate():
         
         logger.info(f"Starting image generation with {num_steps} steps")
         # Generate new views
-        result = pipe(image=input_image, num_inference_steps=num_steps)
+        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}")