Update app.py

app.py

@@ -1,38 +1,22 @@
 import os
 from flask import Flask, request, jsonify, send_file
-from optimum.intel.openvino.modeling_diffusion import OVStableDiffusionPipeline
+from optimum.intel.openvino.modeling_diffusion import OVStableDiffusionXLPipeline
 from PIL import Image
 import io
 import torch
 import logging
+from datetime import datetime
+
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 app = Flask(__name__)
-
-# Set cache directories
-os.environ["HF_HOME"] = "/app/cache/huggingface"
-os.environ["MPLCONFIGDIR"] = "/app/matplotlib_cache"
-os.environ["OPENVINO_TELEMETRY_DIR"] = "/app/openvino_cache"
-
-# Ensure cache directories exist
-for cache_dir in ["/app/cache/huggingface", "/app/matplotlib_cache", "/app/openvino_cache"]:
-    os.makedirs(cache_dir, exist_ok=True)
-
-# Load the base pre-converted OpenVINO SDXL model
-base_model_id = "rupeshs/hyper-sd-sdxl-1-step-openvino-int8"
-try:
-    pipeline = OVStableDiffusionPipeline.from_pretrained(
-        base_model_id,
-        ov_config={"CACHE_DIR": "/app/cache/openvino"},
-        device="CPU"
-    )
-    logger.info("Base model loaded successfully")
-except Exception as e:
-    logger.error(f"Failed to load base model: {str(e)}")
-    raise
+pipeline = OVStableDiffusionXLPipeline.from_pretrained(
+    "rupeshs/hyper-sd-sdxl-1-step-openvino-int8",
+    ov_config={"CACHE_DIR": ""},
+)
 
 @app.route('/generate', methods=['POST'])
 def generate_image():
@@ -48,10 +32,10 @@ def generate_image():
         # Generate image
         image = pipeline(
             prompt=prompt,
-            width=768,
-            height=768,
-            num_inference_steps=1,
-            guidance_scale=1.0,
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale
         ).images[0]
 
         # Save image to a bytes buffer
@@ -69,6 +53,35 @@ def generate_image():
         logger.error(f"Image generation failed: {str(e)}")
         return jsonify({'error': str(e)}), 500
 
+
+@app.route('/health', methods=['GET'])
+def health_check():
+    """
+    Health check endpoint that returns the status of the application
+    """
+    try:
+        # Basic health check response
+        health_status = {
+            'status': 'healthy',
+            'message': 'Application is running',
+            'timestamp': datetime.now().isoformat()
+        }
+
+        # Optional: Add additional checks (e.g., database connectivity)
+        # Example: Check database connection
+        # db_status = check_database_connection()
+        # health_status['database'] = db_status
+
+        return jsonify(health_status), 200
+
+    except Exception as e:
+        error_response = {
+            'status': 'unhealthy',
+            'message': f'Health check failed: {str(e)}',
+            'timestamp': datetime.now().isoformat()
+        }
+        return jsonify(error_response), 500
+
 if __name__ == '__main__':
     port = int(os.getenv('PORT', 7860))
-    app.run(host='0.0.0.0', port=port)
+    app.run(host='localhost', port=port)