Update app.py

app.py

@@ -434,26 +434,46 @@ 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=""):
-    import gc
-    import torch
-    
-    if not audio or not model_keys:
-        raise ValueError("Audio or models missing.")
-    
+def auto_ensemble_process(audio, model_keys, seg_size=128, overlap=0.1, out_format="wav", use_tta="False", model_dir="/tmp/audio-separator-models/", output_dir="output", norm_thresh=0.9, amp_thresh=0.9, batch_size=1, ensemble_method="avg_wave", exclude_stems="", weights_str=""):
     temp_audio_path = None
+    chunk_paths = []
     try:
-        # Limit to 2 models for testing
-        model_keys = model_keys[:2]
+        if not audio or not model_keys:
+            raise ValueError("Audio or models missing.")
         
+        # Handle tuple input (sample_rate, data)
         if isinstance(audio, tuple):
             sample_rate, data = audio
             temp_audio_path = os.path.join("/tmp", "temp_audio.wav")
             scipy.io.wavfile.write(temp_audio_path, sample_rate, data)
             audio = temp_audio_path
         
+        # Load audio to check duration
+        audio_data, sr = librosa.load(audio, sr=None, mono=False)
+        duration = librosa.get_duration(y=audio_data, sr=sr)
+        logger.info(f"Audio duration: {duration:.2f} seconds")
+        
+        # Split audio if longer than 15 minutes (900 seconds)
+        chunk_duration = 300  # 5 minutes in seconds
+        chunks = []
+        if duration > 900:
+            logger.info(f"Audio exceeds 15 minutes, splitting into {chunk_duration}-second chunks")
+            num_chunks = int(np.ceil(duration / chunk_duration))
+            for i in range(num_chunks):
+                start = i * chunk_duration * sr
+                end = min((i + 1) * chunk_duration * sr, audio_data.shape[-1])
+                chunk_data = audio_data[:, start:end] if audio_data.ndim == 2 else audio_data[start:end]
+                chunk_path = os.path.join("/tmp", f"chunk_{i}.wav")
+                sf.write(chunk_path, chunk_data.T if audio_data.ndim == 2 else chunk_data, sr)
+                chunks.append(chunk_path)
+                chunk_paths.append(chunk_path)
+                logger.info(f"Created chunk {i}: {chunk_path}")
+        else:
+            chunks = [audio]
+        
         use_tta = use_tta == "True"
         
+        # Create output directory
         if os.path.exists(output_dir):
             shutil.rmtree(output_dir)
         os.makedirs(output_dir, exist_ok=True)
