Update app.py

app.py

@@ -4,42 +4,45 @@ import shutil
 import logging
 from huggingface_hub import hf_hub_download
 from audio_separator.separator import Separator
+# Import the module where models are actually defined
+from audio_separator.separator.architectures import bs_roformer_separator
 
 # --- Configuration ---
 REPO_ID = "anvuew/dereverb_room"
-MODEL_FILENAME = "dereverb_room_anvuew_sdr_13.7432.ckpt" 
+MODEL_FILENAME = "dereverb_room_anvuew_sdr_13.7432.ckpt"
 CONFIG_FILENAME = "dereverb_room_anvuew.yaml"
 # ---------------------
 
-# Setup basic logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 def inference(audio_path):
-    if not audio_path:
-        return None
+    if not audio_path: return None
 
-    # 1. Setup strictly defined local models directory
     local_models_dir = os.path.abspath("models")
     os.makedirs(local_models_dir, exist_ok=True)
 
-    logger.info(f"Downloading model files to {local_models_dir}...")
+    logger.info(f"Downloading model files...")
     model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME, local_dir=local_models_dir)
     config_path = hf_hub_download(repo_id=REPO_ID, filename=CONFIG_FILENAME, local_dir=local_models_dir)
 
-    # 2. Ensure config has the exact same base name as the model for auto-discovery
-    # (Required by some underlying loaders even if we register it)
     expected_config_path = os.path.splitext(model_path)[0] + ".yaml"
     if config_path != expected_config_path:
-        logger.info(f"Renaming config to {expected_config_path}...")
         shutil.copyfile(config_path, expected_config_path)
-        # Update config filename to match what we just created
-        actual_config_filename = os.path.basename(expected_config_path)
-    else:
-        actual_config_filename = CONFIG_FILENAME
+
+    logger.info("Registering custom model...")
+    # --- FIX: Register directly with the architecture module ---
+    bs_roformer_separator.BS_ROFORMER_MODELS[MODEL_FILENAME] = {
+        "model_type": "bs_roformer",
+        "config_filename": os.path.basename(expected_config_path),
+        "model_filename": MODEL_FILENAME,
+        "model_friendly_name": "Custom Dereverb",
+        "domain": "dereverb",
+        "source": "local"
+    }
+    # -----------------------------------------------------------
 
     logger.info("Initializing separator...")
-    # Initialize with our custom model directory
     separator = Separator(
         model_file_dir=local_models_dir,
         output_dir=".",
@@ -47,42 +50,19 @@ def inference(audio_path):
         log_level=logging.INFO
     )
 
-    # --- CRITICAL FIX: MANUAL REGISTRATION ---
-    # We manually inject the model metadata into the separator's internal registry.
-    # This tricks the library into thinking it's an officially supported model.
-    logger.info("Registering custom model at runtime...")
-    separator.model_constants.BS_ROFORMER_MODELS[MODEL_FILENAME] = {
-        "model_type": "bs_roformer",
-        "config_filename": actual_config_filename,
-        "model_filename": MODEL_FILENAME,
-        "model_friendly_name": "Custom Dereverb Room",
-        "domain": "dereverb",
-        "source": "local" # strictly local, don't try to download
-    }
-    # -----------------------------------------
-
-    logger.info(f"Loading validated model: {MODEL_FILENAME}...")
-    # Now we can load it normally by filename, as it exists in the registry
+    logger.info(f"Loading model: {MODEL_FILENAME}...")
     separator.load_model(model_filename=MODEL_FILENAME)
 
     logger.info("Starting separation...")
     output_files = separator.separate(audio_path)
-    
-    logger.info(f"Separation complete: {output_files}")
     return output_files[0]
 
-# --- Gradio UI ---
 with gr.Blocks(title="Dereverb Room Web UI") as demo:
     gr.Markdown("# Dereverb Room Inference")
-    gr.Markdown(f"**Model:** `{REPO_ID}` (BS-RoFormer)")
-    
     with gr.Row():
-        input_audio = gr.Audio(label="Input Audio", type="filepath")
-        # Using interactive=False prevents user from editing output
-        output_audio = gr.Audio(label="Cleaned Audio (No Reverb)", type="filepath", interactive=False)
-    
-    process_btn = gr.Button("Remove Reverb", variant="primary")
-    process_btn.click(fn=inference, inputs=input_audio, outputs=output_audio)
+        input_audio = gr.Audio(label="Input", type="filepath")
+        output_audio = gr.Audio(label="Output (Dereverbed)", type="filepath", interactive=False)
+    gr.Button("Remove Reverb").click(fn=inference, inputs=input_audio, outputs=output_audio)
 
 if __name__ == "__main__":
     demo.launch()