Update app.py

app.py

@@ -5,26 +5,46 @@ from werkzeug.utils import secure_filename
 from PIL import Image
 import io
 import zipfile
+import uuid
+import traceback
 from diffusers import ShapEImg2ImgPipeline
 from diffusers.utils import export_to_obj
 
 app = Flask(__name__)
 
-# Configure upload folder
-UPLOAD_FOLDER = 'uploads'
-RESULTS_FOLDER = 'results'
+# Configure directories - use /tmp for Hugging Face Spaces which is writable
+UPLOAD_FOLDER = '/tmp/uploads'
+RESULTS_FOLDER = '/tmp/results'
+CACHE_DIR = '/tmp/huggingface'
 ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
 
+# Create necessary directories
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 os.makedirs(RESULTS_FOLDER, exist_ok=True)
+os.makedirs(CACHE_DIR, exist_ok=True)
+
+# Set Hugging Face cache environment variables
+os.environ['HF_HOME'] = CACHE_DIR
+os.environ['TRANSFORMERS_CACHE'] = os.path.join(CACHE_DIR, 'transformers')
+os.environ['HF_DATASETS_CACHE'] = os.path.join(CACHE_DIR, 'datasets')
 
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max
 
-# Initialize the model (will download on first run)
+# Lazy loading for the model - only load when needed
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = ShapEImg2ImgPipeline.from_pretrained("openai/shap-e-img2img", torch_dtype=torch.float16)
-pipe = pipe.to(device)
+pipe = None
+
+def load_model():
+    global pipe
+    if pipe is None:
+        pipe = ShapEImg2ImgPipeline.from_pretrained(
+            "openai/shap-e-img2img", 
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+            cache_dir=CACHE_DIR  # Explicitly set cache directory
+        )
+        pipe = pipe.to(device)
+    return pipe
 
 def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
@@ -64,6 +84,9 @@ def convert_image_to_3d():
         # Open image
         image = Image.open(filepath).convert("RGB")
         
+        # Load model (lazy loading)
+        pipe = load_model()
+        
         # Generate 3D model
         images = pipe(
             image,
@@ -73,7 +96,6 @@ def convert_image_to_3d():
         ).images
         
         # Create unique output directory
-        import uuid
         output_id = str(uuid.uuid4())
         output_dir = os.path.join(RESULTS_FOLDER, output_id)
         os.makedirs(output_dir, exist_ok=True)
@@ -110,7 +132,9 @@ def convert_image_to_3d():
             return send_file(glb_path, as_attachment=True, download_name="model.glb")
     
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        # Enhanced error reporting with traceback
+        error_details = traceback.format_exc()
+        return jsonify({"error": str(e), "details": error_details}), 500
 
 @app.route('/', methods=['GET'])
 def index():
@@ -162,11 +186,13 @@ def index():
                     <li><code>output_format</code>: "obj" or "glb" (default: "obj")</li>
                 </ul>
                 <p>Example curl request:</p>
-                <pre>curl -X POST -F "image=@your_image.jpg" -F "output_format=obj" http://localhost:5000/convert -o model.zip</pre>
+                <pre>curl -X POST -F "image=@your_image.jpg" -F "output_format=obj" http://localhost:7860/convert -o model.zip</pre>
             </div>
         </body>
     </html>
     """
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 5000)))
+    # Use port 7860 which is standard for Hugging Face Spaces
+    port = int(os.environ.get('PORT', 7860))
+    app.run(host='0.0.0.0', port=port)