Update app.py

app.py

@@ -1,92 +1,91 @@
-import os
-from flask import Flask, request, jsonify, send_file
-from optimum.intel.openvino.modeling_diffusion import OVStableDiffusionPipeline
-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__)
-
-# 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"
-    )
-    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 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', 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 = local_pipeline(
-            prompt=prompt,
-            width=width,
-            height=height,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale
-        ).images[0]
-
-        # Save image to a bytes buffer
-        img_io = io.BytesIO()
-        image.save(img_io, 'PNG')
-        img_io.seek(0)
-
-        return send_file(
-            img_io,
-            mimetype='image/png',
-            as_attachment=True,
-            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__':
-    port = int(os.getenv('PORT', 7860))
+import os
+from flask import Flask, request, jsonify, send_file
+from optimum.intel.openvino.modeling_diffusion import OVStableDiffusionPipeline
+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__)
+
+# 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
+
+@app.route('/generate', methods=['POST'])
+def generate_image():
+    try:
+        # 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', 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 = local_pipeline(
+            prompt=prompt,
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale
+        ).images[0]
+
+        # Save image to a bytes buffer
+        img_io = io.BytesIO()
+        image.save(img_io, 'PNG')
+        img_io.seek(0)
+
+        return send_file(
+            img_io,
+            mimetype='image/png',
+            as_attachment=True,
+            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__':
+    port = int(os.getenv('PORT', 7860))
     app.run(host='0.0.0.0', port=port)