Enable YuE2 CUDA graphs and expose stage timings

pipeline.py

@@ -288,6 +288,7 @@ def generate_all(source: str, instrument: str, styles: dict[str, str], seed: int
     outputs = {name: None for name in VARIANTS}
     errors = {}
     durations = {}
+    yue_timing = {}
     drum_count = 0
     drums_midi = None
 
@@ -301,6 +302,7 @@ def generate_all(source: str, instrument: str, styles: dict[str, str], seed: int
         metadata = {"instrument": instrument, "source_sha256": digest,
                     "source_seconds": round(duration, 3),
                     "result_seconds": dict(durations), "drum_hits": drum_count,
+                    "yue_timing": dict(yue_timing),
                     "errors": dict(errors), "complete": complete,
                     "source_timing_preserved": ["instrument", "instrument_drums", "contour"],
                     "yue_vocal_absence_guaranteed": False,
@@ -346,6 +348,9 @@ def generate_all(source: str, instrument: str, styles: dict[str, str], seed: int
         if raw.is_file():
             try:
                 deliver(name, raw)
+                receipt = folder / name / "yue.json"
+                if receipt.is_file():
+                    yue_timing[name] = json.loads(receipt.read_text())["timing"]
             except (ValueError, RuntimeError, OSError) as error:
                 errors[name] = str(error)
         else:

yue_worker.py

@@ -24,7 +24,7 @@ def run(request_path: Path, output: Path) -> None:
     pipe = YuE2Pipeline.from_pretrained(
         "m-a-p/YuE2-3B", vae="m-a-p/YuE2-Vae",
         revision=YUE_REV, vae_revision=VAE_REV,
-        device="cuda", backend="torch-eager", quantization="none",
+        device="cuda", backend="torch", quantization="none",
         offload_ar=False, memory_budget_gib=44, progress=True,
     )
     try:
@@ -44,6 +44,7 @@ def run(request_path: Path, output: Path) -> None:
                 sf.write(target / "raw.wav", audio, song.sample_rate, subtype="FLOAT")
                 (target / "yue.json").write_text(json.dumps({
                     "model": "m-a-p/YuE2-3B", "revision": YUE_REV,
+                    "backend": "torch", "seed": request["seed"],
                     "vae_revision": VAE_REV, "duration_s": len(audio) / song.sample_rate,
                     "truncated": song.truncated, "timing": song.timing,
                     "source_timeline": False, "vocal_absence_guaranteed": False,