cleanup

app.py

@@ -46,9 +46,8 @@ models_on_device = False
 model_loading = True
 model_error: str | None = None
 
-
+# from hyperparameters.py
 def build_params(encodec_weights: str, lm_weights: str) -> hp.MusicgenParams:
-    """Assemble MusicgenParams with explicit weight paths."""
     return hp.MusicgenParams(
         encodec_params=hp.EncodecParams(
             sample_rate=32_000,
@@ -74,9 +73,7 @@ def build_params(encodec_weights: str, lm_weights: str) -> hp.MusicgenParams:
 
 
 def load_checkpoint(params: hp.MusicgenParams, ckp_path: str) -> LightningMusicgen:
-    """Instantiate model from params and load fine-tuned checkpoint onto CPU."""
-    model = LightningMusicgen(params)   # direct instantiation bcs params.instantiate() uses 
-                                        # pydoc.locate("stage.*") which fails after flattening
+    model = LightningMusicgen(params)  # model_class strings in hyperparameters.py use the old "stage.*" prefix
     sft.load_model(model, ckp_path)
     return model.cpu().eval()
 
@@ -128,7 +125,6 @@ def process_fn(
     gen_seconds: int,
     description: str,
 ) -> str:
-    """Load context audio, run STAGE autoregressive generation, return the generated stem."""
     global drums_model, bass_model, models_on_device
 
     if model_loading:
@@ -136,16 +132,13 @@ def process_fn(
     if model_error:
         raise gr.Error(f"Model failed to load: {model_error}")
 
-    model = drums_model if instrument == "Drums" else bass_model
-    if model is None:
-        raise gr.Error(f"Model for {instrument} not available.")
-
     if not models_on_device:
         drums_model = drums_model.to(DEVICE)  # type: ignore[union-attr]
         bass_model = bass_model.to(DEVICE)    # type: ignore[union-attr]
         models_on_device = True
 
-    # load and prepare context audio
+    model = drums_model if instrument == "Drums" else bass_model
+
     audio_np, orig_sr = sf.read(input_audio_path, always_2d=True)
     context = torch.from_numpy(audio_np.T).float()
     context = torchaudio.functional.resample(context, orig_sr, SAMPLE_RATE)
@@ -153,20 +146,18 @@ def process_fn(
         context = context.mean(dim=0, keepdim=True)
     context = context.reshape(1, 1, -1).to(DEVICE)
 
-    desc = [description.strip()] if description and description.strip() else [None]
-
-    out = model.generate(
+    generated = model.generate(
         n_samples=1,
         gen_seconds=gen_seconds,
         prompt=None,
         context=context,
         style=None,
         beat=None,
-        description=desc,
+        description=[description.strip() or None],
     )
 
-    # out: (1, 1, T) or (1, 2, T) — save as stereo wav
-    audio_out = out.squeeze(0).cpu().float()
+    # generated: (1, 1, T) — squeeze to (1, T) and duplicate to stereo
+    audio_out = generated.squeeze(0).cpu().float()
     if audio_out.shape[0] == 1:
         audio_out = audio_out.repeat(2, 1)