v8: Parameter locking for auto-tune — lock any param to constrain the search

sample_extractor.py

@@ -1,11 +1,9 @@
 #!/usr/bin/env python3
 """
-Sample Extractor v7 — Auto-tuning via reconstruction quality.
+Sample Extractor v8 — Auto-tuning with parameter locking.
 
-New: auto_tune() optimizes parameters against the uploaded audio itself.
-No ground truth needed — measures spectral envelope correlation between
-the rendered reconstruction and the original stem. Sweeps onset detection
-params and cluster counts, using cached NCC matrices for speed.
+auto_tune() accepts a `locks` dict: any param name mapped to its fixed value
+will be held constant during the search. All other params are swept freely.
 """
 
 import argparse, json, os, sys, warnings, hashlib
@@ -20,278 +18,205 @@ from scipy.signal import fftconvolve
 warnings.filterwarnings("ignore", category=FutureWarning)
 warnings.filterwarnings("ignore", category=UserWarning)
 
-
-# ─── Data structures ─────────────────────────────────────────────────────────
-
 @dataclass
 class Hit:
     audio: np.ndarray; sr: int; onset_time: float; duration: float; index: int
     rms_energy: float = 0.0; spectral_centroid: float = 0.0
     label: str = ""; embedding: Optional[np.ndarray] = None; cluster_id: int = -1
-    def save(self, path: str): sf.write(path, self.audio, self.sr, subtype='PCM_24')
+    def save(self, path): sf.write(path, self.audio, self.sr, subtype='PCM_24')
 
 @dataclass
 class Cluster:
     cluster_id: int; label: str; hits: list = field(default_factory=list)
     best_hit_idx: int = 0; synthesized: Optional[np.ndarray] = None; midi_note: int = 60
     @property
-    def best_hit(self) -> Hit: return self.hits[self.best_hit_idx]
+    def best_hit(self): return self.hits[self.best_hit_idx]
     @property
-    def count(self) -> int: return len(self.hits)
-
+    def count(self): return len(self.hits)
 
-DEMUCS_MODELS = ["htdemucs", "htdemucs_ft", "htdemucs_6s", "mdx", "mdx_extra", "mdx_extra_q"]
+DEMUCS_MODELS = ["htdemucs","htdemucs_ft","htdemucs_6s","mdx","mdx_extra","mdx_extra_q"]
 DEMUCS_STEMS = {
-    "htdemucs": ["drums","bass","other","vocals"], "htdemucs_ft": ["drums","bass","other","vocals"],
-    "htdemucs_6s": ["drums","bass","other","vocals","guitar","piano"],
-    "mdx": ["drums","bass","other","vocals"], "mdx_extra": ["drums","bass","other","vocals"],
-    "mdx_extra_q": ["drums","bass","other","vocals"],
+    "htdemucs":["drums","bass","other","vocals"],"htdemucs_ft":["drums","bass","other","vocals"],
+    "htdemucs_6s":["drums","bass","other","vocals","guitar","piano"],
+    "mdx":["drums","bass","other","vocals"],"mdx_extra":["drums","bass","other","vocals"],
+    "mdx_extra_q":["drums","bass","other","vocals"],
 }
 
-
-# ─── Caching ──────────────────────────────────────────────────────────────────
-
+# ─── Cache ────────────────────────────────────────────────────────────────────
 _cache = {}
-
-def _audio_hash(audio: np.ndarray) -> str:
-    return hashlib.md5(audio[:4000].tobytes()).hexdigest()
-
-def cache_get(key): return _cache.get(key)
-def cache_set(key, value): _cache[key] = value; return value
+def _audio_hash(a): return hashlib.md5(a[:4000].tobytes()).hexdigest()
+def cache_get(k): return _cache.get(k)
+def cache_set(k,v): _cache[k]=v; return v
 def cache_clear(): _cache.clear()
 
-
-# ─── Stage 1: Stem separation ────────────────────────────────────────────────
-
-def extract_stem(audio_path, stem="drums", device="cpu",
-                 model_name="htdemucs_ft", shifts=1, overlap=0.25):
-    if stem == "all":
-        y, sr = librosa.load(audio_path, sr=44100, mono=True)
-        return y.astype(np.float32), sr
-    with open(audio_path, 'rb') as f:
-        fh = hashlib.md5(f.read(200000)).hexdigest()
-    ck = ("stem", fh, stem, model_name, shifts, overlap)
-    c = cache_get(ck)
-    if c is not None: print(f"[Stage 1] Cached {stem} stem"); return c
-    from demucs.pretrained import get_model
-    from demucs.apply import apply_model
-    print(f"[Stage 1] Extracting '{stem}' with {model_name}...")
-    model = get_model(model_name); model.eval().to(device); sr = model.samplerate
+# ─── Stage 1 ──────────────────────────────────────────────────────────────────
+def extract_stem(audio_path,stem="drums",device="cpu",model_name="htdemucs_ft",shifts=1,overlap=0.25):
+    if stem=="all":
+        y,sr=librosa.load(audio_path,sr=44100,mono=True); return y.astype(np.float32),sr
+    with open(audio_path,'rb') as f: fh=hashlib.md5(f.read(200000)).hexdigest()
+    ck=("stem",fh,stem,model_name,shifts,overlap); c=cache_get(ck)
+    if c: print(f"[Stage 1] Cached {stem}"); return c
+    from demucs.pretrained import get_model; from demucs.apply import apply_model
+    print(f"[Stage 1] {stem} with {model_name}...")
+    model=get_model(model_name); model.eval().to(device); sr=model.samplerate
     if stem not in model.sources: raise ValueError(f"'{stem}' not in {model.sources}")
-    a, _ = librosa.load(audio_path, sr=sr, mono=False)
+    a,_=librosa.load(audio_path,sr=sr,mono=False)
     if a.ndim==1: a=np.stack([a,a])
     elif a.shape[0]>2: a=a[:2]
     elif a.shape[0]==1: a=np.concatenate([a,a],axis=0)
