Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import logging
+import pip
 from flask import Flask, request, jsonify, send_file
 import torch
 from PIL import Image
@@ -13,6 +14,14 @@ logging.basicConfig(
 )
 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'
@@ -32,15 +41,36 @@ def load_model():
         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
+            )
         
-        # 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",
-            safety_checker=None,
-            low_cpu_mem_usage=True
-        )
         pipe.to("cpu")
         logger.info("Model loaded successfully")
         return True
@@ -48,18 +78,17 @@ def load_model():
         logger.error(f"Error loading model: {str(e)}")
         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
-
+# Initialize Flask app
 app = Flask(__name__)
 
 # Load the model immediately
 load_model()
 
-app = Flask(__name__)
-
 @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."})
 
 @app.route("/generate", methods=["POST"])