diff --git a/acestep/engine/stream.py b/acestep/engine/stream.py index b03c8aa1..bc2dbee4 100644 --- a/acestep/engine/stream.py +++ b/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 diff --git a/acestep/nodes/diffusion_nodes.py b/acestep/nodes/diffusion_nodes.py index 9b8f9192..f5e1c11d 100644 --- a/acestep/nodes/diffusion_nodes.py +++ b/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, ) diff --git a/acestep/streaming/ace_backend.py b/acestep/streaming/ace_backend.py index 75225010..fd273487 100644 --- a/acestep/streaming/ace_backend.py +++ b/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"] diff --git a/acestep/streaming/audio_edit.py b/acestep/streaming/audio_edit.py new file mode 100644 index 00000000..d54e65ef --- /dev/null +++ b/acestep/streaming/audio_edit.py @@ -0,0 +1,288 @@ +"""Backend-neutral live audio-edit state and mask/composite helpers. + +Edits are controls on the ordinary streaming pipeline: each generation +request snapshots one immutable :class:`LiveAudioEdit`, carries it through the +ring buffer, and uses that same snapshot when its latent is window-decoded. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Iterable + +import numpy as np + + +SOURCE_MODES = ("waveform", "structure") +EDIT_CROSSFADE_S = 0.025 +_EPS = 1e-6 + + +class AudioEditError(ValueError): + """The requested live edit cannot be represented exactly.""" + + +@dataclass(frozen=True) +class EditRegion: + start_s: float + end_s: float + + def to_wire(self) -> dict: + return {"start_s": self.start_s, "end_s": self.end_s} + + +@dataclass(frozen=True) +class LiveAudioEdit: + enabled: bool = False + regions: tuple[EditRegion, ...] = () + source_mode: str = "waveform" + strength: float = 1.0 + + def to_wire(self) -> dict: + return { + "enabled": self.enabled, + "regions": [r.to_wire() for r in self.regions], + "source_mode": self.source_mode, + "strength": self.strength, + } + + +DISABLED_AUDIO_EDIT = LiveAudioEdit() + + +def constrain_audio_edit( + current: LiveAudioEdit | None, + emerged: LiveAudioEdit | None, +) -> LiveAudioEdit: + """Limit one decoded request to the currently armed waveform mask. + + A request can finish after the user has changed or cleared the selected + regions. Only spans present in both snapshots are allowed to reach the + playback buffer. In particular, an enabled current edit with no regions + is an all-preserve mask even when the emerged request predates Edit mode. + """ + if current is None or not current.enabled or current.source_mode != "waveform": + return emerged or DISABLED_AUDIO_EDIT + if emerged is None or not emerged.enabled or emerged.source_mode != "waveform": + return LiveAudioEdit(True, (), "waveform", current.strength) + + intersections: list[EditRegion] = [] + i = j = 0 + current_regions = current.regions + emerged_regions = emerged.regions + while i < len(current_regions) and j < len(emerged_regions): + left = current_regions[i] + right = emerged_regions[j] + start = max(left.start_s, right.start_s) + end = min(left.end_s, right.end_s) + if end > start + _EPS: + intersections.append(EditRegion(start, end)) + if left.end_s < right.end_s: + i += 1 + else: + j += 1 + return LiveAudioEdit( + True, + tuple(intersections), + "waveform", + min(current.strength, emerged.strength), + ) + + +def parse_live_audio_edit( + regions: Iterable[dict] | None, + *, + enabled: bool, + source_mode: str, + strength: float, + canvas_duration_s: float, + source_duration_s: float, + left_extension_s: float = 0.0, + right_extension_s: float = 0.0, +) -> LiveAudioEdit: + """Validate the wire vocabulary and return an immutable tick snapshot.""" + if not enabled: + return DISABLED_AUDIO_EDIT + if source_mode not in SOURCE_MODES: + raise AudioEditError( + f"source_mode must be one of {SOURCE_MODES}, got {source_mode!r}" + ) + if not math.isfinite(strength) or not 0.0 <= strength <= 1.0: + raise AudioEditError(f"strength must be finite and in [0, 1], got {strength!r}") + if not math.isfinite(canvas_duration_s) or canvas_duration_s <= 0: + raise AudioEditError("the session has no finite editable canvas") + if ( + not math.isfinite(left_extension_s) + or not math.isfinite(right_extension_s) + or left_extension_s < 0.0 + or right_extension_s < 0.0 + or left_extension_s + right_extension_s > canvas_duration_s + _EPS + ): + raise AudioEditError("explicit extension spans are invalid for this canvas") + + parsed: list[EditRegion] = [] + for i, raw in enumerate(regions or ()): + if not isinstance(raw, dict): + raise AudioEditError(f"regions[{i}] must be an object") + try: + start = float(raw["start_s"]) + end = float(raw["end_s"]) + except (KeyError, TypeError, ValueError) as exc: + raise AudioEditError( + f"regions[{i}] requires numeric start_s and end_s" + ) from exc + if not math.isfinite(start) or not math.isfinite(end): + raise AudioEditError(f"regions[{i}] has non-finite bounds") + if start < -_EPS or end <= start + _EPS: + raise AudioEditError(f"regions[{i}] is negative, empty, or inverted") + if end > canvas_duration_s + _EPS: + raise AudioEditError( + f"regions[{i}].end_s={end} exceeds canvas {canvas_duration_s}s" + ) + if parsed and start < parsed[-1].end_s - _EPS: + raise AudioEditError("regions must be ordered and non-overlapping") + parsed.append(EditRegion(max(0.0, start), min(canvas_duration_s, end))) + + if source_mode == "structure": + if len(parsed) != 1 or parsed[0].start_s > _EPS or parsed[0].end_s < canvas_duration_s - _EPS: + raise AudioEditError("structure mode requires one region covering the full canvas") + else: + # The model canvas is routinely a little (and for SA3 sometimes much) + # longer than the uploaded content because of latent/profile padding. + # That is preserved silence, not a user-requested extension. Require + # coverage only for explicit extension spans supplied by the serving + # layer; ordinary canvas/source geometry differences remain editable + # with any valid interior region. + def _covered(required_start: float, required_end: float) -> bool: + if required_end <= required_start + _EPS: + return True + covered_until = required_start + for region in parsed: + if region.end_s <= covered_until + _EPS: + continue + if region.start_s > covered_until + _EPS: + return False + covered_until = max(covered_until, region.end_s) + if covered_until >= required_end - _EPS: + return True + return False + + # Anchor the right span to the end of uploaded content, not the + # backend canvas end: latent/model padding may continue beyond the + # exact requested duration and is not part of the extension. + right_extension_end_s = min( + canvas_duration_s, source_duration_s + right_extension_s, + ) + if not _covered(0.0, left_extension_s) or not _covered( + source_duration_s, right_extension_end_s, + ): + raise AudioEditError( + "regenerate regions must cover every explicit waveform extension tail" + ) + return LiveAudioEdit(True, tuple(parsed), source_mode, float(strength)) + + +def regenerate_mask( + edit: LiveAudioEdit, + *, + total_frames: int, + rate_hz: float, + offset_s: float = 0.0, + device=None, + dtype=None, +): + """ACE-polarity mask: 1 regenerates, 0 preserves.""" + import torch + + mask = torch.zeros((1, total_frames, 1), device=device, dtype=dtype) + if not edit.enabled or edit.source_mode != "waveform": + return mask + for region in edit.regions: + start = max(0, int(math.floor((region.start_s - offset_s) * rate_hz + _EPS))) + end = min(total_frames, int(math.ceil((region.end_s - offset_s) * rate_hz - _EPS))) + if end > start: + mask[:, start:end, :] = edit.strength + return mask + + +def sa3_inpaint_bundle(base: dict, source_btc, edit: LiveAudioEdit) -> dict: + """Return an SA3 conditioning bundle with its live binary inpaint input.""" + import torch + + if not edit.enabled or edit.source_mode != "waveform": + return base + frames = source_btc.shape[1] + regenerate = regenerate_mask( + edit, + total_frames=frames, + rate_hz=44100.0 / 4096.0, + device=source_btc.device, + dtype=source_btc.dtype, + ).movedim(1, 2) + preserve = 1.0 - regenerate.clamp(0, 1) + local_add = torch.cat( + [preserve, source_btc.movedim(1, 2) * preserve], dim=1, + ) + old = base.get("local_add_cond") + if old is not None: + local_add = local_add.to(device=old.device, dtype=old.dtype) + out = dict(base) + out["local_add_cond"] = local_add + return out + + +def composite_window( + generated: np.ndarray, + *, + start_sample: int, + source, + edit: LiveAudioEdit | None, + sample_rate: int = 48000, + crossfade_s: float = EDIT_CROSSFADE_S, +) -> np.ndarray: + """Restore preserved source samples in one absolute decoded window.""" + if edit is None or not edit.enabled or edit.source_mode != "waveform": + return generated + pcm = np.asarray(generated, dtype=np.float32) + if pcm.ndim != 2 or pcm.shape[0] == 0: + return pcm + if hasattr(source, "detach"): + src = source.detach().cpu().float().numpy() + else: + src = np.asarray(source, dtype=np.float32) + if src.ndim != 2: + return pcm + # Source providers use [C,N]; playback windows use [N,C]. + if src.shape[0] <= 8 and src.shape[1] > src.shape[0]: + src = src.T + if src.shape[1] == 1 and pcm.shape[1] > 1: + src = np.repeat(src, pcm.shape[1], axis=1) + elif src.shape[1] != pcm.shape[1]: + src = src[:, :pcm.shape[1]] + + absolute = start_sample + np.arange(pcm.shape[0]) + generated_weight = np.zeros(pcm.shape[0], dtype=np.float32) + fade = max(0, int(round(crossfade_s * sample_rate))) + for region in edit.regions: + r0 = int(math.floor(region.start_s * sample_rate + 0.5)) + r1 = int(math.floor(region.end_s * sample_rate + 0.5)) + inside = (absolute >= r0) & (absolute < r1) + generated_weight[inside] = 1.0 + width = min(fade, max((r1 - r0) // 2, 0)) + if width: + left = (absolute >= r0) & (absolute < r0 + width) + generated_weight[left] = np.minimum( + generated_weight[left], (absolute[left] - r0) / max(1, width - 1), + ) + right = (absolute >= r1 - width) & (absolute < r1) + generated_weight[right] = np.minimum( + generated_weight[right], (r1 - 1 - absolute[right]) / max(1, width - 1), + ) + + valid = (absolute >= 0) & (absolute < src.shape[0]) + out = pcm.copy() + if np.any(valid): + w = generated_weight[valid, None] + out[valid] = pcm[valid] * w + src[absolute[valid]] * (1.0 - w) + return out diff --git a/acestep/streaming/config.py b/acestep/streaming/config.py index f2eaf59d..a12600d1 100644 --- a/acestep/streaming/config.py +++ b/acestep/streaming/config.py @@ -65,6 +65,9 @@ class SessionConfig: # create-time, never hot-swapped. When absent on the wire, the # server's resolved --checkpoint family is the default. backend: str = "acestep" + # Optional fixed edit canvas. A value longer than the uploaded source + # preallocates a zero-padded tail so extension remains a live edit. + audio_edit_duration_s: float | None = None # --- sa3_* family fields (flat + prefixed per plan §3.5) --- # Fixed generation duration for sa3 sessions, seconds. None derives # it from the uploaded source audio length (the audio-to-audio diff --git a/acestep/streaming/encode.py b/acestep/streaming/encode.py index 4b20c6be..6bad5a14 100644 --- a/acestep/streaming/encode.py +++ b/acestep/streaming/encode.py @@ -17,6 +17,7 @@ def encode_cond_pair( duration, key, time_signature, + instruction=None, ): # WYSIWYG: the encoder sees exactly the text the UI sent. LoRA # trigger words land in `tags` via the client's visible-prepend @@ -27,7 +28,7 @@ def encode_cond_pair( cs = session.encode_text( tags=tags, lyrics="[Instrumental]", - instruction=TASK_INSTRUCTIONS["cover"], + instruction=instruction or TASK_INSTRUCTIONS["cover"], refer_latent=None, bpm=bpm, duration=duration, key=key, time_signature=time_signature, @@ -35,7 +36,7 @@ def encode_cond_pair( cf = session.encode_text( tags=tags, lyrics="[Instrumental]", - instruction=TASK_INSTRUCTIONS["cover"], + instruction=instruction or TASK_INSTRUCTIONS["cover"], refer_latent=refer_latent, bpm=bpm, duration=duration, key=key, time_signature=time_signature, diff --git a/acestep/streaming/events.py b/acestep/streaming/events.py index cf4bfaa1..1c04cd06 100644 --- a/acestep/streaming/events.py +++ b/acestep/streaming/events.py @@ -56,6 +56,8 @@ "ParamsUpdate", "ParamsEcho", "PromptApplied", + "AudioEditApplied", + "AudioEditFailed", "PromptBlendEcho", "LoraCatalogUpdate", "DepthApplied", @@ -154,6 +156,19 @@ class PromptApplied: tags: str +@dataclass(frozen=True) +class AudioEditApplied: + enabled: bool + regions: list + source_mode: str + strength: float + + +@dataclass(frozen=True) +class AudioEditFailed: + error: str + + @dataclass(frozen=True) class PromptBlendEcho: """Echo of a prompt-blend slider target from a non-primary origin. diff --git a/acestep/streaming/families.py b/acestep/streaming/families.py index 845c38ff..1270fb0f 100644 --- a/acestep/streaming/families.py +++ b/acestep/streaming/families.py @@ -43,6 +43,7 @@ def _make_acestep(ss): walk_window_s=ss.walk_window_s, neg_conditioning=ss.cond_negative, steering=steering, + source_waveform=(lambda: ss.canvas.wf), ) @@ -74,6 +75,7 @@ def _make_sa3(ss): # payload predating the blend surface stays blend-neutral. cond_b=init.get("cond_b"), source_latent_bct=init["source_latent_bct"], + source_waveform=init.get("source_waveform"), # Resolved accel values (compile already normalized to eager by # the create path); .get so an in-process payload without them # stays on the eager default. diff --git a/acestep/streaming/generator_backend.py b/acestep/streaming/generator_backend.py index 112091a1..8d1db727 100644 --- a/acestep/streaming/generator_backend.py +++ b/acestep/streaming/generator_backend.py @@ -158,6 +158,12 @@ class Capabilities: # True only when the backend has a steering controller with a # reachable vector bundle for its checkpoint. steering: bool = False + # Live audio editing through the ordinary ring-buffer pipeline. + # audio_edit_extend means a preallocated canvas may outlive its source; + # audio_edit_strength means waveform repaint accepts fractional masks. + audio_edit: bool = False + audio_edit_extend: bool = False + audio_edit_strength: bool = False @dataclass(frozen=True) diff --git a/acestep/streaming/sa3_backend.py b/acestep/streaming/sa3_backend.py index b3edba56..a21691a4 100644 --- a/acestep/streaming/sa3_backend.py +++ b/acestep/streaming/sa3_backend.py @@ -63,6 +63,12 @@ from acestep.engine.obs import logger from acestep.nodes.interpolation import INTERPOLATIONS, slerp from acestep.streaming.diffusion_backend import DiffusionBackend +from acestep.streaming.audio_edit import ( + DISABLED_AUDIO_EDIT, + composite_window, + constrain_audio_edit, + sa3_inpaint_bundle, +) from acestep.streaming.generator_backend import ( AudioChunk, AudioGeometry, @@ -174,6 +180,7 @@ def __init__( # SA3Context); None on directly-constructed test backends, where # handle_swap_source fails loudly instead. source_encoder: Optional[Callable] = None, + source_waveform: Optional[torch.Tensor] = None, ): super().__init__(adapter=adapter, codec=codec) self._cond = cond @@ -196,6 +203,8 @@ def __init__( self._prompt_rebuilder = prompt_rebuilder # Source re-encode hook for handle_swap_source (see ctor arg). self._source_encoder = source_encoder + self._source_waveform = source_waveform + self._audio_edit = DISABLED_AUDIO_EDIT self.knob_state = knob_state self.state = state self._steps = int(steps) @@ -240,6 +249,8 @@ def __init__( # GIL-atomic reference-swap argument doesn't cover it. The runner # only ever holds this to snapshot, never across pipeline work. self._control_lock = threading.Lock() + self._edit_bundle_cache: dict = {} + self._edit_bundle_history: list[tuple[dict, dict]] = [] # Rendered-audio cache: one full decode+resample per fresh # latent (SAME-S decodes the whole window in ~11 ms); window @@ -273,6 +284,7 @@ def from_context( prompt_b: Optional[str] = None, cond_b=None, source_latent_bct=None, + source_waveform=None, dit_backend: str = "eager", codec_backend: str = "eager", **kwargs, @@ -367,6 +379,7 @@ def _source_encoder(waveform, sample_rate, sample_size): prompt_rebuilder=_prompt_rebuilder, prompt_tags=prompt, source_encoder=_source_encoder, + source_waveform=source_waveform, **kwargs, ) @@ -398,7 +411,34 @@ def capabilities(self) -> Capabilities: # swap: backend-owned in-place re-anchor (handle_swap_source) — the # session's _apply_swap_if_pending dispatches there instead of the # ACE prepare_source body, so duration/conditioning stay fixed. - return Capabilities(refines_audio=True, loop_band=True, swap=True) + return Capabilities( + refines_audio=True, + loop_band=True, + swap=True, + audio_edit=True, + audio_edit_extend=True, + audio_edit_strength=False, + ) + + def handle_set_audio_edit(self, edit) -> None: + with self._control_lock: + self._audio_edit = edit + self._edit_bundle_cache.clear() + + 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 + with self._control_lock: + current = self._audio_edit + emerged = getattr(self._emerged_request, "audio_edit", DISABLED_AUDIO_EDIT) + return composite_window( + pcm, + start_sample=start_sample, + source=self._source_waveform, + edit=constrain_audio_edit(current, emerged), + sample_rate=DELIVERY_SAMPLE_RATE, + ) def geometry(self) -> AudioGeometry: return AudioGeometry( @@ -478,6 +518,7 @@ def handle_set_prompt(self, tags: str, *, tags_b: Optional[str] = None) -> None: self._cond = cond self._cond_b = cond_b self._active_bundle = self._blend_bundles(self._blend) + self._edit_bundle_cache.clear() # Emerged-generation labeling (see __init__): the new bundle gets # the next cond epoch; keep a short identity history so latents # still in flight on the OLD bundle stay attributable. @@ -518,13 +559,27 @@ def handle_swap_source(self, waveform, sample_rate) -> None: latent_bct = self._source_encoder(waveform, sample_rate, sample_size) encode_ms = (time.perf_counter() - t0) * 1000 latent_btc = latent_bct.movedim(1, 2).contiguous() + preserve_waveform = waveform.detach().cpu().float() + preserve_samples = int(round(self.playable_duration_s() * DELIVERY_SAMPLE_RATE)) + if preserve_waveform.shape[-1] < preserve_samples: + preserve_waveform = torch.nn.functional.pad( + preserve_waveform, + (0, preserve_samples - preserve_waveform.shape[-1]), + ) + else: + preserve_waveform = preserve_waveform[:, :preserve_samples] # Publish atomically w.r.t. the command thread's conditioning # swaps (same lock discipline as handle_set_prompt); the runner # reads the anchor on its own thread, which is also the thread # calling this hook. with self._control_lock: self._source_latent_btc = latent_btc + # Match the fixed client canvas just like the create path. Ordinary + # SA3 geometry padding is preserved silence unless it belongs to an + # explicitly selected extension region. + self._source_waveform = preserve_waveform self._latent_history.clear() + self._edit_bundle_cache.clear() logger.info( "sa3_source_swapped samples={} sample_rate={} encode_ms={:.1f}", int(waveform.shape[-1]), int(sample_rate), encode_ms, @@ -546,6 +601,7 @@ def handle_set_prompt_blend(self, value: float) -> None: with self._control_lock: self._blend = v self._active_bundle = self._blend_bundles(v) + self._edit_bundle_cache.clear() def _blend_bundles(self, v: float) -> dict: """The active cond bundle for blend value ``v``: A verbatim at @@ -588,6 +644,8 @@ def _blend_bundles(self, v: float) -> dict: # ---- produce hooks --------------------------------------------------------- def _prepare_tick(self, knobs: dict, ctx: TickContext) -> dict: + with self._control_lock: + audio_edit = self._audio_edit # Schedule warp: hot-applied, but cache-coupled — the pipeline # caches schedules per denoise value, so a changed alpha must # invalidate or already-seen denoise values keep the old warp. @@ -610,8 +668,11 @@ def _prepare_tick(self, knobs: dict, ctx: TickContext) -> dict: fb_depth_raw = float(knobs.get("feedback_depth", 1.0)) except (TypeError, ValueError): fb_depth_raw = 1.0 + denoise = float(knobs.get("sa3_denoise", 1.0)) + if audio_edit.enabled and audio_edit.source_mode == "structure": + denoise = audio_edit.strength return { - "denoise": float(knobs.get("sa3_denoise", 1.0)), + "denoise": denoise, "seed": int(knobs.get("seed", self._default_seed)), "steps": int(knobs.get("steps_override", self._steps)), "shift": shift, @@ -620,6 +681,7 @@ def _prepare_tick(self, knobs: dict, ctx: TickContext) -> dict: "feedback_depth": max( 1, min(MAX_FEEDBACK_DEPTH, int(round(fb_depth_raw))), ), + "audio_edit": audio_edit, } def _generate(self, prep: dict): @@ -663,6 +725,22 @@ def _generate(self, prep: dict): with self._control_lock: aux_cond = self._active_bundle latent_frames = self._cond.latent_frames + edit = prep["audio_edit"] + if ( + edit.enabled + and edit.source_mode == "waveform" + and self._source_latent_btc is not None + ): + cache_key = (id(aux_cond), edit) + edited = self._edit_bundle_cache.get(cache_key) + if edited is None: + edited = sa3_inpaint_bundle( + aux_cond, self._source_latent_btc, edit, + ) + self._edit_bundle_cache[cache_key] = edited + self._edit_bundle_history.append((edited, aux_cond)) + del self._edit_bundle_history[:-16] + aux_cond = edited self.pipeline.submit(SlotRequest( seed=prep["seed"], @@ -680,6 +758,7 @@ def _generate(self, prep: dict): x0_target_strength=prep["x0_target"], aux_cond=aux_cond, latent_frames=latent_frames, + audio_edit=edit, # Deterministic pingpong: identical requests must replay the # same trajectory or advancing windows splice different # realizations (incoherent audio). See SlotRequest. @@ -700,6 +779,11 @@ def _cond_meta_for(self, bundle) -> tuple: # non-atomic append + truncate on this list. with self._control_lock: history = tuple(self._cond_history) + edit_history = tuple(self._edit_bundle_history) + for edited, parent in edit_history: + if bundle is edited: + bundle = parent + break for b, epoch, tags in history: if b is bundle: return epoch, tags @@ -780,12 +864,14 @@ def render_window(self, t_start_s: float): if decode_src is None: return None if self._windowed_codec: - return self._render_window_via_codec(decode_src, t_start_s) - audio = self._rendered_audio(decode_src) - n = int(round(self.vae_window * DELIVERY_SAMPLE_RATE)) - start = int(round(t_start_s * DELIVERY_SAMPLE_RATE)) - start = max(0, min(start, max(0, audio.shape[0] - n))) - return AudioChunk(pcm=audio[start:start + n], start_sample=start) + chunk = self._render_window_via_codec(decode_src, t_start_s) + else: + audio = self._rendered_audio(decode_src) + n = int(round(self.vae_window * DELIVERY_SAMPLE_RATE)) + start = int(round(t_start_s * DELIVERY_SAMPLE_RATE)) + start = max(0, min(start, max(0, audio.shape[0] - n))) + chunk = AudioChunk(pcm=audio[start:start + n], start_sample=start) + return chunk def _render_window_via_codec(self, latent_btc: torch.Tensor, t_start_s: float): """Windowed-codec render (SAME-L / medium): decode ONLY a small @@ -833,9 +919,8 @@ def _render_window_via_codec(self, latent_btc: torch.Tensor, t_start_s: float): def render_full(self): if self._current_result is None: return None - return AudioChunk( - pcm=self._rendered_audio(self._current_result), start_sample=0, - ) + pcm = self._rendered_audio(self._current_result) + return AudioChunk(pcm=pcm, start_sample=0) # ---- bookkeeping ------------------------------------------------------------- diff --git a/acestep/streaming/sa3_session.py b/acestep/streaming/sa3_session.py index 55473de4..d822306e 100644 --- a/acestep/streaming/sa3_session.py +++ b/acestep/streaming/sa3_session.py @@ -139,7 +139,9 @@ def create_sa3_session( if waveform.shape[0] == 1: waveform = waveform.repeat(2, 1) source_duration_s = waveform.shape[-1] / SAMPLE_RATE - duration_s = float(config.sa3_duration_s or 0.0) or source_duration_s + duration_s = float( + config.audio_edit_duration_s or config.sa3_duration_s or 0.0 + ) or source_duration_s duration_s = min(duration_s, SA3_MAX_DURATION_S) # Land on the TRT DiT fast path when engines are built (medium): # a duration whose padded latent window exceeds every engine @@ -206,6 +208,7 @@ def create_sa3_session( prompt_text=prompt, prompt_text_b=prompt_b, current_depth=depth, + source_content_duration_s=source_duration_s, ) # Same transactional create shape as StreamingSession.create's ACE @@ -262,6 +265,11 @@ def create_sa3_session( "cond": cond, "cond_b": cond_b, "source_latent_bct": source_latent, + # Composite against the exact initial client buffer, including + # ordinary model/canvas padding. Without the zero-padded tail, + # decoded samples outside an edit mask but past the short upload + # were incorrectly left generated instead of preserved silence. + "source_waveform": torch.from_numpy(src_np.T.copy()), "duration_s": duration_s, "dit_backend": dit_backend, "codec_backend": codec_backend, diff --git a/acestep/streaming/session.py b/acestep/streaming/session.py index 1b10ed06..a4478977 100644 --- a/acestep/streaming/session.py +++ b/acestep/streaming/session.py @@ -72,12 +72,19 @@ from acestep.engine.canvas import SourceCanvas from acestep.streaming.audio_engine import AudioEngine +from acestep.streaming.audio_edit import ( + AudioEditError, + DISABLED_AUDIO_EDIT, + parse_live_audio_edit, +) from acestep.streaming.commands import CommandOrigin from acestep.streaming.config import SessionConfig from acestep.streaming.encode import blend_for_strength, encode_cond_pair from acestep.steering import CapacityError, EmptyError from acestep.streaming.events import ( AudioReady, + AudioEditApplied, + AudioEditFailed, AudioWriteFailed, AudioWritten, CommandFailed, @@ -492,6 +499,11 @@ def __init__( self.stream = stream self.state = state self.audio_eng = audio_eng + # Serializes edit-mask activation with the last compositor + event + # publication boundary. A runner callback that began under the old + # mask must publish before the command's full-source restoration, never + # race past it and overwrite the restored client mirror afterwards. + self._audio_edit_delivery_lock = threading.RLock() # Staleness estimator for client playhead reports (params channel # ``client_time`` stamps). Only ever touched from set_knobs. self._report_staleness = ReportStalenessEstimator() @@ -699,6 +711,10 @@ def snapshot(self) -> dict: "timbre_name": state.timbre_name, "timbre_strength": state.timbre_strength, "structure_name": state.struct_name, + "audio_edit": ( + state.audio_edit.to_wire() + if state.audio_edit is not None else DISABLED_AUDIO_EDIT.to_wire() + ), "interp_prompt": state.interp_prompt, "interp_timbre": state.interp_timbre, "interp_structure": state.interp_structure, @@ -853,27 +869,37 @@ def _on_audio_ready(self, wav_np, win_start=None, win_end=None): The audio array passed in the event is the same numpy array the runner produced; subscribers must treat it as immutable. """ - state = self.state - if win_start is not None: - ss = int(win_start) - se = ss + len(wav_np) - else: - self.audio_eng.swap(wav_np) - ss = 0 - se = len(wav_np) - - params_snapshot = dict(state.params) - self.bus.publish(AudioReady( - audio=wav_np, - start_sample=ss, - num_samples=int(se - ss), - channels=state.n_channels, - tick_ms=float(params_snapshot.get("tick_ms", 0) or 0), - dec_ms=float(params_snapshot.get("dec_ms", 0) or 0), - num_gens=int(params_snapshot.get("num_gens", 0) or 0), - params=params_snapshot, - published_wall_s=time.monotonic(), - )) + with self._audio_edit_delivery_lock: + state = self.state + finalize = getattr(self.backend, "finalize_audio_edit_window", None) + if win_start is not None: + ss = int(win_start) + if finalize is not None: + wav_np = finalize(wav_np, ss) + # The runner writes before calling us so it can trim emitted + # windows from its live buffer. Reassert the final edit mask + # after its edge crossfade, immediately before publication. + self.audio_eng.patch_window(wav_np, ss) + se = ss + len(wav_np) + else: + if finalize is not None: + wav_np = finalize(wav_np, 0) + self.audio_eng.swap(wav_np) + ss = 0 + se = len(wav_np) + + params_snapshot = dict(state.params) + self.bus.publish(AudioReady( + audio=wav_np, + start_sample=ss, + num_samples=int(se - ss), + channels=state.n_channels, + tick_ms=float(params_snapshot.get("tick_ms", 0) or 0), + dec_ms=float(params_snapshot.get("dec_ms", 0) or 0), + num_gens=int(params_snapshot.get("num_gens", 0) or 0), + params=params_snapshot, + published_wall_s=time.monotonic(), + )) # ---- Pending drain (runs inside before_tick) ----------------------- @@ -1006,6 +1032,16 @@ def _apply_swap_if_pending(self) -> None: # take advantage of every built engine profile, not a # stale 60 s default. new_wf = new_wf[:, :int(self.max_seconds * SAMPLE_RATE)] + new_source_content_duration_s = new_wf.shape[-1] / SAMPLE_RATE + edit_duration_s = float(self.config.audio_edit_duration_s or 0.0) + if edit_duration_s > new_source_content_duration_s: + target_samples = min( + int(round(edit_duration_s * SAMPLE_RATE)), + int(self.max_seconds * SAMPLE_RATE), + ) + new_wf = torch.nn.functional.pad( + new_wf, (0, max(0, target_samples - new_wf.shape[-1])), + ) rem = new_wf.shape[-1] % self.pool if rem: new_wf = new_wf[:, :new_wf.shape[-1] - rem] @@ -1131,11 +1167,15 @@ def _apply_swap_if_pending(self) -> None: # epoch-checks at commit and discards itself. self.canvas = new_canvas state.source_epoch += 1 + state.source_content_duration_s = new_source_content_duration_s + state.audio_edit = DISABLED_AUDIO_EDIT + self.backend.handle_set_audio_edit(DISABLED_AUDIO_EDIT) tl = state.timbre_latent refer = tl if tl is not None else new_source.latent state.cond_pair = encode_cond_pair( self.session, tags, refer, new_bpm, new_audio_duration_s, new_key, new_time_sig, + instruction=self._audio_edit_instruction(), ) # Carry promptB across the swap so the blend slider keeps # its meaning. If B was identical to A pre-swap, keep it @@ -1144,6 +1184,7 @@ def _apply_swap_if_pending(self) -> None: state.cond_pair_b = encode_cond_pair( self.session, state.prompt_text_b, refer, new_bpm, new_audio_duration_s, new_key, new_time_sig, + instruction=self._audio_edit_instruction(), ) else: state.cond_pair_b = state.cond_pair @@ -1265,6 +1306,9 @@ def _apply_swap_backend_owned(self, handle, new_wf, fixture_name) -> None: else: src_np = src_np[:n_play] state.n_channels = int(src_np.shape[1]) + state.source_content_duration_s = wf.shape[-1] / SAMPLE_RATE + state.audio_edit = DISABLED_AUDIO_EDIT + self.backend.handle_set_audio_edit(DISABLED_AUDIO_EDIT) # Retire anything staged against the old source (ACE # parity; inert for backends without write_audio). state.source_epoch += 1 @@ -1457,6 +1501,7 @@ def _apply_timbre_waveform(self, t_wf: torch.Tensor, name: str) -> float: self.session, state.prompt_text, timbre_latent, state.bpm, state.duration, state.key, state.time_signature, + instruction=self._audio_edit_instruction(), ) # Re-encode B against the new timbre too. if state.prompt_text_b != state.prompt_text: @@ -1464,6 +1509,7 @@ def _apply_timbre_waveform(self, t_wf: torch.Tensor, name: str) -> float: self.session, state.prompt_text_b, timbre_latent, state.bpm, state.duration, state.key, state.time_signature, + instruction=self._audio_edit_instruction(), ) else: state.cond_pair_b = state.cond_pair @@ -1660,6 +1706,110 @@ def set_loop_band( except (TypeError, ValueError): self.audio_eng.loop_band = None + def _audio_edit_instruction(self, edit=None) -> str: + active = edit if edit is not None else self.state.audio_edit + if active is not None and active.enabled and active.source_mode == "waveform": + return TASK_INSTRUCTIONS["repaint"] + return TASK_INSTRUCTIONS["cover"] + + def _reencode_ace_edit_instruction(self, edit) -> None: + """Retarget ACE's task prefix while preserving the visible prompt.""" + state = self.state + refer = self._active_refer_latent() + instruction = self._audio_edit_instruction(edit) + cond_pair = encode_cond_pair( + self.session, state.prompt_text, refer, + state.bpm, state.duration, state.key, state.time_signature, + instruction=instruction, + ) + cond_pair_b = ( + cond_pair if state.prompt_text_b == state.prompt_text + else encode_cond_pair( + self.session, state.prompt_text_b, refer, + state.bpm, state.duration, state.key, state.time_signature, + instruction=instruction, + ) + ) + state.cond_pair = cond_pair + state.cond_pair_b = cond_pair_b + self._refresh_conditioning() + + @requires_capability("audio_edit", "set_audio_edit") + def set_audio_edit( + self, + regions: list | None, + *, + enabled: bool = True, + source_mode: str = "waveform", + strength: float = 1.0, + origin: CommandOrigin = CommandOrigin.PRIMARY, + ) -> None: + """Change the live repaint/extend/cover control without restarting. + + The backend snapshots this immutable value into each new ring slot; + already in-flight slots finish under their previous edit. + """ + self.state.last_activity_ts = time.monotonic() + try: + edit = parse_live_audio_edit( + regions, + enabled=bool(enabled), + source_mode=str(source_mode), + strength=float(strength), + canvas_duration_s=float(self.backend.geometry().duration_s), + source_duration_s=float( + self.state.source_content_duration_s or self.state.duration + ), + # audio_edit_duration_s represents a client-requested right + # extension. Backend/model padding beyond source content is + # ordinary geometry and must never trigger tail validation. + right_extension_s=max( + 0.0, + min( + float(self.config.audio_edit_duration_s or 0.0), + float(self.backend.geometry().duration_s), + ) + - float( + self.state.source_content_duration_s or self.state.duration + ), + ), + ) + caps = self.backend.capabilities() + if ( + edit.enabled + and edit.source_mode == "waveform" + and edit.strength < 1.0 - 1e-6 + and not caps.audio_edit_strength + ): + raise AudioEditError( + f"backend {self.backend.name!r} supports only binary waveform repaint strength" + ) + previous_instruction = self._audio_edit_instruction() + with self.state._lock: + if ( + self.backend.name == "acestep" + and previous_instruction != self._audio_edit_instruction(edit) + ): + self._reencode_ace_edit_instruction(edit) + self.state.audio_edit = edit + with self._audio_edit_delivery_lock: + self.backend.handle_set_audio_edit(edit) + # Immediately restore every currently preserved sample and + # publish that full buffer so the transport mirror cannot + # retain audio from an earlier unmasked generation while new + # masked slots drain. The re-entrant delivery lock makes this + # restoration indivisible from edit activation. + if edit.enabled and edit.source_mode == "waveform": + finalize = getattr( + self.backend, "finalize_audio_edit_window", None, + ) + if finalize is not None: + self._on_audio_ready(self.audio_eng.current.copy()) + except Exception as exc: + self.bus.publish(AudioEditFailed(error=str(exc))) + return + self.bus.publish(AudioEditApplied(**edit.to_wire())) + def set_prompt( self, tags: str, @@ -1710,12 +1860,14 @@ def set_prompt( state.cond_pair = encode_cond_pair( self.session, tags, refer, state.bpm, state.duration, key_used, state.time_signature, + instruction=self._audio_edit_instruction(), ) state.prompt_text = tags if tags_b and tags_b != tags: state.cond_pair_b = encode_cond_pair( self.session, tags_b, refer, state.bpm, state.duration, key_used, state.time_signature, + instruction=self._audio_edit_instruction(), ) state.prompt_text_b = tags_b else: @@ -1982,12 +2134,14 @@ def clear_timbre_source( self.session, state.prompt_text, refer, state.bpm, state.duration, state.key, state.time_signature, + instruction=self._audio_edit_instruction(), ) if state.prompt_text_b != state.prompt_text: state.cond_pair_b = encode_cond_pair( self.session, state.prompt_text_b, refer, state.bpm, state.duration, state.key, state.time_signature, + instruction=self._audio_edit_instruction(), ) else: state.cond_pair_b = state.cond_pair @@ -2208,9 +2362,11 @@ def write_audio( refer = self.stream.source.latent cp = encode_cond_pair( self.session, prompt_a, refer, bpm, dur, key, tsig, + instruction=self._audio_edit_instruction(), ) cp_b = cp if prompt_b == prompt_a else encode_cond_pair( self.session, prompt_b, refer, bpm, dur, key, tsig, + instruction=self._audio_edit_instruction(), ) with state._lock: if state.source_epoch == epoch: @@ -2299,6 +2455,16 @@ def create( max_seconds = max_profile_duration_s() waveform = waveform[:, :int(max_seconds * SAMPLE_RATE)] + source_content_duration_s = waveform.shape[-1] / SAMPLE_RATE + edit_duration_s = float(config.audio_edit_duration_s or 0.0) + if edit_duration_s > source_content_duration_s: + target_samples = min( + int(round(edit_duration_s * SAMPLE_RATE)), + int(max_seconds * SAMPLE_RATE), + ) + waveform = torch.nn.functional.pad( + waveform, (0, max(0, target_samples - waveform.shape[-1])), + ) rem = waveform.shape[-1] % _POOL if rem: waveform = waveform[:, :waveform.shape[-1] - rem] @@ -2599,6 +2765,7 @@ def create( prompt_text=prompt, prompt_text_b=prompt_b, current_depth=int(depth), + source_content_duration_s=source_content_duration_s, ) streaming = cls( diff --git a/acestep/streaming/state.py b/acestep/streaming/state.py index 788e55a5..a94dfe52 100644 --- a/acestep/streaming/state.py +++ b/acestep/streaming/state.py @@ -74,6 +74,8 @@ class SessionState: prompt_text: str # prompt A (read by runner each tick) prompt_text_b: str # prompt B (dispatcher only) current_depth: int # active pipeline_depth + source_content_duration_s: float | None = None + audio_edit: Any = None # === Slider values driving the engine === prompt_blend: float = 0.0 diff --git a/demos/realtime_motion_graph_web/protocol.py b/demos/realtime_motion_graph_web/protocol.py index 94a5ccfe..c17af21d 100644 --- a/demos/realtime_motion_graph_web/protocol.py +++ b/demos/realtime_motion_graph_web/protocol.py @@ -289,6 +289,20 @@ class EventSpec: ), description="Re-encode the live prompt (text-encoder pass).", ), + CommandSpec( + "set_audio_edit", + fields=( + FieldSpec("enabled", "bool", default=True), + FieldSpec("regions", "list", + description="Ordered {start_s,end_s} regenerate spans on the fixed session canvas."), + FieldSpec("source_mode", "enum", default="waveform", + options=("waveform", "structure")), + FieldSpec("strength", "float", default=1.0, + description="Regenerate strength in [0,1]; SA3 waveform repaint is binary (1 only)."), + ), + requires="audio_edit", + description="Live repaint/extend/cover control. New ring slots snapshot the edit; no restart or offline generation.", + ), CommandSpec( "set_prompt_blend", fields=(FieldSpec("value", "float", required=True, default=0.0, @@ -678,6 +692,22 @@ class EventSpec: description="The clamped applied depth."),), description="Ack for set_depth.", ), + EventSpec( + "audio_edit_applied", + fields=( + FieldSpec("enabled", "bool", required=True), + FieldSpec("regions", "list", required=True), + FieldSpec("source_mode", "enum", required=True, + options=("waveform", "structure")), + FieldSpec("strength", "float", required=True), + ), + description="Ack for a live edit state now feeding new ring-buffer requests.", + ), + EventSpec( + "audio_edit_failed", + fields=(FieldSpec("error", "str", required=True),), + description="The live edit request was invalid or unsupported by the active backend.", + ), EventSpec( "manual_slot_count", fields=(FieldSpec("count", "int", required=True, diff --git a/demos/realtime_motion_graph_web/ws_adapter.py b/demos/realtime_motion_graph_web/ws_adapter.py index 0ca0d725..ab14ba7b 100644 --- a/demos/realtime_motion_graph_web/ws_adapter.py +++ b/demos/realtime_motion_graph_web/ws_adapter.py @@ -62,6 +62,8 @@ from acestep.streaming.config import SessionConfig from acestep.streaming.events import ( AudioReady, + AudioEditApplied, + AudioEditFailed, AudioWriteFailed, AudioWritten, CommandFailed, @@ -1421,6 +1423,15 @@ def on_event(event) -> None: _send_json({"type": "prompt_blend_echo", "value": event.value}) elif isinstance(event, PromptApplied): _send_json({"type": "prompt_applied", "tags": event.tags}) + elif isinstance(event, AudioEditApplied): + _send_json({"type": "audio_edit_applied", **{ + "enabled": event.enabled, + "regions": event.regions, + "source_mode": event.source_mode, + "strength": event.strength, + }}) + elif isinstance(event, AudioEditFailed): + _send_json({"type": "audio_edit_failed", "error": event.error}) elif isinstance(event, LoraCatalogUpdate): _send_json({"type": "lora_catalog", "catalog": event.catalog}) elif isinstance(event, DepthApplied): @@ -1730,6 +1741,14 @@ def _recv_binary_payload(fail_type: str): time_signature=data.get("time_signature"), origin=origin, ) + elif mtype == "set_audio_edit": + streaming.set_audio_edit( + data.get("regions") or [], + enabled=bool(data.get("enabled", True)), + source_mode=str(data.get("source_mode", "waveform")), + strength=float(data.get("strength", 1.0)), + origin=origin, + ) elif mtype == "set_prompt_blend": try: v = float(data.get("value", 0.0)) diff --git a/integrations/ableton/release/DEMON for Live/demon-preset-silent.wav.asd b/integrations/ableton/release/DEMON for Live/demon-preset-silent.wav.asd new file mode 100644 index 00000000..c3728f1c Binary files /dev/null and b/integrations/ableton/release/DEMON for Live/demon-preset-silent.wav.asd differ diff --git a/packages/demon-client/dist/demon-client.js b/packages/demon-client/dist/demon-client.js index c2ec4dc8..c4cd9424 100644 --- a/packages/demon-client/dist/demon-client.js +++ b/packages/demon-client/dist/demon-client.js @@ -5,6 +5,7 @@ var COMMAND_NAMES = [ "params", "loop_band", "prompt", + "set_audio_edit", "set_prompt_blend", "set_interp_method", "set_depth", @@ -36,6 +37,8 @@ var EVENT_NAMES = [ "stem_assets", "stem_failed", "depth_applied", + "audio_edit_applied", + "audio_edit_failed", "manual_slot_count", "timbre_set", "timbre_cleared", @@ -1908,6 +1911,25 @@ var RemoteBackend = class extends EventTarget { } catch { } } + /** + * Change live repaint/extend/cover state. This does not launch a one-shot + * job: new requests enter the existing ring and emerge through ordinary + * windowed slices. Disable with `enabled: false` (regions may be empty). + */ + sendSetAudioEdit(regions, options = {}) { + if (this.ws?.readyState !== this._wsOpen) return; + try { + const msg = { + type: "set_audio_edit", + enabled: options.enabled ?? true, + regions, + source_mode: options.sourceMode ?? "waveform", + strength: Math.max(0, Math.min(1, options.strength ?? 1)) + }; + this.ws.send(JSON.stringify(msg)); + } catch { + } + } /** * Live prompt A/B blend knob. Backend keeps cached cond pairs for both * prompts (encoded by the most recent ``sendPrompt`` that carried a diff --git a/packages/demon-client/dist/demon-client.node.cjs b/packages/demon-client/dist/demon-client.node.cjs index 10b96981..3dcc2620 100644 --- a/packages/demon-client/dist/demon-client.node.cjs +++ b/packages/demon-client/dist/demon-client.node.cjs @@ -1368,6 +1368,25 @@ var RemoteBackend = class extends EventTarget { } catch { } } + /** + * Change live repaint/extend/cover state. This does not launch a one-shot + * job: new requests enter the existing ring and emerge through ordinary + * windowed slices. Disable with `enabled: false` (regions may be empty). + */ + sendSetAudioEdit(regions, options = {}) { + if (this.ws?.readyState !== this._wsOpen) return; + try { + const msg = { + type: "set_audio_edit", + enabled: options.enabled ?? true, + regions, + source_mode: options.sourceMode ?? "waveform", + strength: Math.max(0, Math.min(1, options.strength ?? 1)) + }; + this.ws.send(JSON.stringify(msg)); + } catch { + } + } /** * Live prompt A/B blend knob. Backend keeps cached cond pairs for both * prompts (encoded by the most recent ``sendPrompt`` that carried a @@ -1775,6 +1794,7 @@ var COMMAND_NAMES = [ "params", "loop_band", "prompt", + "set_audio_edit", "set_prompt_blend", "set_interp_method", "set_depth", @@ -1806,6 +1826,8 @@ var EVENT_NAMES = [ "stem_assets", "stem_failed", "depth_applied", + "audio_edit_applied", + "audio_edit_failed", "manual_slot_count", "timbre_set", "timbre_cleared", diff --git a/packages/demon-client/index.ts b/packages/demon-client/index.ts index 9eb1b6d4..f960f6ad 100644 --- a/packages/demon-client/index.ts +++ b/packages/demon-client/index.ts @@ -35,7 +35,12 @@ export * from "./controls"; // WebSocket session client (binary slice stream, swap/stem state // machines, typed senders). export { RemoteBackend, float16ArrayToFloat32 } from "./protocol"; -export type { RemoteBackendOptions, WsTrace, WsTracePhase } from "./protocol"; +export type { + AudioEditRegion, + RemoteBackendOptions, + WsTrace, + WsTracePhase, +} from "./protocol"; // Realtime audio playback (worklet + ScriptProcessor fallback, loudness // matcher, stem overlays). diff --git a/packages/demon-client/node.ts b/packages/demon-client/node.ts index d1183a60..26ff02d5 100644 --- a/packages/demon-client/node.ts +++ b/packages/demon-client/node.ts @@ -12,6 +12,7 @@ export { RemoteBackend, float16ArrayToFloat32 } from "./protocol"; export type { + AudioEditRegion, RemoteBackendOptions, WsTrace, WsTracePhase, diff --git a/packages/demon-client/protocol.ts b/packages/demon-client/protocol.ts index 8b594c36..28f2fca3 100644 --- a/packages/demon-client/protocol.ts +++ b/packages/demon-client/protocol.ts @@ -44,6 +44,7 @@ import type { ParamsCommand, PromptCommand, SetDepthCommand, + SetAudioEditCommand, SetInterpMethodCommand, SetPromptBlendCommand, SetStructureFixtureCommand, @@ -56,6 +57,11 @@ import type { WriteAudioCommand, } from "./types/wireContract.gen"; +export interface AudioEditRegion { + start_s: number; + end_s: number; +} + /** Optional behaviors the host app injects into RemoteBackend. */ export interface RemoteBackendOptions { /** Applied to `tags` and `tags_b` on every `sendPrompt` before they hit @@ -998,6 +1004,32 @@ export class RemoteBackend extends EventTarget { } catch {} } + /** + * Change live repaint/extend/cover state. This does not launch a one-shot + * job: new requests enter the existing ring and emerge through ordinary + * windowed slices. Disable with `enabled: false` (regions may be empty). + */ + sendSetAudioEdit( + regions: AudioEditRegion[], + options: { + enabled?: boolean; + sourceMode?: "waveform" | "structure"; + strength?: number; + } = {}, + ): void { + if (this.ws?.readyState !== this._wsOpen) return; + try { + const msg: SetAudioEditCommand = { + type: "set_audio_edit", + enabled: options.enabled ?? true, + regions, + source_mode: options.sourceMode ?? "waveform", + strength: Math.max(0, Math.min(1, options.strength ?? 1)), + }; + this.ws.send(JSON.stringify(msg)); + } catch {} + } + /** * Live prompt A/B blend knob. Backend keeps cached cond pairs for both * prompts (encoded by the most recent ``sendPrompt`` that carried a diff --git a/packages/demon-client/types/wireContract.gen.hpp b/packages/demon-client/types/wireContract.gen.hpp index e00fc272..cfc21e88 100644 --- a/packages/demon-client/types/wireContract.gen.hpp +++ b/packages/demon-client/types/wireContract.gen.hpp @@ -75,6 +75,21 @@ namespace command { inline constexpr const char* kTimeSignature = "time_signature"; } // namespace prompt + namespace set_audio_edit { + inline constexpr const char* kType = "set_audio_edit"; + inline constexpr const char* kEnabled = "enabled"; + /** Ordered {start_s,end_s} regenerate spans on the fixed session canvas. */ + inline constexpr const char* kRegions = "regions"; + inline constexpr const char* kSourceMode = "source_mode"; + /** Regenerate strength in [0,1]; SA3 waveform repaint is binary (1 only). */ + inline constexpr const char* kStrength = "strength"; + + namespace source_mode { + inline constexpr const char* kWaveform = "waveform"; + inline constexpr const char* kStructure = "structure"; + } // namespace source_mode + } // namespace set_audio_edit + namespace set_prompt_blend { inline constexpr const char* kType = "set_prompt_blend"; /** 0.0 = A, 1.0 = B. Clamped to [0,1]. */ @@ -344,6 +359,24 @@ namespace event { inline constexpr const char* kValue = "value"; } // namespace depth_applied + namespace audio_edit_applied { + inline constexpr const char* kType = "audio_edit_applied"; + inline constexpr const char* kEnabled = "enabled"; + inline constexpr const char* kRegions = "regions"; + inline constexpr const char* kSourceMode = "source_mode"; + inline constexpr const char* kStrength = "strength"; + + namespace source_mode { + inline constexpr const char* kWaveform = "waveform"; + inline constexpr const char* kStructure = "structure"; + } // namespace source_mode + } // namespace audio_edit_applied + + namespace audio_edit_failed { + inline constexpr const char* kType = "audio_edit_failed"; + inline constexpr const char* kError = "error"; + } // namespace audio_edit_failed + namespace manual_slot_count { inline constexpr const char* kType = "manual_slot_count"; /** The live manual steering slot count after the command. */ @@ -431,6 +464,7 @@ namespace config { inline constexpr const char* kLoraPaths = "lora_paths"; inline constexpr const char* kClientId = "client_id"; inline constexpr const char* kBackend = "backend"; + inline constexpr const char* kAudioEditDurationS = "audio_edit_duration_s"; inline constexpr const char* kSa3DurationS = "sa3_duration_s"; } // namespace config diff --git a/packages/demon-client/types/wireContract.gen.ts b/packages/demon-client/types/wireContract.gen.ts index 76e8d691..9bf63840 100644 --- a/packages/demon-client/types/wireContract.gen.ts +++ b/packages/demon-client/types/wireContract.gen.ts @@ -24,6 +24,7 @@ export type CommandName = | "params" | "loop_band" | "prompt" + | "set_audio_edit" | "set_prompt_blend" | "set_interp_method" | "set_depth" @@ -45,6 +46,7 @@ export const COMMAND_NAMES: readonly CommandName[] = [ "params", "loop_band", "prompt", + "set_audio_edit", "set_prompt_blend", "set_interp_method", "set_depth", @@ -77,6 +79,8 @@ export type EventName = | "stem_assets" | "stem_failed" | "depth_applied" + | "audio_edit_applied" + | "audio_edit_failed" | "manual_slot_count" | "timbre_set" | "timbre_cleared" @@ -102,6 +106,8 @@ export const EVENT_NAMES: readonly EventName[] = [ "stem_assets", "stem_failed", "depth_applied", + "audio_edit_applied", + "audio_edit_failed", "manual_slot_count", "timbre_set", "timbre_cleared", @@ -161,6 +167,16 @@ export interface PromptCommand { time_signature?: string; } +export interface SetAudioEditCommand { + type: "set_audio_edit"; + enabled?: boolean; + /** Ordered {start_s,end_s} regenerate spans on the fixed session canvas. */ + regions?: unknown[]; + source_mode?: "waveform" | "structure"; + /** Regenerate strength in [0,1]; SA3 waveform repaint is binary (1 only). */ + strength?: number; +} + export interface SetPromptBlendCommand { type: "set_prompt_blend"; /** 0.0 = A, 1.0 = B. Clamped to [0,1]. */ @@ -392,6 +408,19 @@ export interface DepthAppliedEvent { value: number; } +export interface AudioEditAppliedEvent { + type: "audio_edit_applied"; + enabled: boolean; + regions: unknown[]; + source_mode: "waveform" | "structure"; + strength: number; +} + +export interface AudioEditFailedEvent { + type: "audio_edit_failed"; + error: string; +} + export interface ManualSlotCountEvent { type: "manual_slot_count"; /** The live manual steering slot count after the command. */ @@ -477,6 +506,7 @@ export interface SessionConfigPayload { lora_paths?: unknown[]; client_id?: string | null; backend?: string; + audio_edit_duration_s?: number | null; sa3_duration_s?: number | null; // SessionConfig is permissive; extras pass through. [k: string]: unknown; @@ -518,6 +548,7 @@ export type WireCommand = | ParamsCommand | LoopBandCommand | PromptCommand + | SetAudioEditCommand | SetPromptBlendCommand | SetInterpMethodCommand | SetDepthCommand @@ -549,6 +580,8 @@ export type WireEvent = | StemAssetsEvent | StemFailedEvent | DepthAppliedEvent + | AudioEditAppliedEvent + | AudioEditFailedEvent | ManualSlotCountEvent | TimbreSetEvent | TimbreClearedEvent diff --git a/tests/unit/test_live_audio_edit.py b/tests/unit/test_live_audio_edit.py new file mode 100644 index 00000000..a11fcc34 --- /dev/null +++ b/tests/unit/test_live_audio_edit.py @@ -0,0 +1,311 @@ +from types import SimpleNamespace +import threading + +import numpy as np +import pytest +import torch + +from acestep.engine.stream import SlotRequest +from acestep.nodes.diffusion_nodes import StreamDenoise +from acestep.nodes.types import Latent +from acestep.streaming.audio_edit import ( + AudioEditError, + LiveAudioEdit, + EditRegion, + composite_window, + constrain_audio_edit, + parse_live_audio_edit, + regenerate_mask, + sa3_inpaint_bundle, +) +from acestep.streaming.ace_backend import ACEStepBackend +from acestep.streaming.audio_engine import AudioEngine +from acestep.streaming.session import StreamingSession + + +def test_stream_denoise_registers_hidden_audio_edit_snapshot(): + param = next( + p for p in StreamDenoise.get_definition().params + if p.name == "audio_edit" + ) + assert param.type == "any" + assert param.default is None + assert param.hidden is True + + +def test_explicit_right_extension_requires_only_the_added_tail(): + edit = parse_live_audio_edit( + [{"start_s": 10, "end_s": 15}], + enabled=True, + source_mode="waveform", + strength=1, + canvas_duration_s=15, + source_duration_s=10, + right_extension_s=5, + ) + assert edit.regions == (EditRegion(10, 15),) + + with pytest.raises(AudioEditError, match="explicit waveform extension"): + parse_live_audio_edit( + [{"start_s": 11, "end_s": 15}], + enabled=True, + source_mode="waveform", + strength=1, + canvas_duration_s=15, + source_duration_s=10, + right_extension_s=5, + ) + + +def test_ordinary_canvas_padding_does_not_require_tail_coverage(): + edit = parse_live_audio_edit( + [{"start_s": 2, "end_s": 3}], + enabled=True, + source_mode="waveform", + strength=1, + canvas_duration_s=60, + source_duration_s=10, + ) + assert edit.regions == (EditRegion(2, 3),) + + +def test_empty_waveform_edit_preserves_the_entire_source(): + edit = parse_live_audio_edit( + [], + enabled=True, + source_mode="waveform", + strength=1, + canvas_duration_s=1, + source_duration_s=1, + ) + assert edit.enabled is True + assert edit.regions == () + + source = np.linspace(-1, 1, 200, dtype=np.float32).reshape(100, 2) + generated = np.full((100, 2), 0.75, dtype=np.float32) + out = composite_window( + generated, + start_sample=0, + source=source, + edit=edit, + sample_rate=100, + crossfade_s=0, + ) + np.testing.assert_array_equal(out, source) + + +def test_empty_current_edit_rejects_stale_unmasked_generation(): + current = LiveAudioEdit(True, (), "waveform", 1) + stale = LiveAudioEdit(False) + effective = constrain_audio_edit(current, stale) + assert effective.enabled is True + assert effective.regions == () + + source = np.arange(200, dtype=np.float32).reshape(100, 2) + generated = np.full((100, 2), -123, dtype=np.float32) + np.testing.assert_array_equal( + composite_window( + generated, + start_sample=0, + source=source, + edit=effective, + sample_rate=100, + crossfade_s=0, + ), + source, + ) + + +def test_changed_regions_only_accept_current_and_emerged_intersection(): + current = LiveAudioEdit(True, (EditRegion(2, 5),), "waveform", 1) + emerged = LiveAudioEdit(True, (EditRegion(0, 3),), "waveform", 1) + effective = constrain_audio_edit(current, emerged) + assert effective.regions == (EditRegion(2, 3),) + + +def test_session_reasserts_empty_mask_after_runner_crossfade(): + source = np.arange(200, dtype=np.float32).reshape(100, 2) + generated = np.full((100, 2), -123, dtype=np.float32) + empty_edit = LiveAudioEdit(True, (), "waveform", 1) + published = [] + + streaming = object.__new__(StreamingSession) + streaming.backend = SimpleNamespace( + finalize_audio_edit_window=lambda pcm, start: composite_window( + pcm, + start_sample=start, + source=source, + edit=empty_edit, + sample_rate=100, + crossfade_s=0, + ), + ) + streaming._audio_edit_delivery_lock = threading.RLock() + streaming.audio_eng = AudioEngine(generated, 100) + streaming.state = SimpleNamespace(params={}, n_channels=2) + streaming.bus = SimpleNamespace(publish=published.append) + + streaming._on_audio_ready(generated.copy(), 0, 100) + + np.testing.assert_array_equal(streaming.audio_eng.current, source) + np.testing.assert_array_equal(published[0].audio, source) + + +def test_explicit_right_extension_excludes_backend_padding_after_request(): + edit = parse_live_audio_edit( + [{"start_s": 10, "end_s": 15}], + enabled=True, + source_mode="waveform", + strength=1, + canvas_duration_s=16, + source_duration_s=10, + right_extension_s=5, + ) + assert edit.regions == (EditRegion(10, 15),) + + +def test_explicit_left_extension_requires_only_the_added_tail(): + edit = parse_live_audio_edit( + [ + {"start_s": 0, "end_s": 5}, + {"start_s": 8, "end_s": 9}, + ], + enabled=True, + source_mode="waveform", + strength=1, + canvas_duration_s=15, + source_duration_s=10, + left_extension_s=5, + ) + assert edit.regions[0] == EditRegion(0, 5) + + with pytest.raises(AudioEditError, match="explicit waveform extension"): + parse_live_audio_edit( + [{"start_s": 1, "end_s": 5}], + enabled=True, + source_mode="waveform", + strength=1, + canvas_duration_s=15, + source_duration_s=10, + left_extension_s=5, + ) + + +def test_ace_mask_uses_absolute_regions_for_a_walk_window(): + edit = LiveAudioEdit(True, (EditRegion(61, 62),), "waveform", 0.5) + mask = regenerate_mask(edit, total_frames=50, rate_hz=25, offset_s=60) + assert torch.count_nonzero(mask == 0.5).item() == 25 + assert torch.count_nonzero(mask).item() == 25 + + +def test_window_compositor_restores_only_preserved_samples(): + source = torch.ones((2, 100)) + generated = np.zeros((40, 2), dtype=np.float32) + edit = LiveAudioEdit(True, (EditRegion(0.04, 0.06),), "waveform", 1) + out = composite_window( + generated, + start_sample=30, + source=source, + edit=edit, + sample_rate=1000, + crossfade_s=0, + ) + np.testing.assert_array_equal(out[:10], 1) + np.testing.assert_array_equal(out[10:30], 0) + np.testing.assert_array_equal(out[30:], 1) + + +def test_window_compositor_regenerates_left_extension_and_preserves_source(): + source = np.concatenate([ + np.zeros((20, 2), dtype=np.float32), + np.ones((80, 2), dtype=np.float32), + ]) + generated = np.full((100, 2), 0.25, dtype=np.float32) + edit = LiveAudioEdit(True, (EditRegion(0, 0.02),), "waveform", 1) + out = composite_window( + generated, + start_sample=0, + source=source, + edit=edit, + sample_rate=1000, + crossfade_s=0, + ) + np.testing.assert_array_equal(out[:20], 0.25) + np.testing.assert_array_equal(out[20:], 1) + + +def test_window_compositor_regenerates_right_extension_and_preserves_source(): + source = np.concatenate([ + np.ones((80, 2), dtype=np.float32), + np.zeros((20, 2), dtype=np.float32), + ]) + generated = np.full((100, 2), 0.25, dtype=np.float32) + edit = LiveAudioEdit(True, (EditRegion(0.08, 0.1),), "waveform", 1) + out = composite_window( + generated, + start_sample=0, + source=source, + edit=edit, + sample_rate=1000, + crossfade_s=0, + ) + np.testing.assert_array_equal(out[:80], 1) + np.testing.assert_array_equal(out[80:], 0.25) + + +def test_sa3_bundle_uses_preserve_polarity_and_masked_source(): + source = torch.ones((1, 20, 256)) + base = {"local_add_cond": torch.zeros((1, 257, 20))} + edit = LiveAudioEdit( + True, + (EditRegion(5 / (44100 / 4096), 10 / (44100 / 4096)),), + "waveform", + 1, + ) + bundle = sa3_inpaint_bundle(base, source, edit) + mask = bundle["local_add_cond"][:, :1] + assert torch.all(mask[..., :5] == 1) + assert torch.all(mask[..., 5:10] == 0) + assert torch.all(mask[..., 10:] == 1) + assert torch.equal(bundle["local_add_cond"][:, 1:], source.movedim(1, 2) * mask) + + +def test_ace_generate_attaches_edit_to_the_normal_stream_tick(): + edit = LiveAudioEdit(True, (EditRegion(0, 0.4),), "waveform", 1) + + class FakeStream: + def __init__(self): + self.pipeline = SimpleNamespace(last_finished_request=None) + self.kwargs = None + + def tick(self, **kwargs): + self.kwargs = kwargs + self.pipeline.last_finished_request = SlotRequest( + audio_edit=kwargs["audio_edit"], + ) + return Latent(tensor=torch.zeros((1, 25, 64))) + + backend = object.__new__(ACEStepBackend) + backend.stream = FakeStream() + backend._walk_active = False + backend._walk_chunk_start_s = 0 + backend._current_shift = 3 + backend._emerged_audio_edit = None + source = Latent(tensor=torch.zeros((1, 25, 64))) + prep = { + "raw": {}, + "source_lat": None, + "live_src_lat": source, + "audio_edit": edit, + "denoise": 1, + "seed": 1, + "x0_tgt": source, + "x0_target_curve": None, + "initial_noise_curve": None, + "tick_kwargs": {}, + } + backend._generate(prep) + sent = backend.stream.kwargs + assert sent["audio_edit"] is edit + assert sent["source_latent"].mask is not None + assert backend._emerged_audio_edit is edit