Spaces:
Sleeping
Sleeping
| """Incrementally FILL latents/restored/{stem}.pt by running the v4b restorer over cached | |
| degraded stem latents (+ degraded_mix). Latent->latent, no decode -> fast. Skips existing. | |
| These restored latents are the realistic CONTEXT for the generator: at inference we never have | |
| clean stems, only v4b restorations. Run: python -m restoflow.fill_restored --root demucs_results_all | |
| """ | |
| from __future__ import annotations | |
| import argparse, csv, os | |
| from pathlib import Path | |
| import torch | |
| from .config import Cfg, STEM_ID, STEMS | |
| from .model import DetRestorer, AttnRestorer | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--root", required=True) | |
| ap.add_argument("--ckpt", default="restoflow_runs/restorer_attn_distvar_w003/ckpt_best.pt") | |
| ap.add_argument("--stats", default="restoflow_runs/restorer_attn_distvar_w003/norm_stats.pt") | |
| ap.add_argument("--device", default="cuda") | |
| ap.add_argument("--overwrite", action="store_true", help="regenerate even if restored/{stem}.pt exists") | |
| a = ap.parse_args() | |
| c = Cfg(); dev = a.device | |
| BASE = Path("/media/maindisk/ksoil/repos/Restoration/Dac_demucs_lyra") | |
| root = BASE / a.root | |
| ck = torch.load(BASE / a.ckpt, map_location="cpu"); cfg = ck["cfg"] | |
| saved = {k: v for k, v in cfg.items() if k in ("latent_dim", "hidden", "depth", "stem_emb_dim", "use_mix")} | |
| if cfg.get("model_kind") == "attn": # advramp/cotrain/w003 are AttnRestorer; v4b was DetRestorer | |
| model = AttnRestorer(saved.get("latent_dim", 256), saved["hidden"], saved["depth"], n_stems=len(STEMS), | |
| stem_emb=saved["stem_emb_dim"], use_mix=saved["use_mix"], heads=cfg.get("heads", 4)) | |
| else: | |
| model = DetRestorer(saved.get("latent_dim", 256), saved["hidden"], saved["depth"], | |
| n_stems=len(STEMS), stem_emb=saved["stem_emb_dim"], use_mix=saved["use_mix"]) | |
| model.load_state_dict(ck["model"]); model.eval().to(dev) | |
| stats = torch.load(BASE / a.stats, map_location="cpu") | |
| mmu, msd = stats["__mix__"]["mu"], stats["__mix__"]["sd"] | |
| print(f"[fill] v4b loaded; root={root}") | |
| rows = list(csv.DictReader(open(root / "metadata.csv"))) | |
| made = skipped = nomix = 0 | |
| with torch.inference_mode(): | |
| for r in rows: | |
| stem = r.get("stem") | |
| if stem not in STEMS or r.get("deg_present") != "1" or not r.get("deg_latent_path"): | |
| continue | |
| pair = r["pair_id"] | |
| out = root / pair / "latents" / "restored" / f"{stem}.pt" | |
| if out.exists() and not a.overwrite: | |
| skipped += 1; continue | |
| mix_p = root / pair / "latents" / "degraded_mix.pt" | |
| if not mix_p.exists(): | |
| nomix += 1; continue | |
| deg = torch.load(root / r["deg_latent_path"], map_location="cpu").float() | |
| mix = torch.load(mix_p, map_location="cpu").float() | |
| mu, sd = stats[stem]["mu"], stats[stem]["sd"] | |
| std_d = ((deg - mu[:, None]) / sd[:, None]).to(dev)[None] | |
| std_m = ((mix - mmu[:, None]) / msd[:, None]).to(dev)[None] | |
| sid = torch.tensor([STEM_ID[stem]], device=dev) | |
| std_r = model(std_d, sid, std_m)[0].cpu() | |
| raw_r = std_r * sd[:, None] + mu[:, None] | |
| out.parent.mkdir(parents=True, exist_ok=True) | |
| torch.save(raw_r.half(), out); made += 1 | |
| if made % 2000 == 0: | |
| print(f" made={made} skipped={skipped}") | |
| print(f"[fill] done. made={made} skipped(existing)={skipped} no_mix={nomix}") | |
| if __name__ == "__main__": | |
| main() | |