Update app.py

app.py

@@ -132,7 +132,7 @@ 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", "6",  # Increased from 4 to 8
+        "--stage2_batch_size", "8",  # Increased from 4 to 8
         "--output_dir", f"{output_dir}",
         "--cuda_idx", "0",
         "--max_new_tokens", "3000",
@@ -147,7 +147,9 @@ def infer(genre_txt_content, lyrics_txt_content):
         "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', '')}",
-        "TORCH_DISTRIBUTED_DEBUG": "DETAIL"  # Added for better debugging
+        "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