Update app.py

app.py

@@ -555,17 +555,13 @@ def auto_ensemble_process(audio, model_keys, state, seg_size=64, overlap=0.1, ou
         permanent_output_dir = os.path.join(output_dir, "permanent_stems")
         os.makedirs(permanent_output_dir, exist_ok=True)
 
-        model_cache = {}
-        all_stems = []
-        total_tasks = len(model_keys)
-        current_idx = state["current_model_idx"]
-        logger.info(f"Current model index: {current_idx}, total models: {len(model_keys)}")
-
-        if current_idx >= len(model_keys):
+        # Check if all models have been processed
+        if state["current_model_idx"] >= len(model_keys):
             logger.info("All models processed, running ensemble...")
             progress(0.9, desc="Running ensemble...")
 
             excluded_stems_list = [s.strip().lower() for s in exclude_stems.split(',')] if exclude_stems.strip() else []
+            all_stems = []
             for model_key, stems_dict in state["model_outputs"].items():
                 for stem_type in ["vocals", "other"]:
                     if stems_dict[stem_type]:
@@ -594,6 +590,7 @@ def auto_ensemble_process(audio, model_keys, state, seg_size=64, overlap=0.1, ou
             if result is None or not os.path.exists(output_file):
                 raise RuntimeError(f"Ensemble failed, output file not created: {output_file}")
 
+            # Reset state after ensemble
             state["current_model_idx"] = 0
             state["current_audio"] = None
             state["processed_stems"] = []
@@ -610,10 +607,12 @@ def auto_ensemble_process(audio, model_keys, state, seg_size=64, overlap=0.1, ou
             status += "</ul>"
             return output_file, status, file_list, state
 
-        model_key = model_keys[current_idx]
-        logger.info(f"Processing model {current_idx + 1}/{len(model_keys)}: {model_key}")
+        # Process the next model
+        model_key = model_keys[state["current_model_idx"]]
+        logger.info(f"Processing model {state['current_model_idx'] + 1}/{len(model_keys)}: {model_key}")
         progress(0.1, desc=f"Processing model {model_key}...")
 
+        model_cache = {}
         with torch.no_grad():
             for attempt in range(max_retries + 1):
                 try:
@@ -692,12 +691,13 @@ def auto_ensemble_process(audio, model_keys, state, seg_size=64, overlap=0.1, ou
         elapsed = time.time() - start_time
         logger.info(f"Model {model_key} completed in {elapsed:.2f}s")
 
-        if state["current_model_idx"] >= len(model_keys):
-            logger.info("Last model processed, running ensemble immediately...")
-            return auto_ensemble_process(audio, model_keys, state, seg_size, overlap, out_format, use_tta, model_dir, output_dir, norm_thresh, amp_thresh, batch_size, ensemble_method, exclude_stems, weights_str, progress)
-
         file_list = state["processed_stems"]
-        status = f"Model {model_key} (Model {current_idx + 1}/{len(model_keys)}) completed in {elapsed:.2f}s<br>Click 'Run Ensemble!' to process the next model.<br>Processed stems:<ul>"
+        status = f"Model {model_key} (Model {state['current_model_idx']}/{len(model_keys)}) completed in {elapsed:.2f}s<br>"
+        if state["current_model_idx"] >= len(model_keys):
+            status += "All models processed. Click 'Run Ensemble!' to combine the stems.<br>"
+        else:
+            status += "Click 'Run Ensemble!' to process the next model.<br>"
+        status += "Processed stems:<ul>"
         for file in file_list:
             file_name = os.path.basename(file)
             status += f"<li><a href='file={file}' download>{file_name}</a></li>"
@@ -725,7 +725,7 @@ def auto_ensemble_process(audio, model_keys, state, seg_size=64, overlap=0.1, ou
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
             logger.info("GPU memory cleared")
-
+            
 def update_roformer_models(category):
     choices = list(ROFORMER_MODELS.get(category, {}).keys()) or []
     logger.debug(f"Updating roformer models for category {category}: {choices}")