Update app.py

app.py

@@ -434,14 +434,18 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, p
                 logger.warning(f"Failed to clean up temporary file {temp_audio_path}: {e}")
 
 @spaces.GPU
-def auto_ensemble_process(audio, model_keys, seg_size, overlap, out_format, use_tta, model_dir, output_dir, norm_thresh, amp_thresh, batch_size, ensemble_method, exclude_stems="", weights_str="", progress=gr.Progress()):
-    """Perform ensemble processing on audio using multiple Roformer models."""
+def auto_ensemble_process(audio, model_keys, seg_size, overlap, out_format, use_tta, model_dir, output_dir, norm_thresh, amp_thresh, batch_size, ensemble_method, exclude_stems="", weights_str=""):
+    import gc
+    import torch
+    
     if not audio or not model_keys:
         raise ValueError("Audio or models missing.")
     
-    temp_audio_path = None  # Initialize to None to avoid undefined variable in finally
+    temp_audio_path = None
     try:
-        # If audio is a tuple (sample_rate, data), save it as a temporary file
+        # Limit to 2 models for testing
+        model_keys = model_keys[:2]
+        
         if isinstance(audio, tuple):
             sample_rate, data = audio
             temp_audio_path = os.path.join("/tmp", "temp_audio.wav")
@@ -478,9 +482,9 @@ def auto_ensemble_process(audio, model_keys, seg_size, overlap, out_format, use_
                 use_autocast=use_autocast,
                 mdxc_params={"segment_size": seg_size, "overlap": overlap, "use_tta": use_tta, "batch_size": batch_size}
             )
-            progress(0.1 + (0.4 / total_models) * i, desc=f"Loading {model_key}")
+            logger.info(f"Loading {model_key}")
             separator.load_model(model_filename=model)
-            progress(0.5 + (0.4 / total_models) * i, desc=f"Separating with {model_key}")
+            logger.info(f"Separating with {model_key}")
             separation = separator.separate(audio)
             stems = [os.path.join(output_dir, file_name) for file_name in separation]
             
@@ -490,6 +494,12 @@ def auto_ensemble_process(audio, model_keys, seg_size, overlap, out_format, use_
                 all_stems.extend(filtered_stems)
             else:
                 all_stems.extend(stems)
+            
+            # Clean up model to free memory
+            separator = None
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
         
         if not all_stems:
             raise ValueError("No valid stems for ensemble after exclusion.")
@@ -505,22 +515,21 @@ def auto_ensemble_process(audio, model_keys, seg_size, overlap, out_format, use_
             "--weights", *[str(w) for w in weights],
             "--output", output_file
         ]
-        progress(0.9, desc="Running ensemble...")
+        logger.info("Running ensemble...")
         ensemble_files(ensemble_args)
         
-        progress(1.0, desc="Ensemble complete")
+        logger.info("Ensemble complete")
         return output_file, f"Ensemble completed with {ensemble_method}, excluded: {exclude_stems if exclude_stems else 'None'}"
     except Exception as e:
         logger.error(f"Ensemble failed: {e}")
         raise RuntimeError(f"Ensemble failed: {e}")
     finally:
-        # Clean up temporary file if it was created
         if temp_audio_path and os.path.exists(temp_audio_path):
             try:
                 os.remove(temp_audio_path)
-                logger.info(f"Cleaned up temporary file: {temp_audio_path}")
+                logger.info(f"Successfully cleaned up {temp_audio_path}")
             except Exception as e:
-                logger.warning(f"Failed to clean up temporary file {temp_audio_path}: {e}")
+                logger.error(f"Failed to clean up {temp_audio_path}: {e}")
                 
 def update_roformer_models(category):
     """Update Roformer model dropdown based on selected category."""