Update app.py

app.py

@@ -1,73 +1,88 @@
 import gradio as gr
 import os
 import shutil
+import logging
 from huggingface_hub import hf_hub_download
 from audio_separator.separator import Separator
 
 # --- Configuration ---
 REPO_ID = "anvuew/dereverb_room"
 MODEL_FILENAME = "dereverb_room_anvuew_sdr_13.7432.ckpt" 
-CONFIG_FILENAME = "dereverb_room_anvuew.yaml" 
+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
 
-    # 1. Setup local models directory
+    # 1. Setup strictly defined local models directory
     local_models_dir = os.path.abspath("models")
     os.makedirs(local_models_dir, exist_ok=True)
 
-    print(f"Downloading model files to {local_models_dir}...")
-    hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME, local_dir=local_models_dir)
-    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
-    model_full_path = os.path.join(local_models_dir, MODEL_FILENAME)
-    original_config_path = os.path.join(local_models_dir, CONFIG_FILENAME)
-    new_config_path = os.path.splitext(model_full_path)[0] + ".yaml"
+    logger.info(f"Downloading model files to {local_models_dir}...")
+    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)
 
-    if original_config_path != new_config_path:
-        print(f"Renaming config to {new_config_path} for auto-discovery...")
-        shutil.copyfile(original_config_path, new_config_path)
+    # 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
 
-    print("Initializing separator with custom model directory...")
-    # IMPORTANT: Set model_file_dir to our local directory
+    logger.info("Initializing separator...")
+    # Initialize with our custom model directory
     separator = Separator(
         model_file_dir=local_models_dir,
         output_dir=".",
-        output_format="FLAC"
+        output_format="FLAC",
+        log_level=logging.INFO
     )
 
-    print(f"Loading model: {MODEL_FILENAME}...")
-    # Now we just pass the FILENAME, not the full path. 
-    # Because we set model_file_dir, it should find it there.
-    # It might still complain about "not found in supported", if so we have one more trick.
-    try:
-        separator.load_model(model_filename=MODEL_FILENAME)
-    except ValueError as e:
-        if "not found in supported model files" in str(e):
-            print("Standard load failed, attempting raw load (might fail if unsupported by this version)...")
-            # Fallback: Some versions might allow direct path loading if specifically formatted,
-            # or we might have to stick to standard models.
-            # Let's try passing the full absolute path as a last resort if above fails,
-            # but usually setting model_file_dir and passing just filename works for custom models
-            # IF they don't strictly enforce the catalog.
-            raise e
+    # --- 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
+    }
+    # -----------------------------------------
 
-    print("Starting inference...")
+    logger.info(f"Loading validated model: {MODEL_FILENAME}...")
+    # Now we can load it normally by filename, as it exists in the registry
+    separator.load_model(model_filename=MODEL_FILENAME)
+
+    logger.info("Starting separation...")
     output_files = separator.separate(audio_path)
     
-    print(f"Separation complete. Outputs: {output_files}")
+    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="Upload Audio", type="filepath")
-        output_audio = gr.Audio(label="Dereverbed Audio", type="filepath", interactive=False)
+        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")
+    process_btn = gr.Button("Remove Reverb", variant="primary")
     process_btn.click(fn=inference, inputs=input_audio, outputs=output_audio)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()