Update app.py

app.py

@@ -100,10 +100,20 @@ def infer(genre_txt_content, lyrics_txt_content):
         "--max_new_tokens", "3000",
         "--disable_offload_model"
     ]
+
+    # Set up environment variables for CUDA
+    env = os.environ.copy()  # Copy current environment
+    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', '')}"
+    })
     
     # Execute the command
     try:
-        subprocess.run(command, check=True)
+        subprocess.run(command, check=True, env=env)
         print("Command executed successfully!")
         
         # Check and print the contents of the output folder