feat(streaming): add backend-agnostic live audio editing

acestep/engine/stream.py

@@ -15,7 +15,7 @@
 import time
 from collections import OrderedDict
 from dataclasses import dataclass, field
-from typing import Callable, Dict, NamedTuple, Optional, List, Tuple, TYPE_CHECKING
+from typing import Any, Callable, Dict, NamedTuple, Optional, List, Tuple, TYPE_CHECKING
 
 from loguru import logger
 import torch
@@ -117,6 +117,12 @@ class SlotRequest:
     extra_conditions: List[SlotCondition] = field(default_factory=list)
     primary_temporal_weight: Optional[torch.Tensor] = None
     primary_step_range: Optional[Tuple[float, float]] = None
+    # Immutable control snapshot attached by the serving layer. The engine
+    # deliberately treats it as opaque; renderers recover it from
+    # ``last_finished_request`` so a decoded latent is composited with the
+    # edit state it was generated from, not whatever is live several ring
+    # ticks later.
+    audio_edit: Any = None
     # --- CFG (Phase 2) ---
     # Flat list of negative conditions for classifier-free guidance. When
     # set together with ``guidance_curve``, each step runs a second forward

acestep/nodes/diffusion_nodes.py

@@ -379,6 +379,7 @@ def _build_slot_request(
     x0_target_strength: float,
     x0_target_gate: float,
     guidance_curve,
+    audio_edit,
     device,
     dtype,
 ) -> SlotRequest:
@@ -441,6 +442,7 @@ def _build_slot_request(
         extra_conditions=extra_conditions,
         primary_temporal_weight=primary.temporal_weight,
         primary_step_range=primary.step_range,
+        audio_edit=audio_edit,
         neg_conditions=neg_conditions,
         guidance_curve=guidance_curve_t,
         rcfg_mode=rcfg_mode,
@@ -618,6 +620,14 @@ def get_definition(cls) -> NodeDefinition:
                     ),
                     hidden=True,
                 ),
+                NodeParam(
+                    name="audio_edit", type="any", default=None,
+                    description=(
+                        "Immutable live audio-edit snapshot injected by the "
+                        "streaming session for this ring slot."
+                    ),
+                    hidden=True,
+                ),
             ),
         )
 
@@ -841,6 +851,7 @@ def execute(self, **kwargs: Any) -> dict[str, Any]:
             guidance_curve=modulation.guidance_curve,
             rcfg_mode=kwargs.get("rcfg_mode"),
             cfg_rescale=kwargs.get("cfg_rescale"),
+            audio_edit=kwargs.get("audio_edit"),
             device=device,
             dtype=dtype,
         )

acestep/streaming/ace_backend.py

@@ -26,12 +26,19 @@
 import torch
 
 from acestep.engine.dcw import DCWAdvanced
+from acestep.engine.masking import LatentNoiseMask
 from acestep.engine.obs import logger
 from acestep.nodes.types import ChannelGuidanceEntry, Latent
 from acestep.nodes.interpolation import INTERPOLATIONS
 from acestep.nodes.vae_nodes import EmptyLatent, LatentBlend
 
 from acestep.streaming.diffusion_backend import DiffusionBackend