-    wav = torch.from_numpy(a).float().unsqueeze(0).to(device)
-    with torch.no_grad():
-        src = apply_model(model, wav, device=device, shifts=shifts, split=True, overlap=overlap)
-    r = src[0, model.sources.index(stem)].mean(dim=0).cpu().numpy()
-    print(f"  ✓ {stem}: {len(r)/sr:.1f}s")
-    return cache_set(ck, (r.astype(np.float32), sr))
-
-
-# ─── Stage 2: Onset detection ────────────────────────────────────────────────
-
-def detect_onsets(y, sr, pre_pad=0.005, min_dur=0.02, max_dur=1.5,
-                  min_gap=0.03, energy_threshold_db=-35.0,
-                  mode="auto", backtrack=True, onset_delta=0.12):
-    print(f"[Stage 2] Onsets (mode={mode}, delta={onset_delta}, energy≥{energy_threshold_db}dB)...")
-    if mode == "percussive":
-        oe = librosa.onset.onset_strength(y=y, sr=sr, aggregate=np.median, fmax=8000)
-    elif mode == "harmonic":
-        yh, _ = librosa.effects.hpss(y)
-        oe = librosa.onset.onset_strength(y=yh, sr=sr, fmax=8000, lag=2, max_size=3)
-    elif mode == "broadband":
-        oe = librosa.onset.onset_strength(y=y, sr=sr)
+    wav=torch.from_numpy(a).float().unsqueeze(0).to(device)
+    with torch.no_grad(): src=apply_model(model,wav,device=device,shifts=shifts,split=True,overlap=overlap)
+    r=src[0,model.sources.index(stem)].mean(dim=0).cpu().numpy()
+    print(f"  ✓ {len(r)/sr:.1f}s"); return cache_set(ck,(r.astype(np.float32),sr))
+
+# ─── Stage 2 ──────────────────────────────────────────────────────────────────
+def detect_onsets(y,sr,pre_pad=0.005,min_dur=0.02,max_dur=1.5,min_gap=0.03,
+                  energy_threshold_db=-35.0,mode="auto",backtrack=True,onset_delta=0.12):
+    print(f"[Stage 2] Onsets (delta={onset_delta}, energy≥{energy_threshold_db}dB)...")
+    if mode=="percussive": oe=librosa.onset.onset_strength(y=y,sr=sr,aggregate=np.median,fmax=8000)
+    elif mode=="harmonic": yh,_=librosa.effects.hpss(y); oe=librosa.onset.onset_strength(y=yh,sr=sr,fmax=8000,lag=2,max_size=3)
+    elif mode=="broadband": oe=librosa.onset.onset_strength(y=y,sr=sr)
     else:
