Update app.py

app.py

@@ -124,7 +124,7 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
         "--genre_txt", f"{genre_txt_path}",
         "--lyrics_txt", f"{lyrics_txt_path}",
         "--run_n_segments", f"{num_segments}",
-        "--stage2_batch_size", "8",  # Increased from 4 to 8
+        "--stage2_batch_size", "2",
         "--output_dir", f"{output_dir}",
         "--cuda_idx", "0",
         "--max_new_tokens", f"{max_new_tokens}",
@@ -135,13 +135,9 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
     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', '')}",
-        "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
+        "LD_LIBRARY_PATH": f"/usr/local/cuda/lib64:{env.get('LD_LIBRARY_PATH', '')}"
     })
     
     # Execute the command