Upload 2 files

Dockerfile

@@ -18,15 +18,29 @@ RUN pip install uv
 COPY requirements.txt .
 RUN uv pip install --system -r requirements.txt
 
-# Copy application code
-COPY app.py .
+# Create cache directories with write permissions
+RUN mkdir -p /app/cache/huggingface /app/cache/openvino /app/matplotlib_cache /app/openvino_cache \
+    && chmod -R 777 /app/cache /app/matplotlib_cache /app/openvino_cache
 
-# Download pre-converted OpenVINO model
+# Set environment variables for cache directories
+ENV HF_HOME=/app/cache/huggingface
+ENV MPLCONFIGDIR=/app/matplotlib_cache
+ENV OPENVINO_TELEMETRY_DIR=/app/openvino_cache
+
+# Pre-download base SDXL model
 RUN python -c "from optimum.intel.openvino import OVStableDiffusionPipeline; \
-    OVStableDiffusionPipeline.from_pretrained('rupeshs/hyper-sd-sdxl-1-step-openvino-int8', ov_config={'CACHE_DIR': ''})"
+    OVStableDiffusionPipeline.from_pretrained('rupeshs/hyper-sd-sdxl-1-step-openvino-int8', ov_config={'CACHE_DIR': '/app/cache/openvino'})"
+
+# Pre-download a default LoRA model
+RUN python -c "from diffusers import LoraLoaderMixin; \
+    LoraLoaderMixin.download_lora_weights('latent-consistency/lcm-lora-sdxl', cache_dir='/app/cache/huggingface')"
+
+# Copy application code
+COPY app.py .
 
-# Expose port 5000
-EXPOSE 7860
+# Expose port (default 5000, configurable via PORT env variable)
+ENV PORT=7860
+EXPOSE $PORT
 
 # Command to run the Flask app
 CMD ["python", "app.py"]

app.py

@@ -1,36 +1,71 @@
 import os
 from flask import Flask, request, jsonify, send_file
 from optimum.intel.openvino.modeling_diffusion import OVStableDiffusionPipeline
+from diffusers import LoraLoaderMixin
 from PIL import Image
 import io
 import torch
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 app = Flask(__name__)
 
-# Load the pre-converted OpenVINO SDXL model
-model_id = "rupeshs/hyper-sd-sdxl-1-step-openvino-int8"
-pipeline = OVStableDiffusionPipeline.from_pretrained(
-    model_id,
-    ov_config={"CACHE_DIR": ""},
-    device="CPU"
-)
+# 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)
 
-# Ensure Tiny Auto Encoder is enabled to reduce memory usage
-pipeline.enable_tiny_auto_encoder()
+# 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"
+    )
+    pipeline.enable_tiny_auto_encoder()
+    logger.info("Base model loaded successfully")
+except Exception as e:
+    logger.error(f"Failed to load base model: {str(e)}")
+    raise
 
 @app.route('/generate', methods=['POST'])
 def generate_image():
     try:
-        # Get prompt from request
+        # Get parameters from request
         data = request.get_json()
         prompt = data.get('prompt', 'A futuristic cityscape at sunset, cyberpunk style, 8k')
         width = data.get('width', 512)
         height = data.get('height', 512)
-        num_inference_steps = data.get('num_inference_steps', 1)
+        num_inference_steps = data.get('num_inference_steps', 4)
         guidance_scale = data.get('guidance_scale', 1.0)
+        lora_model_id = data.get('lora_model_id', None)
+        lora_weight = data.get('lora_weight', 0.8)
+
+        # Load LoRA weights if specified
+        local_pipeline = pipeline
+        if lora_model_id:
+            try:
+                local_pipeline = LoraLoaderMixin.load_lora_weights(
+                    local_pipeline,
+                    lora_model_id,
+                    lora_scale=lora_weight,
+                    cache_dir="/app/cache/huggingface"
+                )
+                logger.info(f"LoRA model {lora_model_id} loaded successfully")
+            except Exception as e:
+                logger.error(f"Failed to load LoRA model: {str(e)}")
+                return jsonify({'error': f"Failed to load LoRA model: {str(e)}"}), 400
 
         # Generate image
-        image = pipeline(
+        image = local_pipeline(
             prompt=prompt,
             width=width,
             height=height,
@@ -50,7 +85,9 @@ def generate_image():
             download_name='generated_image.png'
         )
     except Exception as e:
+        logger.error(f"Image generation failed: {str(e)}")
         return jsonify({'error': str(e)}), 500
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860)
+    port = int(os.getenv('PORT', 7860))
+    app.run(host='0.0.0.0', port=port)