Create app.py

app.py

@@ -0,0 +1,35 @@
+import streamlit as st
+from diffusers import StableDiffusionPipeline
+import torch
+import os
+import logging
+import psutil
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def print_memory_info():
+    # Get free CPU memory
+    virtual_mem = psutil.virtual_memory()
+    free_cpu_mem = virtual_mem.available / (1024 ** 3)  # Convert bytes to GB
+
+    # Get free GPU memory
+    free_gpu_mem = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3) - torch.cuda.memory_reserved(0) / (1024 ** 3)  # Convert bytes to GB
+
+    logger.info(f"Free CPU Memory: {free_cpu_mem:.2f} GB")
+    logger.info(f"Free GPU Memory: {free_gpu_mem:.2f} GB")
+    
+
+# @st.cache_resource
+def load_image_model():
+     print_memory_info()
+     torch.cuda.empty_cache()
+     pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", variant='fp16')
+     pipeline.to("cuda")
+     pipeline.load_lora_weights("Weights/pytorch_lora_weights.safetensors", weight_name="pytorch_lora_weights.safetensors")
+     return pipeline
+
+
+image_service = load_image_model()
+