check disk space

app.py

@@ -32,6 +32,7 @@ audio_model = None
 audio_config = None
 
 def initialize_models():
+    check_disk_space()
     global kosmos_model, kosmos_processor, zephyr_pipe, audio_model, audio_config
     try:
         print("Loading Kosmos-2...")
@@ -47,6 +48,8 @@ def initialize_models():
     except Exception as e:
         print(f"Error loading Kosmos-2: {e}")
         raise
+
+    check_disk_space()
     try:
         print("Loading Zephyr...")
         zephyr_pipe = pipeline(
@@ -59,6 +62,7 @@ def initialize_models():
         print(f"Error loading Zephyr: {e}")
         raise
 
+    check_disk_space()
     try:
         print("Loading Stable Audio...")
         audio_model, audio_config = get_pretrained_model("stabilityai/stable-audio-open-1.0")
@@ -67,12 +71,14 @@ def initialize_models():
     except Exception as e:
         print(f"Error loading Stable Audio: {e}")
         raise
+    check_disk_space()
 
 def get_caption(image_in):
     if not image_in:
         raise gr.Error("Please provide an image")
         
     try:
+        check_disk_space()
         # Convert image to PIL if needed
         if isinstance(image_in, str):
             image = Image.open(image_in)
@@ -116,6 +122,7 @@ def get_musical_prompt(user_prompt, chosen_model):
         raise gr.Error("No image caption provided")
     
     try:
+        check_disk_space()
         standard_sys = """
 You are a musician AI who specializes in translating architectural spaces into musical experiences. Your job is to create concise musical descriptions that capture the essence of architectural photographs.
 
@@ -173,6 +180,7 @@ def get_stable_audio_open(prompt, seconds_total=47, steps=100, cfg_scale=7):
     try:
         torch.cuda.empty_cache()  # Clear GPU memory before generation
 
+        check_disk_space()
         device = "cuda" if torch.cuda.is_available() else "cpu"
         sample_rate = audio_config["sample_rate"]
         sample_size = audio_config["sample_size"]
@@ -226,6 +234,17 @@ def get_storage_info():
     percent = disk_usage.percent
     return f"Storage: {used}/{total} ({percent}% used)"
 
+def check_disk_space(min_gb=10):
+    """Check if there's enough disk space (default: 10GB)"""
+    disk_usage = psutil.disk_usage('/')
+    gb_free = disk_usage.free / (1024 * 1024 * 1024)
+    if gb_free < min_gb:
+        print("Disk space: {gb_free:.1f}GB free")
+        raise RuntimeError(f"Low disk space: {gb_free:.1f}GB free, need {min_gb}GB")
+    else:
+        print("Disk space: {gb_free:.1f}GB free")
+    return True
+
 def smart_cleanup():
     try:
         cache_info = scan_cache_dir()