Update app.py

app.py

@@ -369,29 +369,30 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, p
     if not audio:
         raise ValueError("No audio file provided.")
     
-    # If audio is a tuple (sample_rate, data), save it as a temporary file
-    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
-    
-    override_seg_size = override_seg_size == "True"
-    
-    if os.path.exists(output_dir):
-        shutil.rmtree(output_dir)
-    os.makedirs(output_dir, exist_ok=True)
-    
-    base_name = os.path.splitext(os.path.basename(audio))[0]
-    for category, models in ROFORMER_MODELS.items():
-        if model_key in models:
-            model = models[model_key]
-            break
-    else:
-        raise ValueError(f"Model '{model_key}' not found.")
-    
-    logger.info(f"Separating {base_name} with {model_key} on {device}")
+    temp_audio_path = None  # Initialize to None to avoid undefined variable in finally
     try:
+        # If audio is a tuple (sample_rate, data), save it as a temporary file
+        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
+        
+        override_seg_size = override_seg_size == "True"
+        
+        if os.path.exists(output_dir):
+            shutil.rmtree(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        
+        base_name = os.path.splitext(os.path.basename(audio))[0]
+        for category, models in ROFORMER_MODELS.items():
+            if model_key in models:
+                model = models[model_key]
+                break
+        else:
+            raise ValueError(f"Model '{model_key}' not found.")
+        
+        logger.info(f"Separating {base_name} with {model_key} on {device}")
         separator = Separator(
             log_level=logging.INFO,
             model_file_dir=model_dir,
@@ -417,9 +418,13 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, p
         logger.error(f"Separation failed: {e}")
         raise RuntimeError(f"Separation failed: {e}")
     finally:
-        # Clean up temporary file if created
-        if isinstance(audio, tuple) and os.path.exists(temp_audio_path):
-            os.remove(temp_audio_path)
+        # 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}")
+            except Exception as e:
+                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()):
@@ -427,80 +432,89 @@ def auto_ensemble_process(audio, model_keys, seg_size, overlap, out_format, use_
     if not audio or not model_keys:
         raise ValueError("Audio or models missing.")
     
-    # If audio is a tuple (sample_rate, data), save it as a temporary file
-    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
-    
-    use_tta = use_tta == "True"
-    
-    if os.path.exists(output_dir):
-        shutil.rmtree(output_dir)
-    os.makedirs(output_dir, exist_ok=True)
-    
-    base_name = os.path.splitext(os.path.basename(audio))[0]
-    logger.info(f"Ensemble for {base_name} with {model_keys} on {device}")
-    
-    all_stems = []
-    total_models = len(model_keys)
-    
-    for i, model_key in enumerate(model_keys):
-        for category, models in ROFORMER_MODELS.items():
-            if model_key in models:
-                model = models[model_key]
-                break
-        else:
-            continue
+    temp_audio_path = None  # Initialize to None to avoid undefined variable in finally
+    try:
+        # If audio is a tuple (sample_rate, data), save it as a temporary file
+        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
         
-        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}
-        )
-        progress(0.1 + (0.4 / total_models) * i, desc=f"Loading {model_key}")
-        separator.load_model(model_filename=model)
-        progress(0.5 + (0.4 / total_models) * i, desc=f"Separating with {model_key}")
-        separation = separator.separate(audio)
-        stems = [os.path.join(output_dir, file_name) for file_name in separation]
+        use_tta = use_tta == "True"
         
-        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)
-    
-    if not all_stems:
-        raise ValueError("No valid stems for ensemble after exclusion.")
-    
-    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)
-    
-    output_file = os.path.join(output_dir, f"{base_name}_ensemble_{ensemble_method}.{out_format}")
-    ensemble_args = [
-        "--files", *all_stems,
-        "--type", ensemble_method,
-        "--weights", *[str(w) for w in weights],
-        "--output", output_file
-    ]
-    progress(0.9, desc="Running ensemble...")
-    ensemble_files(ensemble_args)
-    
-    progress(1.0, desc="Ensemble complete")
-    return output_file, f"Ensemble completed with {ensemble_method}, excluded: {exclude_stems if exclude_stems else 'None'}"
-    finally:
-        # Clean up temporary file if created
-        if isinstance(audio, tuple) and os.path.exists(temp_audio_path):
-            os.remove(temp_audio_path)
+        if os.path.exists(output_dir):
+            shutil.rmtree(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        
+        base_name = os.path.splitext(os.path.basename(audio))[0]
+        logger.info(f"Ensemble for {base_name} with {model_keys} on {device}")
+        
+        all_stems = []
+        total_models = len(model_keys)
+        
+        for i, model_key in enumerate(model_keys):
+            for category, models in ROFORMER_MODELS.items():
+                if model_key in models:
+                    model = models[model_key]
+                    break
+            else:
+                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}
+            )
+            progress(0.1 + (0.4 / total_models) * i, desc=f"Loading {model_key}")
+            separator.load_model(model_filename=model)
+            progress(0.5 + (0.4 / total_models) * i, desc=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)
+        
+        if not all_stems:
+            raise ValueError("No valid stems for ensemble after exclusion.")
+        
+        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)
+        
+        output_file = os.path.join(output_dir, f"{base_name}_ensemble_{ensemble_method}.{out_format}")
+        ensemble_args = [
+            "--files", *all_stems,
+            "--type", ensemble_method,
+            "--weights", *[str(w) for w in weights],
+            "--output", output_file
+        ]
+        progress(0.9, desc="Running ensemble...")
+        ensemble_files(ensemble_args)
+        
+        progress(1.0, desc="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}")
+            except Exception as e:
+                logger.warning(f"Failed to clean up temporary file {temp_audio_path}: {e}")
+                
 def update_roformer_models(category):
     """Update Roformer model dropdown based on selected category."""
     choices = list(ROFORMER_MODELS.get(category, {}).keys()) or []