-        yh, yp = librosa.effects.hpss(y)
-        envs = [librosa.onset.onset_strength(y=y,sr=sr,fmin=20,fmax=250,aggregate=np.median),
-                librosa.onset.onset_strength(y=y,sr=sr,fmin=250,fmax=4000,aggregate=np.median),
-                librosa.onset.onset_strength(y=y,sr=sr,fmin=4000,fmax=min(sr//2,20000),aggregate=np.median),
-                librosa.onset.onset_strength(y=yh,sr=sr,lag=2)]
+        yh,_=librosa.effects.hpss(y)
         def _n(x): m=x.max(); return x/m if m>0 else x
-        oe = np.maximum.reduce([_n(e) for e in envs])
-    w = max(1, int(min_gap*sr/512))
-    fr = librosa.onset.onset_detect(onset_envelope=oe, sr=sr, wait=w,
-        pre_avg=3, post_avg=3, pre_max=3, post_max=5,
-        delta=onset_delta, backtrack=backtrack, units='frames')
-    times = librosa.frames_to_time(fr, sr=sr)
-    print(f"  Raw: {len(times)}")
-    thr = 10**(energy_threshold_db/20); hits = []
-    for i, t in enumerate(times):
-        s = max(0, int((t-pre_pad)*sr))
-        e = min(int(times[i+1]*sr) if i+1<len(times) else len(y), s+int(max_dur*sr))
-        seg = y[s:e]
+        oe=np.maximum.reduce([_n(librosa.onset.onset_strength(y=y,sr=sr,fmin=20,fmax=250,aggregate=np.median)),
+            _n(librosa.onset.onset_strength(y=y,sr=sr,fmin=250,fmax=4000,aggregate=np.median)),
+            _n(librosa.onset.onset_strength(y=y,sr=sr,fmin=4000,fmax=min(sr//2,20000),aggregate=np.median)),
+            _n(librosa.onset.onset_strength(y=yh,sr=sr,lag=2))])
+    w=max(1,int(min_gap*sr/512))
+    fr=librosa.onset.onset_detect(onset_envelope=oe,sr=sr,wait=w,pre_avg=3,post_avg=3,
+        pre_max=3,post_max=5,delta=onset_delta,backtrack=backtrack,units='frames')
+    times=librosa.frames_to_time(fr,sr=sr); print(f"  Raw: {len(times)}")
+    thr=10**(energy_threshold_db/20); hits=[]
+    for i,t in enumerate(times):
+        s=max(0,int((t-pre_pad)*sr))
+        e=min(int(times[i+1]*sr) if i+1<len(times) else len(y),s+int(max_dur*sr))
+        seg=y[s:e]
         if len(seg)<int(min_dur*sr): continue
-        rms = np.sqrt(np.mean(seg**2))
+        rms=np.sqrt(np.mean(seg**2))
         if rms<thr: continue
-        fl = min(int(0.005*sr), len(seg)//4)
+        fl=min(int(0.005*sr),len(seg)//4)
         if fl>0: seg=seg.copy(); seg[-fl:]*=np.linspace(1,0,fl)
-        sc = float(librosa.feature.spectral_centroid(y=seg,sr=sr).mean())
+        sc=float(librosa.feature.spectral_centroid(y=seg,sr=sr).mean())
         hits.append(Hit(audio=seg,sr=sr,onset_time=t,duration=len(seg)/sr,
                         index=len(hits),rms_energy=float(rms),spectral_centroid=sc))
     print(f"  ✓ {len(hits)} hits"); return hits
 
-
-# ─── Stage 3: Classification ─────────────────────────────────────────────────
-
-LABEL_RULES = [
-    ("kick",         lambda lr,mr,hr,c,zcr,d: lr>0.5 and c<800),
-    ("hihat_closed", lambda lr,mr,hr,c,zcr,d: hr>0.35 and c>4000 and d<0.15),
-    ("hihat_open",   lambda lr,mr,hr,c,zcr,d: hr>0.35 and c>4000 and d>=0.15),
-    ("cymbal",       lambda lr,mr,hr,c,zcr,d: hr>0.25 and c>3000),
-    ("snare",        lambda lr,mr,hr,c,zcr,d: mr>0.4 and zcr>0.1 and c>1000),
-    ("tom",          lambda lr,mr,hr,c,zcr,d: lr>0.3 and mr>0.3 and c<1500),
-    ("bass",         lambda lr,mr,hr,c,zcr,d: lr>0.6 and c<400 and d>0.2),
-    ("vocal",        lambda lr,mr,hr,c,zcr,d: mr>0.5 and 500<c<3000 and zcr<0.15),
-    ("bright",       lambda lr,mr,hr,c,zcr,d: c>2500),
-    ("mid",          lambda lr,mr,hr,c,zcr,d: c>800),
-]
-
+# ─── Stage 3 ──────────────────────────────────────────────────────────────────
+LABEL_RULES=[("kick",lambda lr,mr,hr,c,zcr,d:lr>0.5 and c<800),
+    ("hihat_closed",lambda lr,mr,hr,c,zcr,d:hr>0.35 and c>4000 and d<0.15),
+    ("hihat_open",lambda lr,mr,hr,c,zcr,d:hr>0.35 and c>4000 and d>=0.15),
+    ("cymbal",lambda lr,mr,hr,c,zcr,d:hr>0.25 and c>3000),
+    ("snare",lambda lr,mr,hr,c,zcr,d:mr>0.4 and zcr>0.1 and c>1000),
+    ("tom",lambda lr,mr,hr,c,zcr,d:lr>0.3 and mr>0.3 and c<1500),
+    ("bass",lambda lr,mr,hr,c,zcr,d:lr>0.6 and c<400 and d>0.2),
+    ("vocal",lambda lr,mr,hr,c,zcr,d:mr>0.5 and 500<c<3000 and zcr<0.15),
+    ("bright",lambda lr,mr,hr,c,zcr,d:c>2500),("mid",lambda lr,mr,hr,c,zcr,d:c>800)]
 def classify_hit(hit):
-    y,sr = hit.audio, hit.sr
-    D = np.abs(librosa.stft(y, n_fft=2048))
-    f = librosa.fft_frequencies(sr=sr, n_fft=2048)
+    y,sr=hit.audio,hit.sr; D=np.abs(librosa.stft(y,n_fft=2048))
+    f=librosa.fft_frequencies(sr=sr,n_fft=2048)
     le=np.sum(D[(f>=20)&(f<200)]**2); me=np.sum(D[(f>=200)&(f<4000)]**2)
-    he=np.sum(D[(f>=4000)]**2); t=le+me+he+1e-10
-    lr,mr,hr = le/t,me/t,he/t
-    zcr = float(librosa.feature.zero_crossing_rate(y=y).mean())
-    for name, fn in LABEL_RULES:
-        if fn(lr,mr,hr,hit.spectral_centroid,zcr,hit.duration): return name
+    he=np.sum(D[(f>=4000)]**2); t=le+me+he+1e-10; lr,mr,hr=le/t,me/t,he/t
+    zcr=float(librosa.feature.zero_crossing_rate(y=y).mean())
+    for n,fn in LABEL_RULES:
+        if fn(lr,mr,hr,hit.spectral_centroid,zcr,hit.duration): return n
     return "other"
-
 def classify_hits(hits):
-    print(f"[Stage 3] Classifying {len(hits)} hits...")
-    for h in hits: h.label = classify_hit(h)
-    counts = defaultdict(int)
-    for h in hits: counts[h.label] += 1
-    for l, c in sorted(counts.items(), key=lambda x: -x[1]): print(f"    {l}: {c}")
+    print(f"[Stage 3] Classifying {len(hits)}...")
+    for h in hits: h.label=classify_hit(h)
+    c=defaultdict(int)
+    for h in hits: c[h.label]+=1
+    for l,n in sorted(c.items(),key=lambda x:-x[1]): print(f"    {l}: {n}")
     return hits
 
-
-# ─── Stage 4: NCC-based clustering ───────────────────────────────────────────
-
-def ncc_max(a, b):
-    n = min(len(a), len(b)); a,b = a[:n].copy(), b[:n].copy()
-    a-=a.mean(); b-=b.mean()
-    norm = np.sqrt(np.dot(a,a)*np.dot(b,b))
+# ─── Stage 4 ──────────────────────────────────────────────────────────────────
+def ncc_max(a,b):
+    n=min(len(a),len(b)); a,b=a[:n].copy(),b[:n].copy(); a-=a.mean(); b-=b.mean()
+    norm=np.sqrt(np.dot(a,a)*np.dot(b,b))
     if norm<1e-10: return 0.0
     return float(np.max(np.abs(fftconvolve(a,b[::-1],mode='full'))))/norm
-
-def build_ncc_distance_matrix(hits, max_compare_samples=8820):
-    key = ("ncc_dist", tuple(_audio_hash(h.audio) for h in hits), max_compare_samples)
-    c = cache_get(key)
+def build_ncc_distance_matrix(hits,max_compare_samples=8820):
+    key=("ncc_dist",tuple(_audio_hash(h.audio) for h in hits),max_compare_samples)
+    c=cache_get(key)
     if c is not None: print(f"  Cached NCC matrix"); return c
     N=len(hits); D=np.zeros((N,N),dtype=np.float32)
     for i in range(N):
         ai=hits[i].audio[:max_compare_samples]
-        for j in range(i+1,N):
-            D[i,j]=D[j,i]=max(0.0, 1.0-ncc_max(ai,hits[j].audio[:max_compare_samples]))
-    return cache_set(key, D)
-
-def _labels_to_clusters(labels, hits):
-    cm = defaultdict(list)
+        for j in range(i+1,N): D[i,j]=D[j,i]=max(0.0,1.0-ncc_max(ai,hits[j].audio[:max_compare_samples]))
+    return cache_set(key,D)
+def _labels_to_clusters(labels,hits):
+    cm=defaultdict(list)
     for i,l in enumerate(labels): cm[l].append(i)
-    clusters = []
-    for _, idx in sorted(cm.items()):
-        v = defaultdict(int)
+    clusters=[]
+    for _,idx in sorted(cm.items()):
+        v=defaultdict(int)
         for i in idx: v[hits[i].label]+=1
-        maj = max(v, key=v.get)
-        ex = sum(1 for c in clusters if c.label.rsplit('_',1)[0]==maj)
-        clusters.append(Cluster(cluster_id=len(clusters), label=f"{maj}_{ex}",
-                                 hits=[hits[i] for i in idx]))
-    clusters.sort(key=lambda c: c.count, reverse=True)
+        maj=max(v,key=v.get); ex=sum(1 for c in clusters if c.label.rsplit('_',1)[0]==maj)
+        clusters.append(Cluster(cluster_id=len(clusters),label=f"{maj}_{ex}",hits=[hits[i] for i in idx]))
+    clusters.sort(key=lambda c:c.count,reverse=True)
     for i,c in enumerate(clusters): c.cluster_id=i
     return clusters
-
-def cluster_hits(hits, ncc_threshold=0.80, max_compare_ms=0,
-                 target_min=0, target_max=0, linkage='average'):
-    from sklearn.cluster import AgglomerativeClustering
+def cluster_hits(hits,ncc_threshold=0.80,max_compare_ms=0,target_min=0,target_max=0,linkage='average'):
+    from sklearn.cluster import AgglomerativeClustering as AC
     if not hits: return []
     N=len(hits); sr=hits[0].sr
-    if N==1: return [Cluster(cluster_id=0, label=f"{hits[0].label}_0", hits=[hits[0]])]
-    if max_compare_ms<=0:
-        ms = max(int(0.03*sr), int(np.median([len(h.audio) for h in hits])))
-    else: ms = int(max_compare_ms/1000.0*sr)
-    print(f"[Stage 4] NCC clustering ({N} hits, {ms/sr*1000:.0f}ms, {linkage})...")
-    print(f"  Computing {N*(N-1)//2} pairwise distances...")
-    D = build_ncc_distance_matrix(hits, max_compare_samples=ms)
-    use_t = target_min>0 and target_max>0 and target_max>=target_min
-    tmin=max(1,min(target_min or 1,N)); tmax=max(tmin,min(target_max or N,N))
+    if N==1: return [Cluster(cluster_id=0,label=f"{hits[0].label}_0",hits=[hits[0]])]
+    ms=max(int(0.03*sr),int(np.median([len(h.audio) for h in hits]))) if max_compare_ms<=0 else int(max_compare_ms/1000.0*sr)
+    print(f"[Stage 4] NCC ({N} hits, {ms/sr*1000:.0f}ms, {linkage})...")
+    print(f"  {N*(N-1)//2} pairs..."); D=build_ncc_distance_matrix(hits,max_compare_samples=ms)
+    use_t=target_min>0 and target_max>0 and target_max>=target_min
+    tmin,tmax=max(1,min(target_min or 1,N)),max(max(1,target_min or 1),min(target_max or N,N))
     if use_t:
         print(f"  Target: {tmin}–{tmax}")
         lo,hi=0.001,1.0; bl,bn,bd=None,-1,0.5
         for _ in range(30):
-            mid=(lo+hi)/2
-            agg=AgglomerativeClustering(n_clusters=None,distance_threshold=max(0.001,mid),
-                                         metric='precomputed',linkage=linkage)
-            lb=agg.fit_predict(D); n=len(set(lb))
+            mid=(lo+hi)/2; lb=AC(n_clusters=None,distance_threshold=max(0.001,mid),metric='precomputed',linkage=linkage).fit_predict(D)
+            n=len(set(lb))
             if tmin<=n<=tmax: bl,bn,bd=lb,n,mid; break
             elif n>tmax: lo=mid
             else: hi=mid
             if bl is None or abs(n-(tmin+tmax)/2)<abs(bn-(tmin+tmax)/2): bl,bn,bd=lb,n,mid
         if bn<tmin or bn>tmax:
-            tm=min((tmin+tmax)//2, N-1)
-            print(f"  Fallback n_clusters={tm}")
-            try:
-                agg=AgglomerativeClustering(n_clusters=tm,metric='precomputed',linkage=linkage)
-                bl=agg.fit_predict(D); bn=tm
+            tm=min((tmin+tmax)//2,N-1); print(f"  Fallback n={tm}")
+            try: bl=AC(n_clusters=tm,metric='precomputed',linkage=linkage).fit_predict(D); bn=tm
             except: pass
-        labels=bl; print(f"  → {bn} clusters")
+        labels=bl; print(f"  → {bn}")
     else:
-        dt=max(0.001, 1.0-ncc_threshold); print(f"  Fixed: dist≤{dt:.3f}")
-        labels=AgglomerativeClustering(n_clusters=None,distance_threshold=dt,
-                                        metric='precomputed',linkage=linkage).fit_predict(D)
+        dt=max(0.001,1.0-ncc_threshold); print(f"  Fixed dist≤{dt:.3f}")
+        labels=AC(n_clusters=None,distance_threshold=dt,metric='precomputed',linkage=linkage).fit_predict(D)
     print(f"  ✓ {len(set(labels))} clusters")
-    cl = _labels_to_clusters(labels, hits)
-    for c in cl: print(f"    {c.label}: {c.count} hits")
+    cl=_labels_to_clusters(labels,hits)
+    for c in cl: print(f"    {c.label}: {c.count}")
     return cl
 
-
-# ─── Stage 5: Quality scoring ────────────────────────────────────────────────
-
-def sample_quality_score(y, sr, label="other"):
+# ─── Stage 5 ────────────────���─────────────────────────────────────────────────
+def sample_quality_score(y,sr,label="other"):
     import scipy.stats
-    rms_env=librosa.feature.rms(y=y,frame_length=512,hop_length=128)[0]
-    if len(rms_env)>=10:
-        pk=np.argmax(rms_env); post=rms_env[pk:]
-        c1=max(0,1.0-np.mean(post[-max(3,len(post)//5):])/( rms_env[pk]+1e-8)*5)
-        if len(post)>=5:
-            sl,_,r,_,_=scipy.stats.linregress(np.arange(len(post)),np.log(post+1e-8))
-            c2=max(0,r**2) if sl<0 else r**2*0.3
-        else: c2=0.0
+    re=librosa.feature.rms(y=y,frame_length=512,hop_length=128)[0]
+    if len(re)>=10:
+        pk=np.argmax(re); po=re[pk:]
+        c1=max(0,1.0-np.mean(po[-max(3,len(po)//5):])/( re[pk]+1e-8)*5)
+        c2=(max(0,scipy.stats.linregress(np.arange(len(po)),np.log(po+1e-8))[2]**2) if scipy.stats.linregress(np.arange(len(po)),np.log(po+1e-8))[0]<0 else 0.0) if len(po)>=5 else 0.0
     else: c1,c2=0.5,0.0
-    comp=c1*0.6+c2*0.4
-    snr=10*np.log10(np.percentile(y**2,99)/(np.percentile(y**2,10)+1e-12))
-    ns=np.clip((snr-10)/40,0,1)
-    ons=librosa.onset.onset_detect(y=y,sr=sr,units='samples',backtrack=True)
+    comp=c1*0.6+c2*0.4; snr=10*np.log10(np.percentile(y**2,99)/(np.percentile(y**2,10)+1e-12))
+    ns=np.clip((snr-10)/40,0,1); ons=librosa.onset.onset_detect(y=y,sr=sr,units='samples',backtrack=True)
     if len(ons)>0:
         o=int(ons[0]); pre=y[max(0,o-int(sr*.02)):o]; sig=y[o:o+int(sr*.1)]
-        np2=np.clip((-10*np.log10(np.mean(pre**2+1e-12)/np.mean(sig**2+1e-12))-5)/30,0,1) \
-            if len(pre)>10 and len(sig)>10 else 0.5
+        np2=np.clip((-10*np.log10(np.mean(pre**2+1e-12)/np.mean(sig**2+1e-12))-5)/30,0,1) if len(pre)>10 and len(sig)>10 else 0.5
     else: np2=0.5
-    clean=ns*0.5+np2*0.5
-    oe=librosa.onset.onset_strength(y=y,sr=sr)
-    sh=float(np.max(oe)/(np.mean(oe)+1e-8)) if len(oe)>1 else 1.0
-    oq=float(np.clip((sh-1.0)/5.0,0,1))
-    tot=(comp*0.30+clean*0.40+oq*0.20+0.5*0.10)*100
-    return {'total':float(tot),'completeness':float(comp),'cleanness':float(clean),'onset_quality':float(oq)}
-
+    cl=ns*0.5+np2*0.5; oe=librosa.onset.onset_strength(y=y,sr=sr)
+    oq=float(np.clip((float(np.max(oe)/(np.mean(oe)+1e-8)) if len(oe)>1 else 1.0-1.0)/5.0,0,1))
+    return {'total':float((comp*0.30+cl*0.40+oq*0.20+0.5*0.10)*100),'completeness':float(comp),'cleanness':float(cl),'onset_quality':float(oq)}
 def select_best(clusters):
     print(f"[Stage 5] Selecting best...")
     for c in clusters:
         if c.count<=1: c.best_hit_idx=0; continue
-        sc=[sample_quality_score(h.audio,h.sr,c.label.rsplit('_',1)[0])['total'] for h in c.hits]
-        c.best_hit_idx=int(np.argmax(sc))
-
-
-# ─── Stage 6: Synthesis ──────────────────────────────────────────────────────
+        c.best_hit_idx=int(np.argmax([sample_quality_score(h.audio,h.sr,c.label.rsplit('_',1)[0])['total'] for h in c.hits]))
 
+# ─── Stage 6 ──────────────────────────────────────────────────────────────────
 def synthesize_from_cluster(cluster):
     if cluster.count<2: return None
-    tl=int(np.median([len(h.audio) for h in cluster.hits]))
-    al,wt=[],[]
-    pp=None
+    tl=int(np.median([len(h.audio) for h in cluster.hits])); al,wt=[],[]; pp=None
     for i,h in enumerate(cluster.hits):
         a=h.audio.copy(); p=np.argmax(np.abs(a))
         if pp is None: pp=p
@@ -302,33 +227,24 @@ def synthesize_from_cluster(cluster):
         pk=np.abs(a).max()
         if pk>0: a=a/pk
         al.append(a); wt.append(2.0 if i==cluster.best_hit_idx else 1.0)
-    al=np.array(al); w=np.array(wt); w/=w.sum()
-    sy=np.average(al,axis=0,weights=w); pk=np.abs(sy).max()
+    al=np.array(al); w=np.array(wt); w/=w.sum(); sy=np.average(al,axis=0,weights=w); pk=np.abs(sy).max()
     return (sy*0.95/pk).astype(np.float32) if pk>0 else sy.astype(np.float32)
 
-
-# ─── Stage 7: MIDI + rendering ───────────────────────────────────────────────
-
-def build_midi(clusters, bpm=120.0):
-    import pretty_midi
-    pm=pretty_midi.PrettyMIDI(initial_tempo=bpm)
+# ─── Stage 7 ──────────────────────────────────────────────────────────────────
+def build_midi(clusters,bpm=120.0):
+    import pretty_midi; pm=pretty_midi.PrettyMIDI(initial_tempo=bpm)
     for i,c in enumerate(clusters): c.midi_note=min(36+i,127)
-    inst=pretty_midi.Instrument(program=0,is_drum=True,name='Extracted Samples')
-    pm.instruments.append(inst)
+    inst=pretty_midi.Instrument(program=0,is_drum=True,name='Samples'); pm.instruments.append(inst)
     for c in clusters:
         for h in c.hits:
-            v=max(1,min(127,int(h.rms_energy/0.3*127)))
-            inst.notes.append(pretty_midi.Note(velocity=v,pitch=c.midi_note,
-                start=h.onset_time,end=h.onset_time+max(h.duration,0.05)))
-    inst.notes.sort(key=lambda n: n.start); return pm
-
-def export_midi(clusters, path, bpm=120.0):
-    pm=build_midi(clusters,bpm); pm.write(path)
-    print(f"  ✓ MIDI: {path} ({len(pm.instruments[0].notes)} notes)"); return pm
-
-def detect_bpm(y, sr):
+            inst.notes.append(pretty_midi.Note(velocity=max(1,min(127,int(h.rms_energy/0.3*127))),
+                pitch=c.midi_note,start=h.onset_time,end=h.onset_time+max(h.duration,0.05)))
+    inst.notes.sort(key=lambda n:n.start); return pm
+def export_midi(clusters,path,bpm=120.0):
+    pm=build_midi(clusters,bpm); pm.write(path); print(f"  ✓ MIDI: {len(pm.instruments[0].notes)} notes"); return pm
+def detect_bpm(y,sr):
     ck=("bpm",_audio_hash(y),sr); c=cache_get(ck)
-    if c is not None: return c
+    if c: return c
     oe=librosa.onset.onset_strength(y=y,sr=sr,aggregate=np.median)
     bpm=float(librosa.feature.tempo(onset_envelope=oe,sr=sr).item())
     _,beats=librosa.beat.beat_track(onset_envelope=oe,sr=sr,units='time')
@@ -339,48 +255,34 @@ def detect_bpm(y, sr):
         else:
             if bpm<70: bpm*=2
             elif bpm>200: bpm/=2
-    return cache_set(ck, round(bpm,1))
-
-def render_midi_with_samples(clusters, sr=44100):
+    return cache_set(ck,round(bpm,1))
+def render_midi_with_samples(clusters,sr=44100):
     me=max((h.onset_time+h.duration for c in clusters for h in c.hits),default=1.0)
     buf=np.zeros(int((me+1.0)*sr),dtype=np.float64)
     for c in clusters:
-        s=c.best_hit.audio.astype(np.float64)
-        re=c.best_hit.rms_energy if c.best_hit.rms_energy>0 else 0.1
+        s=c.best_hit.audio.astype(np.float64); re=c.best_hit.rms_energy if c.best_hit.rms_energy>0 else 0.1
         for h in c.hits:
-            vs=min(2.0,h.rms_energy/(re+1e-8))**0.5
-            i=int(h.onset_time*sr); e=i+len(s)
+            vs=min(2.0,h.rms_energy/(re+1e-8))**0.5; i=int(h.onset_time*sr); e=i+len(s)
             if e>len(buf): buf=np.concatenate([buf,np.zeros(e-len(buf))])
             buf[i:e]+=s*vs
-    pk=np.abs(buf).max()
-    return (buf/pk*0.9).astype(np.float32) if pk>1e-8 else buf.astype(np.float32)
-
+    pk=np.abs(buf).max(); return (buf/pk*0.9).astype(np.float32) if pk>1e-8 else buf.astype(np.float32)
 def build_sample_map(clusters):
-    return {c.midi_note:{'label':c.label,'count':c.count,
-            'duration_ms':int(c.best_hit.duration*1000)} for c in clusters}
-
-def build_archive(clusters, bpm, sr, midi_path=None, rendered_audio=None):
-    import zipfile, tempfile, io
-    zp=tempfile.mktemp(suffix='.zip')
-    idx={'bpm':round(bpm,1),'sample_rate':sr,'total_clusters':len(clusters),
-         'total_hits':sum(c.count for c in clusters),'samples':{}}
+    return {c.midi_note:{'label':c.label,'count':c.count,'duration_ms':int(c.best_hit.duration*1000)} for c in clusters}
+def build_archive(clusters,bpm,sr,midi_path=None,rendered_audio=None):
+    import zipfile,tempfile,io; zp=tempfile.mktemp(suffix='.zip')
+    idx={'bpm':round(bpm,1),'sample_rate':sr,'total_clusters':len(clusters),'total_hits':sum(c.count for c in clusters),'samples':{}}
     with zipfile.ZipFile(zp,'w',compression=zipfile.ZIP_STORED) as zf:
         for c in clusters:
-            b=c.best_hit; fn=f"samples/{c.label}.wav"
-            buf=io.BytesIO(); sf.write(buf,b.audio,sr,format='WAV',subtype='PCM_24')
-            zf.writestr(fn,buf.getvalue())
+            b=c.best_hit; fn=f"samples/{c.label}.wav"; buf=io.BytesIO()
+            sf.write(buf,b.audio,sr,format='WAV',subtype='PCM_24'); zf.writestr(fn,buf.getvalue())
             ot=sorted([h.onset_time for h in c.hits])
-            idx['samples'][c.label]={
-                'file':fn,'classification':c.label.rsplit('_',1)[0],
-                'midi_note':c.midi_note,'occurrences':c.count,
-                'onset_times_sec':[round(t,4) for t in ot],
-                'duration_sec':round(b.duration,4),
-                'rms_energy':round(b.rms_energy,6),
-                'spectral_centroid_hz':round(b.spectral_centroid,1)}
+            idx['samples'][c.label]={'file':fn,'classification':c.label.rsplit('_',1)[0],'midi_note':c.midi_note,
+                'occurrences':c.count,'onset_times_sec':[round(t,4) for t in ot],'duration_sec':round(b.duration,4),
+                'rms_energy':round(b.rms_energy,6),'spectral_centroid_hz':round(b.spectral_centroid,1)}
             if c.synthesized is not None:
-                sf2=f"samples/{c.label}__synthesized.wav"; b2=io.BytesIO()
-                sf.write(b2,c.synthesized,sr,format='WAV',subtype='PCM_24')
-                zf.writestr(sf2,b2.getvalue()); idx['samples'][c.label]['synthesized_file']=sf2
+                sf2=f"samples/{c.label}__synth.wav"; b2=io.BytesIO()
+                sf.write(b2,c.synthesized,sr,format='WAV',subtype='PCM_24'); zf.writestr(sf2,b2.getvalue())
+                idx['samples'][c.label]['synthesized_file']=sf2
         zf.writestr('index.json',json.dumps(idx,indent=2))
         if midi_path and os.path.exists(midi_path): zf.write(midi_path,'reconstruction.mid')
         if rendered_audio is not None:
@@ -388,66 +290,57 @@ def build_archive(clusters, bpm, sr, midi_path=None, rendered_audio=None):
             zf.writestr('rendered_reconstruction.wav',rb.getvalue())
     return zp
 
+# ─── Auto-tuner with parameter locking ───────────────────────────────────────
 
-# ─── Auto-tuner ──────────────────────────────────────────────────────────────
-
-def _spectral_envelope_corr(original, rendered, sr, n_fft=4096, hop=2048):
-    """Spectral envelope correlation between two signals. Phase-insensitive.
-    Returns correlation coefficient [−1, 1]. Higher = better reconstruction."""
-    n = min(len(original), len(rendered))
-    if n < n_fft: return 0.0
-    S_orig = np.abs(librosa.stft(original[:n], n_fft=n_fft, hop_length=hop))
-    S_rend = np.abs(librosa.stft(rendered[:n], n_fft=n_fft, hop_length=hop))
-    # Average over time to get spectral envelope
-    env_o = S_orig.mean(axis=1)
-    env_r = S_rend.mean(axis=1)
-    if env_o.std() < 1e-10 or env_r.std() < 1e-10: return 0.0
-    return float(np.corrcoef(env_o, env_r)[0, 1])
-
+def _spectral_envelope_corr(orig,rend,sr,n_fft=4096,hop=2048):
+    n=min(len(orig),len(rend))
+    if n<n_fft: return 0.0
+    eo=np.abs(librosa.stft(orig[:n],n_fft=n_fft,hop_length=hop)).mean(axis=1)
+    er=np.abs(librosa.stft(rend[:n],n_fft=n_fft,hop_length=hop)).mean(axis=1)
+    if eo.std()<1e-10 or er.std()<1e-10: return 0.0
+    return float(np.corrcoef(eo,er)[0,1])
 
-def _rms_envelope_corr(original, rendered, sr, hop=1024):
-    """RMS amplitude envelope correlation. Measures timing/dynamics match."""
-    n = min(len(original), len(rendered))
-    if n < hop * 4: return 0.0
-    rms_o = librosa.feature.rms(y=original[:n], hop_length=hop)[0]
-    rms_r = librosa.feature.rms(y=rendered[:n], hop_length=hop)[0]
-    n2 = min(len(rms_o), len(rms_r))
-    if n2 < 4: return 0.0
-    rms_o, rms_r = rms_o[:n2], rms_r[:n2]
-    if rms_o.std() < 1e-10 or rms_r.std() < 1e-10: return 0.0
-    return float(np.corrcoef(rms_o, rms_r)[0, 1])
+def _rms_envelope_corr(orig,rend,sr,hop=1024):
+    n=min(len(orig),len(rend))
+    if n<hop*4: return 0.0
+    ro=librosa.feature.rms(y=orig[:n],hop_length=hop)[0]; rr=librosa.feature.rms(y=rend[:n],hop_length=hop)[0]
+    n2=min(len(ro),len(rr))
+    if n2<4: return 0.0
+    ro,rr=ro[:n2],rr[:n2]
+    if ro.std()<1e-10 or rr.std()<1e-10: return 0.0
+    return float(np.corrcoef(ro,rr)[0,1])
 
+def _reconstruction_score(orig,rend,sr):
+    sc=_spectral_envelope_corr(orig,rend,sr); rc=_rms_envelope_corr(orig,rend,sr)
+    lr=min(len(rend),len(orig))/(max(len(rend),len(orig))+1)
+    return max(0.0,(sc*0.5+rc*0.4+lr*0.1)*100)
 
-def _reconstruction_score(original, rendered, sr):
-    """Combined score [0, 100] measuring how well the reconstruction matches."""
-    spec_corr = _spectral_envelope_corr(original, rendered, sr)
-    rms_corr = _rms_envelope_corr(original, rendered, sr)
-    # Penalize if reconstruction is much shorter/longer
-    len_ratio = min(len(rendered), len(original)) / (max(len(rendered), len(original)) + 1)
-    score = (spec_corr * 0.5 + rms_corr * 0.4 + len_ratio * 0.1) * 100
-    return max(0.0, score)
 
+def auto_tune(stem_audio, sr, mode="auto", locks=None, log_fn=None):
+    """Find best extraction parameters. Locked params are held constant.
 
-def auto_tune(stem_audio, sr, mode="auto", log_fn=None):
-    """Automatically find the best extraction parameters for this audio.
+    locks: dict mapping param names to fixed values. Supported keys:
+        'onset_delta', 'energy_threshold_db', 'min_gap',
+        'target_min', 'target_max'
+    Any key present in locks will NOT be swept — its value is used as-is.
 
-    Sweeps onset detection params and cluster counts. Uses the cached NCC matrix
-    so re-clustering is near-instant after the first onset detection.
+    Example: locks={'target_min': 5, 'target_max': 20}
+      → cluster count will always be between 5 and 20, only onset params are tuned.
 
-    Returns: (best_params: dict, best_score: float, log: list[str])
-
-    The returned params can be passed directly to the extraction pipeline.
+    Returns: (best_params, best_score, log_lines)
     """
+    locks = locks or {}
     log = []
-    def _log(msg):
-        log.append(msg)
-        if log_fn: log_fn(msg)
-        print(msg)
+    def _log(m):
+        log.append(m)
+        if log_fn: log_fn(m)
+        print(m)
 
-    _log(f"[Auto-tune] Starting parameter search on {len(stem_audio)/sr:.1f}s audio...")
+    locked_names = list(locks.keys())
+    _log(f"[Auto-tune] {len(stem_audio)/sr:.1f}s audio, locked: {locked_names or 'none'}")
 
-    # Parameter grid — coarse sweep first
-    onset_configs = [
+    # Build onset config grid — replace locked values with fixed ones
+    base_onset_configs = [
         {"onset_delta": 0.08, "energy_threshold_db": -40, "min_gap": 0.02},
         {"onset_delta": 0.10, "energy_threshold_db": -35, "min_gap": 0.025},
         {"onset_delta": 0.12, "energy_threshold_db": -35, "min_gap": 0.03},
@@ -456,85 +349,88 @@ def auto_tune(stem_audio, sr, mode="auto", log_fn=None):
         {"onset_delta": 0.25, "energy_threshold_db": -25, "min_gap": 0.06},
     ]
 
-    cluster_targets = [3, 5, 8, 10, 15, 20, 30]
+    # Apply locks to onset configs
+    onset_configs = []
+    seen = set()
+    for oc in base_onset_configs:
+        for k in ['onset_delta', 'energy_threshold_db', 'min_gap']:
+            if k in locks:
+                oc[k] = locks[k]
+        # Deduplicate configs that become identical after locking
+        key = (oc['onset_delta'], oc['energy_threshold_db'], oc['min_gap'])
+        if key not in seen:
+            seen.add(key)
+            onset_configs.append(oc)
+
+    # Build cluster target grid — respect locks
+    if 'target_min' in locks and 'target_max' in locks:
+        tmin_locked, tmax_locked = int(locks['target_min']), int(locks['target_max'])
+        # Generate targets within the locked range
+        cluster_targets = sorted(set([tmin_locked, tmax_locked,
+            (tmin_locked+tmax_locked)//2,
+            tmin_locked + (tmax_locked-tmin_locked)//3,
+            tmin_locked + 2*(tmax_locked-tmin_locked)//3]))
+        cluster_targets = [t for t in cluster_targets if tmin_locked <= t <= tmax_locked]
+        _log(f"  Cluster targets (locked {tmin_locked}–{tmax_locked}): {cluster_targets}")
+    elif 'target_min' in locks:
+        tmin_locked = int(locks['target_min'])
+        cluster_targets = [t for t in [3,5,8,10,15,20,30] if t >= tmin_locked]
+        if not cluster_targets: cluster_targets = [tmin_locked]
+    elif 'target_max' in locks:
+        tmax_locked = int(locks['target_max'])
+        cluster_targets = [t for t in [3,5,8,10,15,20,30] if t <= tmax_locked]
+        if not cluster_targets: cluster_targets = [tmax_locked]
+    else:
+        cluster_targets = [3, 5, 8, 10, 15, 20, 30]
 
-    best_score = -1
-    best_params = {}
-    best_clusters = None
-    results = []
+    best_score, best_params = -1, {}
 
     for oc_idx, oc in enumerate(onset_configs):
-        _log(f"\n  Config {oc_idx+1}/{len(onset_configs)}: delta={oc['onset_delta']}, "
-             f"energy={oc['energy_threshold_db']}dB, gap={oc['min_gap']}s")
-
+        _log(f"\n  [{oc_idx+1}/{len(onset_configs)}] delta={oc['onset_delta']}, "
+             f"energy={oc['energy_threshold_db']}dB, gap={oc['min_gap']}")
         hits = detect_onsets(stem_audio, sr, mode=mode, **oc)
-        if len(hits) < 2:
-            _log(f"    → Only {len(hits)} hits, skipping")
-            continue
-
+        if len(hits) < 2: _log(f"    {len(hits)} hits, skip"); continue
         hits = classify_hits(hits)
 
-        # NCC matrix computed once per onset config (cached within)
-        for n_target in cluster_targets:
-            if n_target >= len(hits): continue
-
-            clusters = cluster_hits(hits, target_min=n_target, target_max=n_target,
-                                     linkage='average')
+        for nt in cluster_targets:
+            if nt >= len(hits): continue
+            clusters = cluster_hits(hits, target_min=nt, target_max=nt, linkage='average')
             if not clusters: continue
-
-            # Quick select best + render
             for c in clusters:
-                if c.count <= 1: c.best_hit_idx = 0
-                else:
-                    energies = [h.rms_energy for h in c.hits]
-                    c.best_hit_idx = int(np.argmax(energies))  # fast: pick loudest
-
+                if c.count<=1: c.best_hit_idx=0
+                else: c.best_hit_idx=int(np.argmax([h.rms_energy for h in c.hits]))
             rendered = render_midi_with_samples(clusters, sr=sr)
             score = _reconstruction_score(stem_audio, rendered, sr)
-
-            results.append({**oc, 'n_clusters': len(clusters),
-                            'target': n_target, 'n_hits': len(hits), 'score': score})
-
             if score > best_score:
                 best_score = score
-                best_params = {**oc, 'n_clusters': len(clusters), 'target_min': n_target,
-                               'target_max': n_target}
-                best_clusters = clusters
-                _log(f"    ★ target={n_target} → {len(clusters)} clusters, "
-                     f"score={score:.1f} (NEW BEST)")
+                best_params = {**oc, 'n_clusters': len(clusters),
+                               'target_min': nt, 'target_max': nt}
+                _log(f"    ★ target={nt} → {len(clusters)} cl, score={score:.1f} BEST")
             else:
-                _log(f"    target={n_target} → {len(clusters)} clusters, score={score:.1f}")
+                _log(f"    target={nt} → {len(clusters)} cl, score={score:.1f}")
 
-    # Fine-tune: try ±1 around best target with best onset config
+    # Fine-tune around best
     if best_params:
         bt = best_params.get('target_min', 10)
-        _log(f"\n  Fine-tuning around best (delta={best_params['onset_delta']}, "
-             f"target≈{bt})...")
-
-        fine_oc = {k: best_params[k] for k in ['onset_delta', 'energy_threshold_db', 'min_gap']}
+        fine_oc = {k: best_params[k] for k in ['onset_delta','energy_threshold_db','min_gap']}
+        _log(f"\n  Fine-tuning ±3 around target={bt}...")
         hits = detect_onsets(stem_audio, sr, mode=mode, **fine_oc)
         if len(hits) >= 2:
             hits = classify_hits(hits)
-            for ft in range(max(2, bt-3), bt+4):
-                if ft >= len(hits): continue
+            lo = max(2, int(locks.get('target_min', bt-3)))
+            hi = min(len(hits)-1, int(locks.get('target_max', bt+3)))
+            for ft in range(lo, hi+1):
                 clusters = cluster_hits(hits, target_min=ft, target_max=ft, linkage='average')
                 if not clusters: continue
                 for c in clusters:
-                    if c.count <= 1: c.best_hit_idx = 0
-                    else: c.best_hit_idx = int(np.argmax([h.rms_energy for h in c.hits]))
+                    if c.count<=1: c.best_hit_idx=0
+                    else: c.best_hit_idx=int(np.argmax([h.rms_energy for h in c.hits]))
                 rendered = render_midi_with_samples(clusters, sr=sr)
                 score = _reconstruction_score(stem_audio, rendered, sr)
                 if score > best_score:
-                    best_score = score
-                    best_params = {**fine_oc, 'n_clusters': len(clusters),
-                                   'target_min': ft, 'target_max': ft}
-                    best_clusters = clusters
-                    _log(f"    ★ target={ft} → {len(clusters)} clusters, "
-                         f"score={score:.1f} (NEW BEST)")
-
-    _log(f"\n[Auto-tune] Best: score={best_score:.1f}, "
-         f"delta={best_params.get('onset_delta')}, "
-         f"energy={best_params.get('energy_threshold_db')}dB, "
-         f"clusters={best_params.get('n_clusters')}")
+                    best_score=score
+                    best_params={**fine_oc,'n_clusters':len(clusters),'target_min':ft,'target_max':ft}
+                    _log(f"    ★ target={ft} → {len(clusters)} cl, score={score:.1f} BEST")
 
+    _log(f"\n[Auto-tune] Best: score={best_score:.1f}, params={best_params}")
     return best_params, best_score, log