Add files using upload-large-folder tool

code/models/autoports/minimaxai_minimax_music3/reference/pipeline.py

@@ -230,10 +230,13 @@ class Music3Reference:
         guidance_scale: float = DIT_GUIDANCE_SCALE,
         record_steps: Sequence[int] = (),
         dit_forward=None,
+        on_chunk=None,
     ) -> dict:
         """before_denoise.py + denoise.py. Returns dict(latent_chunks [list of [1,128,L]], chunk_starts, and when
         record_steps: per-chunk dicts with condition / noise / timesteps / (t, latents_in, pred_cond, pred_uncond)).
-        `dit_forward(latents [B,128,L], t [B], cond [B,L,2048]) -> velocity` lets a TT DiT drive the same loop."""
+        `dit_forward(latents [B,128,L], t [B], cond [B,L,2048]) -> velocity` lets a TT DiT drive the same loop.
+        `on_chunk(k, latents)` is called as soon as window k is final (the caller may start vocoding it while the
+        next window denoises); the latents handed over are never mutated afterwards."""
         dit_forward = dit_forward or (lambda x, t, c: self.dit(x, t, c))
         starts = chunk_starts(frame_hiddens.shape[1])
         timesteps, sigmas = flow_schedule(num_steps)
@@ -288,25 +291,37 @@ class Music3Reference:
             oe = max(os_, latents.shape[-1] - OVERLAP_LATENT_LENGTH)
             prev_latent, prev_cond = latents[..., os_:oe], condition[:, os_:oe]
             latent_chunks.append(latents)
+            if on_chunk is not None:
+                on_chunk(k, latents)
             if rec is not None:
                 rec["latents_out"] = latents.detach().cpu()
                 chunk_records.append(rec)
         return {"latent_chunks": latent_chunks, "chunk_starts": starts, "chunk_records": chunk_records}
 
     @torch.no_grad()
-    def decode(self, latent_chunks: List[torch.Tensor], vocoder=None) -> torch.Tensor:
-        """decoders.py: vocode each window, crop overlaps, stitch -> [1, 2, samples] float32 in [-1, 1] at 44.1 kHz."""
+    def vocode(self, latents: torch.Tensor, vocoder=None) -> torch.Tensor:
+        """decoders.py, one window: latents [1,128,L] -> uncropped waveform [1, 2, L*hop]."""
+        vocoder = vocoder or self.vocoder
+        return vocoder(latents.to(vocoder.dec_in_proj.weight.dtype).to(self.device))
+
+    @torch.no_grad()
+    def stitch(self, wavs: List[torch.Tensor], vocoder=None) -> torch.Tensor:
+        """decoders.py: crop the window overlaps and concatenate -> [1, 2, samples] float32 in [-1, 1] at 44.1 kHz."""
         vocoder = vocoder or self.vocoder
         hop = vocoder.hop_length
-        n = len(latent_chunks)
+        n = len(wavs)
         chunks = []
-        for i, lat in enumerate(latent_chunks):
-            wav = vocoder(lat.to(vocoder.dec_in_proj.weight.dtype).to(self.device))
+        for i, wav in enumerate(wavs):
             left = 0 if i == 0 else CROP_LEFT_LATENT * hop
             right = 0 if i == n - 1 else CROP_RIGHT_LATENT * hop
             chunks.append(wav[..., left : wav.shape[-1] - right])
         return torch.cat(chunks, dim=-1).float().clamp(-1.0, 1.0).cpu()
 
+    @torch.no_grad()
+    def decode(self, latent_chunks: List[torch.Tensor], vocoder=None) -> torch.Tensor:
+        """decoders.py: vocode each window, crop overlaps, stitch -> [1, 2, samples] float32 in [-1, 1] at 44.1 kHz."""
+        return self.stitch([self.vocode(lat, vocoder) for lat in latent_chunks], vocoder)
+
     # ------------------------------------------------------------------ end to end
     @torch.no_grad()
     def generate(

code/models/autoports/minimaxai_minimax_music3/tt/config_defaults.json

@@ -11,4 +11,4 @@
     "depth_top1": 0.9,
     "depth_top5": 0.98
   }
-}
+}

code/models/autoports/minimaxai_minimax_music3/tt/generator.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 
 import os
 import time
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, Dict, List, Optional
@@ -49,6 +50,7 @@ class GenStats:
     depth_s: float = 0.0
     denoise_s: float = 0.0
     decode_s: float = 0.0
+    vocoder_s: float = 0.0  # summed per-window vocoder compute (overlaps denoise_s when pipelined)
     total_s: float = 0.0
     audio_s: float = 0.0
     chunks: int = 0
@@ -183,6 +185,8 @@ class Music3Generator:
         ]
         if load_vocoder and self.vocoder_dtype != torch.float32:
             self.ref.vocoder = self.ref.vocoder.to(self.vocoder_dtype)
+        # vocode window k on a worker thread while the DiT denoises window k+1 (MUSIC3_PIPELINE_VOCODER=0 disables)
+        self.pipeline_vocoder = os.environ.get("MUSIC3_PIPELINE_VOCODER", "1") not in ("0", "false", "no")
         self.dit = None
         if load_dit and dit_device == "tt":
             from models.autoports.minimaxai_minimax_music3.tt.dit import TTDiT
@@ -366,9 +370,16 @@ class Music3Generator:
         return self.dit(latents, timestep, condition)
 
     @torch.inference_mode()
-    def denoise(self, frame_hiddens: torch.Tensor, generator, num_steps: int = DIT_NUM_STEPS, record_steps=()):
+    def denoise(
+        self, frame_hiddens: torch.Tensor, generator, num_steps: int = DIT_NUM_STEPS, record_steps=(), on_chunk=None
+    ):
         return self.ref.denoise(
-            frame_hiddens, generator, num_steps=num_steps, record_steps=record_steps, dit_forward=self.dit_forward
+            frame_hiddens,
+            generator,
+            num_steps=num_steps,
+            record_steps=record_steps,
+            dit_forward=self.dit_forward,
+            on_chunk=on_chunk,
         )
 
     @torch.inference_mode()
@@ -400,22 +411,53 @@ class Music3Generator:
         sem = self.semantic_generation(text_ids, mf, generator, stats=stats, on_frame=on_frame)
         out = {**sem, "text_ids": text_ids}
         if "denoise" in stages and "frame_hiddens" in sem:
+            pipelined = "decode" in stages and self.pipeline_vocoder
+            # Pipelined: the (host, torch) vocoder runs window k on one worker thread while the DiT denoises window
+            # k+1 on the chip. Each window is vocoded by the identical call in the identical order, and the crop /
+            # stitch happens once at the end exactly as in the reference decode(), so the audio bytes are unchanged
+            # (the torch.Generator is only consumed on the denoise side).
+            wav_futures: List = []
+            vocoder_s = [0.0]
+
+            def vocode_window(lat):
+                t = time.time()
+                with torch.inference_mode():  # inference_mode is thread-local
+                    wav = self.ref.vocode(lat)
+                vocoder_s[0] += time.time() - t
+                return wav
+
             t1 = time.time()
-            den = self.denoise(sem["frame_hiddens"], generator, num_steps=num_steps)
-            stats.denoise_s = time.time() - t1
+            if pipelined:
+                with ThreadPoolExecutor(max_workers=1, thread_name_prefix="music3-vocoder") as pool:
+                    den = self.denoise(
+                        sem["frame_hiddens"],
+                        generator,
+                        num_steps=num_steps,
+                        on_chunk=lambda k, lat: wav_futures.append(pool.submit(vocode_window, lat)),
+                    )
+                    stats.denoise_s = time.time() - t1
+                    t2 = time.time()
+                    wavs = [f.result() for f in wav_futures]
+                out["audio"] = self.ref.stitch(wavs)
+            else:
+                den = self.denoise(sem["frame_hiddens"], generator, num_steps=num_steps)
+                stats.denoise_s = time.time() - t1
+                t2 = time.time()
+                if "decode" in stages:
+                    out["audio"] = self.ref.stitch([vocode_window(lat) for lat in den["latent_chunks"]])
             stats.chunks = len(den["latent_chunks"])
             out.update(den)
             if "decode" in stages:
-                t2 = time.time()
-                out["audio"] = self.decode(den["latent_chunks"])
                 stats.decode_s = time.time() - t2
+                stats.vocoder_s = vocoder_s[0]
                 stats.audio_s = out["audio"].shape[-1] / self.cfg.vocoder.sampling_rate
                 out["sample_rate"] = self.cfg.vocoder.sampling_rate
         stats.total_s = time.time() - t0
         out["stats"] = stats
         self.log(
             f"generate: {stats.frames} frames, {stats.audio_s:.1f}s audio in {stats.total_s:.1f}s (RTF {stats.rtf:.2f}; prefill {stats.prefill_s:.1f}s, "
-            f"llm {stats.ms_per_frame_llm:.1f} ms/f, depth {stats.ms_per_frame_depth:.1f} ms/f, denoise {stats.denoise_s:.1f}s, vocoder {stats.decode_s:.1f}s)"
+            f"llm {stats.ms_per_frame_llm:.1f} ms/f, depth {stats.ms_per_frame_depth:.1f} ms/f, denoise {stats.denoise_s:.1f}s, "
+            f"vocoder {stats.vocoder_s:.1f}s{' pipelined' if self.pipeline_vocoder else ''}, decode tail {stats.decode_s:.1f}s)"
         )
         return out