Update app.py

app.py

@@ -56,6 +56,23 @@ except FileNotFoundError:
     print(f"Directory not found: {inference_dir}")
     exit(1)
 
+def empty_output_folder(output_dir):
+    # List all files in the output directory
+    files = os.listdir(output_dir)
+    
+    # Iterate over the files and remove them
+    for file in files:
+        file_path = os.path.join(output_dir, file)
+        try:
+            if os.path.isdir(file_path):
+                # If it's a directory, remove it recursively
+                shutil.rmtree(file_path)
+            else:
+                # If it's a file, delete it
+                os.remove(file_path)
+        except Exception as e:
+            print(f"Error deleting file {file_path}: {e}")
+
 # Function to create a temporary file with string content
 def create_temp_file(content, prefix, suffix=".txt"):
     temp_file = tempfile.NamedTemporaryFile(delete=False, mode="w", prefix=prefix, suffix=suffix)
@@ -72,6 +89,26 @@ def create_temp_file(content, prefix, suffix=".txt"):
     
     return temp_file.name
 
+def get_last_mp3_file(output_dir):
+    # List all files in the output directory
+    files = os.listdir(output_dir)
+    
+    # Filter only .mp3 files
+    mp3_files = [file for file in files if file.endswith('.mp3')]
+    
+    if not mp3_files:
+        print("No .mp3 files found in the output folder.")
+        return None
+    
+    # Get the full path for the mp3 files
+    mp3_files_with_path = [os.path.join(output_dir, file) for file in mp3_files]
+    
+    # Sort the files based on the modification time (most recent first)
+    mp3_files_with_path.sort(key=lambda x: os.path.getmtime(x), reverse=True)
+    
+    # Return the most recent .mp3 file
+    return mp3_files_with_path[0]
+
 def infer(genre_txt_content, lyrics_txt_content):
     # Create temporary files
     genre_txt_path = create_temp_file(genre_txt_content, prefix="genre_")
@@ -85,7 +122,8 @@ def infer(genre_txt_content, lyrics_txt_content):
     os.makedirs(output_dir, exist_ok=True)
     print(f"Output folder ensured at: {output_dir}")
 
-    
+    empty_output_folder(output_dir)
+ 
     # Command and arguments
     command = [
         "python", "infer.py",
@@ -105,7 +143,6 @@ def infer(genre_txt_content, lyrics_txt_content):
     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', '')}"
@@ -122,9 +159,17 @@ def infer(genre_txt_content, lyrics_txt_content):
             print("Output folder contents:")
             for file in output_files:
                 print(f"- {file}")
+
+            last_mp3 = get_last_mp3_file(output_dir)
+
+            if last_mp3:
+                print("Last .mp3 file:", last_mp3)
+                return last_mp3
+            else:
+                return None
         else:
             print("Output folder is empty.")
-        return None
+            return None
     except subprocess.CalledProcessError as e:
         print(f"Error occurred: {e}")
         return None