Update app.py

app.py

@@ -86,7 +86,7 @@ def infer(genre_txt_content, lyrics_txt_content):
     print(f"Output folder ensured at: {output_dir}")
 
     
-    # Command and arguments
+    # Command and arguments with additional parameters
     command = [
         "python", "infer.py",
         "--stage1_model", "m-a-p/YuE-s1-7B-anneal-en-cot",
@@ -97,25 +97,39 @@ def infer(genre_txt_content, lyrics_txt_content):
         "--stage2_batch_size", "4",
         "--output_dir", f"{output_dir}",
         "--cuda_idx", "0",
-        "--max_new_tokens", "3000"
+        "--max_new_tokens", "3000",
+        "--disable_offload_model"  # Add this to prevent GPU/CPU transfer issues
     ]
     
-    # Execute the command
     try:
-        subprocess.run(command, check=True)
-        print("Command executed successfully!")
+        # Run with environment variables to control CUDA
+        env = os.environ.copy()
+        env["CUDA_VISIBLE_DEVICES"] = "0"
+        result = subprocess.run(
+            command, 
+            check=True,
+            env=env,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True
+        )
+        print("Command output:", result.stdout)
+        print("Command errors:", result.stderr)
         
-        # Check and print the contents of the output folder
+        # Check output files
         output_files = os.listdir(output_dir)
         if output_files:
             print("Output folder contents:")
             for file in output_files:
                 print(f"- {file}")
+            # TODO: Return the generated audio file path
+            return None
         else:
             print("Output folder is empty.")
-        return None
+            return None
     except subprocess.CalledProcessError as e:
         print(f"Error occurred: {e}")
+        print("Error output:", e.stderr)
         return None
     finally:
         # Clean up temporary files