Update app.py

app.py

@@ -38,18 +38,30 @@ if np.__version__ != "1.23.5":
 if not torch.__version__.startswith(("2.1.0", "2.3.1")):
     print(f"WARNING: PyTorch version {torch.__version__} may not be compatible. Expected torch==2.1.0 or 2.3.1.")
 
-# 1) DEVICE SETUP
+# 1) DEVICE SETUP WITH ENHANCED DEBUGGING
+print("Debugging GPU and CUDA setup...")
+print(f"PyTorch CUDA available: {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    print(f"CUDA device count: {torch.cuda.device_count()}")
+    print(f"CUDA current device: {torch.cuda.current_device()}")
+    print(f"CUDA device name: {torch.cuda.get_device_name(0)}")
+    print(f"CUDA version: {torch.version.cuda}")
+else:
+    print("WARNING: CUDA is not available. Checking for potential issues...")
+    print(f"CUDA library loaded: {torch._C._cuda_getCompiledVersion() > 0}")
+    print("Please ensure the Hugging Face Space is configured to use a GPU runtime.")
+    print("Falling back to CPU rendering. Note: Performance may be significantly slower.")
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 if device != "cuda":
-    print("ERROR: CUDA is required for GPU rendering. CPU rendering is disabled.")
-    sys.exit(1)
-print(f"CUDA is available. Using GPU: {torch.cuda.get_device_name(0)}")
+    print("WARNING: Running on CPU. This is a fallback due to CUDA unavailability. GPU rendering is strongly recommended for performance.")
 
 # Pre-run memory cleanup
-torch.cuda.empty_cache()
-gc.collect()
-torch.cuda.ipc_collect()
-torch.cuda.synchronize()
+if device == "cuda":
+    torch.cuda.empty_cache()
+    gc.collect()
+    torch.cuda.ipc_collect()
+    torch.cuda.synchronize()
 
 # 2) LOAD MUSICGEN INTO VRAM
 try:
@@ -67,13 +79,17 @@ except Exception as e:
 # 3) RESOURCE MONITORING FUNCTION
 def print_resource_usage(stage: str):
     print(f"--- {stage} ---")
-    print(f"GPU Memory Allocated: {torch.cuda.memory_allocated() / (1024**3):.2f} GB")
-    print(f"GPU Memory Reserved: {torch.cuda.memory_reserved() / (1024**3):.2f} GB")
+    if device == "cuda":
+        print(f"GPU Memory Allocated: {torch.cuda.memory_allocated() / (1024**3):.2f} GB")
+        print(f"GPU Memory Reserved: {torch.cuda.memory_reserved() / (1024**3):.2f} GB")
     print(f"CPU Memory Used: {psutil.virtual_memory().percent}%")
     print("---------------")
 
-# Check available GPU memory
+# Check available GPU memory if using CUDA
 def check_vram_availability(required_gb=3.5):
+    if device != "cuda":
+        print("Skipping VRAM check as running on CPU.")
+        return True
     total_vram = torch.cuda.get_device_properties(0).total_memory / (1024**3)
     allocated_vram = torch.cuda.memory_allocated() / (1024**3)
     available_vram = total_vram - allocated_vram
@@ -237,7 +253,7 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
 
         for i in range(num_chunks):
             chunk_prompt = instrumental_prompt
-            print(f"Generating chunk {i+1}/{num_chunks} on GPU (prompt: {chunk_prompt})...")
+            print(f"Generating chunk {i+1}/{num_chunks} on {device} (prompt: {chunk_prompt})...")
             musicgen_model.set_generation_params(
                 duration=generation_duration,
                 use_sampling=True,
@@ -274,10 +290,11 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
             os.unlink(temp_wav_path)
             audio_segments.append(segment)
 
-            torch.cuda.empty_cache()
-            gc.collect()
-            torch.cuda.ipc_collect()
-            torch.cuda.synchronize()
+            if device == "cuda":
+                torch.cuda.empty_cache()
+                gc.collect()
+                torch.cuda.ipc_collect()
+                torch.cuda.synchronize()
             time.sleep(0.5)
             print_resource_usage(f"After Chunk {i+1} Generation")
 
@@ -311,10 +328,11 @@ def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p
     except Exception as e:
         return None, f"❌ Generation failed: {e}"
     finally:
-        torch.cuda.empty_cache()
-        gc.collect()
-        torch.cuda.ipc_collect()
-        torch.cuda.synchronize()
+        if device == "cuda":
+            torch.cuda.empty_cache()
+            gc.collect()
+            torch.cuda.ipc_collect()
+            torch.cuda.synchronize()
 
 # Function to clear inputs
 def clear_inputs():