Update app.py

app.py

@@ -124,7 +124,7 @@ def infer(genre_txt_content, lyrics_txt_content):
 
     empty_output_folder(output_dir)
  
-    # Command and arguments
+    # Command and arguments with optimized settings
     command = [
         "python", "infer.py",
         "--stage1_model", "m-a-p/YuE-s1-7B-anneal-en-cot",
@@ -132,20 +132,24 @@ def infer(genre_txt_content, lyrics_txt_content):
         "--genre_txt", f"{genre_txt_path}",
         "--lyrics_txt", f"{lyrics_txt_path}",
         "--run_n_segments", "2",
-        "--stage2_batch_size", "4",
+        "--stage2_batch_size", "8",  # Increased from 4 to 8
         "--output_dir", f"{output_dir}",
         "--cuda_idx", "0",
         "--max_new_tokens", "3000",
         "--disable_offload_model"
     ]
 
-    # Set up environment variables for CUDA
-    env = os.environ.copy()  # Copy current environment
+    # Set up environment variables for CUDA with optimized settings
+    env = os.environ.copy()
     env.update({
         "CUDA_VISIBLE_DEVICES": "0",
+        "PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:512",
         "CUDA_HOME": "/usr/local/cuda",
         "PATH": f"/usr/local/cuda/bin:{env.get('PATH', '')}",
-        "LD_LIBRARY_PATH": f"/usr/local/cuda/lib64:{env.get('LD_LIBRARY_PATH', '')}"
+        "LD_LIBRARY_PATH": f"/usr/local/cuda/lib64:{env.get('LD_LIBRARY_PATH', '')}",
+        "PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:512,garbage_collection_threshold:0.8",  # Added garbage collection threshold
+        "TORCH_DISTRIBUTED_DEBUG": "DETAIL",  # Added for better debugging
+        "CUDA_LAUNCH_BLOCKING": "0"  # Ensure asynchronous CUDA operations
     })
     
     # Execute the command