+from acestep.streaming.audio_edit import (
+    DISABLED_AUDIO_EDIT,
+    composite_window,
+    constrain_audio_edit,
+    regenerate_mask,
+)
 from acestep.streaming.generator_backend import (
     AudioChunk,
     AudioGeometry,
@@ -180,6 +187,7 @@ def __init__(
         walk_window_s=60.0,
         neg_conditioning=None,
         steering: SteeringController | None = None,
+        source_waveform=None,
     ):
         # The family codec is the engine Session: its windowed VAE
         # decode is what render_window()/render_full() drive. The
@@ -202,6 +210,9 @@ def __init__(
         self.k1_name = k1_name
         self.SEED = seed
         self.skip_threshold = skip_threshold
+        self._source_waveform = source_waveform
+        self._audio_edit = DISABLED_AUDIO_EDIT
+        self._emerged_audio_edit = DISABLED_AUDIO_EDIT
 
         # Walk-window mode: drive the DiT with a fixed-T window sliced
         # from a longer pre-encoded source so the 60s TRT engine can
@@ -337,6 +348,25 @@ def capabilities(self) -> Capabilities:
             curves=True,
             notes_conditioning=False,
             steering=self.steering.is_loaded,
+            audio_edit=True,
+            audio_edit_extend=True,
+            audio_edit_strength=True,
+        )
+
+    def handle_set_audio_edit(self, edit) -> None:
+        self._audio_edit = edit
+
+    def finalize_audio_edit_window(self, pcm, start_sample: int):
+        """Apply the current/request mask after runner edge crossfades."""
+        if self._source_waveform is None:
+            return pcm
+        edit = constrain_audio_edit(self._audio_edit, self._emerged_audio_edit)
+        return composite_window(
+            pcm,
+            start_sample=start_sample,
+            source=self._source_waveform(),
+            edit=edit,
+            sample_rate=SAMPLE_RATE,
         )
 
     def geometry(self) -> AudioGeometry:
@@ -635,6 +665,7 @@ def _prepare_tick(self, knobs: dict, ctx: TickContext) -> dict:
         :meth:`_generate` consumes; the mode dispatch / caching /
         timing skeleton lives on :class:`DiffusionBackend`."""
         raw = knobs
+        audio_edit = self._audio_edit
         walk_active = self._walk_active
         full_src_T = self._full_src_T
 
@@ -787,6 +818,8 @@ def _prepare_tick(self, knobs: dict, ctx: TickContext) -> dict:
         else:
             denoise = k1
             self.state.sde_curve_display = None
+        if audio_edit.enabled and audio_edit.source_mode == "structure":
+            denoise = audio_edit.strength
 
         # Source lock: x0_target_curve from client overrides the
         # scalar x0_target_strength knob. The latent is attached
@@ -875,6 +908,7 @@ def _prepare_tick(self, knobs: dict, ctx: TickContext) -> dict:
             "x0_target_curve": x0_target_curve,
             "initial_noise_curve": initial_noise_curve,
             "tick_kwargs": tick_kwargs,
+            "audio_edit": audio_edit,
             "echo": {
                 "k1": k1, "seed": seed, "feedback": feedback,
                 "fb_depth": fb_depth, "shift_val": shift_val,
@@ -885,13 +919,33 @@ def _prepare_tick(self, knobs: dict, ctx: TickContext) -> dict:
     def _generate(self, prep: dict):
         raw = prep["raw"]
         source_lat = prep["source_lat"]
-        return self.stream.tick(
+        edit = prep["audio_edit"]
+        source_tensor = (
+            source_lat if source_lat is not None
+            else prep["live_src_lat"].tensor
+        )
+        source = Latent(tensor=source_tensor)
+        if edit.enabled and edit.source_mode == "waveform":
+            mask = regenerate_mask(
+                edit,
+                total_frames=source_tensor.shape[1],
+                rate_hz=25.0,
+                offset_s=self._walk_chunk_start_s if self._walk_active else 0.0,
+                device=source_tensor.device,
+                dtype=source_tensor.dtype,
+            )
+            source = Latent(
+                tensor=source_tensor,
+                mask=LatentNoiseMask(
+                    mask=mask,
+                    original_latents=prep["live_src_lat"].tensor,
+                ),
+            )
+        result = self.stream.tick(
             denoise=prep["denoise"],
             seed=prep["seed"],
-            source_latent=(
-                Latent(tensor=source_lat) if source_lat is not None
-                else prep["live_src_lat"]
-            ),
+            source_latent=source,
+            audio_edit=edit,
             x0_target=prep["x0_tgt"],
             x0_target_curve=prep["x0_target_curve"],
             shift=self._current_shift,
@@ -909,6 +963,12 @@ def _generate(self, prep: dict):
             dcw_wavelet=str(raw.get("dcw_wavelet", "haar")),
             dcw_advanced=_build_dcw_advanced(raw),
         )
+        if result is not None:
+            request = getattr(self.stream.pipeline, "last_finished_request", None)
+            self._emerged_audio_edit = getattr(
+                request, "audio_edit", DISABLED_AUDIO_EDIT,
+            )
+        return result
 
     def _after_produce(self, prep: dict, result_latent, is_fresh: bool) -> None:
         self.last_denoise = prep["denoise"]