@@ -462,48 +482,70 @@ def auto_ensemble_process(audio, model_keys, seg_size, overlap, out_format, use_
         logger.info(f"Ensemble for {base_name} with {model_keys} on {device}")
         
         all_stems = []
-        total_models = len(model_keys)
+        model_stems = {}  # Store stems per model for concatenation
         
-        for i, model_key in enumerate(model_keys):
+        for model_key in model_keys:
+            model_stems[model_key] = {"vocals": [], "other": []}
             for category, models in ROFORMER_MODELS.items():
                 if model_key in models:
                     model = models[model_key]
                     break
             else:
+                logger.warning(f"Model {model_key} not found, skipping")
                 continue
             
-            separator = Separator(
-                log_level=logging.INFO,
-                model_file_dir=model_dir,
-                output_dir=output_dir,
-                output_format=out_format,
-                normalization_threshold=norm_thresh,
-                amplification_threshold=amp_thresh,
-                use_autocast=use_autocast,
-                mdxc_params={"segment_size": seg_size, "overlap": overlap, "use_tta": use_tta, "batch_size": batch_size}
-            )
-            logger.info(f"Loading {model_key}")
-            separator.load_model(model_filename=model)
-            logger.info(f"Separating with {model_key}")
-            separation = separator.separate(audio)
-            stems = [os.path.join(output_dir, file_name) for file_name in separation]
-            
-            if exclude_stems.strip():
-                excluded = [s.strip().lower() for s in exclude_stems.split(',')]
-                filtered_stems = [stem for stem in stems if not any(ex in os.path.basename(stem).lower() for ex in excluded)]
-                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()
+            for chunk_idx, chunk_path in enumerate(chunks):
+                separator = Separator(
+                    log_level=logging.INFO,
+                    model_file_dir=model_dir,
+                    output_dir=output_dir,
+                    output_format=out_format,
+                    normalization_threshold=norm_thresh,
+                    amplification_threshold=amp_thresh,
+                    use_autocast=use_autocast,
+                    mdxc_params={"segment_size": seg_size, "overlap": overlap, "use_tta": use_tta, "batch_size": batch_size}
+                )
+                logger.info(f"Loading {model_key} for chunk {chunk_idx}")
+                separator.load_model(model_filename=model)
+                logger.info(f"Separating chunk {chunk_idx} with {model_key}")
+                separation = separator.separate(chunk_path)
+                stems = [os.path.join(output_dir, file_name) for file_name in separation]
+                
+                # Store stems for this chunk
+                for stem in stems:
+                    if "vocals" in os.path.basename(stem).lower():
+                        model_stems[model_key]["vocals"].append(stem)
+                    elif "other" in os.path.basename(stem).lower():
+                        model_stems[model_key]["other"].append(stem)
+                
+                # Clean up memory
+                separator = None
+                gc.collect()
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+                    logger.info(f"Cleared CUDA cache after {model_key} chunk {chunk_idx}")
+        
+        # Combine stems for each model
+        for model_key, stems_dict in model_stems.items():
+            for stem_type in ["vocals", "other"]:
+                if stems_dict[stem_type]:
+                    combined_path = os.path.join(output_dir, f"{base_name}_{stem_type}_{model_key.replace(' | ', '_').replace(' ', '_')}.wav")
+                    combined_data = []
+                    for stem_path in stems_dict[stem_type]:
+                        data, _ = librosa.load(stem_path, sr=sr, mono=False)
+                        combined_data.append(data)
+                    combined_data = np.concatenate(combined_data, axis=-1) if combined_data[0].ndim == 2 else np.concatenate(combined_data)
+                    sf.write(combined_path, combined_data.T if combined_data.ndim == 2 else combined_data, sr)
+                    logger.info(f"Combined {stem_type} for {model_key}: {combined_path}")
+                    if exclude_stems.strip() and stem_type.lower() in [s.strip().lower() for s in exclude_stems.split(',')]:
+                        logger.info(f"Excluding {stem_type} for {model_key}")
+                        continue
+                    all_stems.append(combined_path)
         
         if not all_stems:
             raise ValueError("No valid stems for ensemble after exclusion.")
         
+        # Ensemble the combined stems
         weights = [float(w.strip()) for w in weights_str.split(',')] if weights_str.strip() else [1.0] * len(all_stems)
         if len(weights) != len(all_stems):
             weights = [1.0] * len(all_stems)
@@ -515,7 +557,7 @@ def auto_ensemble_process(audio, model_keys, seg_size, overlap, out_format, use_
             "--weights", *[str(w) for w in weights],
             "--output", output_file
         ]
-        logger.info("Running ensemble...")
+        logger.info(f"Running ensemble with args: {ensemble_args}")
         ensemble_files(ensemble_args)
         
         logger.info("Ensemble complete")
@@ -524,12 +566,14 @@ def auto_ensemble_process(audio, model_keys, seg_size, overlap, out_format, use_
         logger.error(f"Ensemble failed: {e}")
         raise RuntimeError(f"Ensemble failed: {e}")
     finally:
-        if temp_audio_path and os.path.exists(temp_audio_path):
+        # Clean up temporary files
+        for path in chunk_paths + ([temp_audio_path] if temp_audio_path and os.path.exists(temp_audio_path) else []):
             try:
-                os.remove(temp_audio_path)
-                logger.info(f"Successfully cleaned up {temp_audio_path}")
+                if os.path.exists(path):
+                    os.remove(path)
+                    logger.info(f"Successfully cleaned up {path}")
             except Exception as e:
-                logger.error(f"Failed to clean up {temp_audio_path}: {e}")
+                logger.error(f"Failed to clean up {path}: {e}")
                 
 def update_roformer_models(category):
     """Update Roformer model dropdown based on selected category."""