diffusers-workflow 0.4.0__py3-none-any.whl

dw/tasks/speech_generation.py

@@ -0,0 +1,228 @@
+"""
+Speech generation via HuggingFace transformers.
+
+Takes a line of text and speaks it with a local text-to-speech model, returning
+the waveform and the rate it was generated at so the rest of the audio plumbing -
+slice_audio, fade_audio and pair_audio - composes with it directly (concat_videos
+and dissolve_videos join videos; pair the track onto a video first).
+
+The role this is built for is voice *timbre reference*, not the track a mouth
+follows. MiniMax H3 lip-syncs well when it generates the speech itself and poorly
+when it must follow supplied audio, so its audio reference takes a few seconds of
+a voice to fix timbre, pitch and delivery while the model still generates the
+line. A generated clip referenced in every shot makes voice consistency an actual
+conditioning signal rather than a prose description that has to land identically
+a dozen times. The other honest uses are a voice that must be matched, and
+narration over shots where nothing has to lip-sync to it.
+"""
+
+import logging
+
+import torch
+from transformers import pipeline as hf_pipeline
+
+from .. import preferred_task_dtype
+from ..result import AudioTrack
+from .audio_utils import as_channels_samples, load_audio, resample_waveform
+from .model_cache import cached_model, hf_pipeline_placement
+
+logger = logging.getLogger("dw")
+
+# Small, loads without a separate speaker-embedding dataset, and its voice presets
+# give distinct speakers - which is the point when two characters have to sound
+# like two people. A single-voice model like facebook/mms-tts-eng (which takes no
+# voice_preset) is a quarter the size and a reasonable override where only one
+# voice is needed
+_DEFAULT_MODEL = "suno/bark-small"
+
+# What SpeechT5's voice-cloning recipes are built around - trained on the same
+# corpus (VoxCeleb) the CMU ARCTIC x-vectors that ship with the model card come
+# from, so a reference clip reduces to a vector in the space the model expects
+_SPEAKER_ENCODER_MODEL = "speechbrain/spkrec-xvect-voxceleb"
+# What the encoder was trained on - resampling anything else to this rate is
+# part of extracting a voiceprint, not an approximation of one
+_SPEAKER_ENCODER_SAMPLE_RATE = 16000
+
+
+def _speaker_embedding_tensor(location, device, dtype):
+    """The x-vector speechbrain's spkrec-xvect-voxceleb extracts from a
+    reference audio file - what SpeechT5 conditions its voice on.
+
+    Args:
+        location: Path to a reference audio file (an 'asset:' reference has
+            already resolved to this by the time a task sees it).
+        device: Where to run the encoder.
+        dtype: The speech pipeline's own dtype - the embedding is concatenated
+            with the pipe's hidden states inside the decoder prenet, so it must
+            match the pipe's precision (fp16 on CUDA) or the linear layer there
+            refuses a Float x Half product.
+    """
+    from speechbrain.inference.speaker import EncoderClassifier
+
+    def load_encoder():
+        return EncoderClassifier.from_hparams(
+            source=_SPEAKER_ENCODER_MODEL,
+            run_opts={"device": str(device)},
+        )
+
+    encoder = cached_model(
+        ("speaker_encoder", _SPEAKER_ENCODER_MODEL, str(device)), load_encoder
+    )
+
+    waveform, sample_rate = load_audio(location)
+    waveform = resample_waveform(waveform, sample_rate, _SPEAKER_ENCODER_SAMPLE_RATE)
+    mono = waveform.mean(axis=0)
+
+    with torch.no_grad():
+        embedding = encoder.encode_batch(torch.as_tensor(mono).unsqueeze(0))
+        embedding = torch.nn.functional.normalize(embedding, dim=2)
+    # SpeechT5's generate() wants (batch, 512); the encoder's raw output is
+    # (1, 1, 512), so squeeze collapses it back to (512,) before restoring
+    # the batch dimension the model actually requires
+    return embedding.squeeze().unsqueeze(0).to(device=device, dtype=dtype)
+
+
+def generate_speech(text=None, device="cpu", seed=None, **kwargs):
+    """Speak a line of text with a local text-to-speech model.
+
+    Args:
+        text: The line to speak. Mutually exclusive with messages - exactly
+            one of the two is required.
+        device: Target device ("cuda", "mps", "cpu").
+        seed: The workflow/step-resolved seed, when one was set. A pipeline
+            step gets an explicit torch.Generator; transformers' generate()
+            takes no such argument, so reproducing Bark (and any other
+            sampling model here) means seeding the global RNG immediately
+            before the call - unseeded when None, matching every other
+            unseeded task (#261).
+        **kwargs:
+            messages: Chat-templated input for a model such as VibeVoice that
+                takes a conversation rather than a bare string - a list of
+                {"role": ..., "content": ...} dicts. Passed straight through
+                as the pipeline's text_inputs, which applies the model's own
+                chat template; a model with no chat template configured (Bark
+                and friends) raises when handed this instead of text.
+            model_name: HuggingFace model ID. Defaults to suno/bark-small.
+            voice_preset: The speaker to use, for a model that has presets -
+                "v2/en_speaker_6" and friends for Bark. This is a preprocessing
+                argument: it selects the speaker before generation rather than
+                parameterizing it, which is why it is named here rather than left
+                to forward_params, where it would be silently dropped.
+            speaker_embedding: Path to a reference audio file (typically an
+                'asset:' reference) whose voice a SpeechT5 model should speak
+                in. Reduced to an x-vector with speechbrain's
+                spkrec-xvect-voxceleb and injected into forward_params as
+                'speaker_embeddings' - SpeechT5 is the only pipeline here that
+                conditions on one, and it is not optional for SpeechT5: the
+                model refuses to generate without a speaker embedding, so this
+                is required whenever model_name is a SpeechT5 checkpoint. A
+                VITS model's speaker instead takes a plain 'speaker_id' int,
+                which already reaches the model unchanged through
+                forward_params and needs no argument of its own.
+            forward_params: Passed to the model's forward/generate call.
+            generate_kwargs: Ad-hoc generation settings for a generative model -
+                temperature, do_sample and so on.
+
+    Returns:
+        An AudioTrack holding the waveform, shaped (channels, samples), and the
+        sample rate the model generated it at.
+
+    Raises:
+        ValueError: If the model reports no sample rate for what it generated,
+            or if text/messages are both given or neither is.
+    """
+    messages = kwargs.get("messages", None)
+    if (text is None) == (messages is None):
+        raise ValueError(
+            "generate_speech needs exactly one of 'text' or 'messages' - "
+            "a plain line to speak, or chat-templated input for a model "
+            "such as VibeVoice"
+        )
+    if messages is not None:
+        valid = isinstance(messages, list) and len(messages) > 0
+        if valid:
+            for message in messages:
+                if (
+                    not isinstance(message, dict)
+                    or not isinstance(message.get("role"), str)
+                    or not isinstance(message.get("content"), str)
+                ):
+                    valid = False
+                    break
+        if not valid:
+            raise ValueError(
+                "generate_speech's 'messages' needs a non-empty list of "
+                "{'role': ..., 'content': ...} dicts, both strings - a bare "
+                "string or a wrong-keyed dict is not chat-templated input"
+            )
+    text_inputs = messages if messages is not None else text
+
+    model_name = kwargs.get("model_name", _DEFAULT_MODEL)
+    voice_preset = kwargs.get("voice_preset", None)
+    speaker_embedding = kwargs.get("speaker_embedding", None)
+    dtype = preferred_task_dtype(device)
+
+    def load_pipe():
+        logger.info(f"Generating speech with {model_name} on {device}")
+        placement = hf_pipeline_placement(device)
+        return hf_pipeline(
+            "text-to-speech",
+            model=model_name,
+            torch_dtype=dtype,
+            **placement,
+        )
+
+    pipe = cached_model(
+        ("speech_generation", model_name, str(device), str(dtype)),
+        load_pipe,
+    )
+
+    if voice_preset and getattr(pipe, "processor", None) is None:
+        # A single-voice model has no processor to hand the preset to;
+        # transformers logs the kwarg as unrecognised and speaks anyway
+        raise ValueError(
+            f"{model_name} takes no 'voice_preset' - it has one voice. Drop the "
+            "preset, or use a model with speaker presets such as suno/bark-small"
+        )
+
+    forward_params = kwargs.get("forward_params") or {}
+    if speaker_embedding:
+        model_type = getattr(getattr(pipe.model, "config", None), "model_type", None)
+        if model_type != "speecht5":
+            # Only SpeechT5 conditions on an x-vector; a model that does not
+            # would drop 'speaker_embeddings' as an unrecognised forward kwarg
+            # and generate in its own voice, same failure mode as voice_preset
+            raise ValueError(
+                f"{model_name} takes no 'speaker_embedding' - only a SpeechT5 "
+                "model conditions on an x-vector. Drop it, or use a SpeechT5 "
+                "model such as microsoft/speecht5_tts"
+            )
+        forward_params = {
+            **forward_params,
+            "speaker_embeddings": _speaker_embedding_tensor(
+                speaker_embedding, device, dtype
+            ),
+        }
+
+    if messages is not None:
+        logger.info(f"Speaking {len(messages)} chat-templated message(s)")
+    else:
+        logger.info(f"Speaking: {text[:100]}{'...' if len(text) > 100 else ''}")
+    if seed is not None:
+        torch.manual_seed(seed)
+    output = pipe(
+        text_inputs,
+        preprocess_params={"voice_preset": voice_preset} if voice_preset else {},
+        forward_params=forward_params,
+        generate_kwargs=kwargs.get("generate_kwargs") or {},
+    )
+
+    sample_rate = output.get("sampling_rate")
+    if sample_rate is None:
+        # Saving at the 44100 default instead would be quietly wrong rather than
+        # loud, and speech at the wrong rate is wrong in pitch as well as length
+        raise ValueError(
+            f"{model_name} reported no sample rate for the speech it generated"
+        )
+
+    return AudioTrack(as_channels_samples(output["audio"]), int(sample_rate))

dw/tasks/stabilize.py

@@ -0,0 +1,129 @@
+"""Hold a generated clip's framing still.
+
+A video model drifts. Ask it for a locked camera and the framing still wanders
+over a few seconds - a slow slide of the whole picture that nothing in the
+prompt asked for. It goes unnoticed inside a single shot, and becomes obvious
+the moment two shots are cut together and the subject jumps back to where it
+started.
+
+The wander is a global translation, so phase correlation measures it: the
+cross-power spectrum of consecutive frames peaks at the shift between them.
+Accumulating those shifts gives the clip's drift, and shifting every frame back
+by it puts the framing where it began.
+"""
+
+import logging
+
+import numpy as np
+from PIL import Image
+
+from ..result import AudioVideo
+from ..shots import carried_shots
+from .video_utils import frames_as_pil_list, load_audio_video
+
+logger = logging.getLogger("dw")
+
+
+def _pair_shift(previous, following, window):
+    """The translation between two grayscale frames, in whole pixels."""
+    height, width = previous.shape
+    a = np.fft.rfft2(previous * window)
+    b = np.fft.rfft2(following * window)
+    cross = a * np.conj(b)
+    cross /= np.abs(cross) + 1e-8
+    correlation = np.fft.irfft2(cross, s=previous.shape)
+
+    # The peak of this cross-power spectrum sits at the negative of the
+    # displacement, so it is negated here and every caller reads (dx, dy) as
+    # "how far the picture moved between these two frames".
+    peak = np.unravel_index(np.argmax(correlation), correlation.shape)
+    dy = peak[0] - height if peak[0] > height // 2 else peak[0]
+    dx = peak[1] - width if peak[1] > width // 2 else peak[1]
+    return -dx, -dy
+
+
+def _moving_average(values, window):
+    """Trajectory smoothed over `window` frames, with the ends held."""
+    pad = window // 2
+    padded = np.pad(values, ((pad, pad), (0, 0)), mode="edge")
+    kernel = np.ones(window, dtype=np.float32) / window
+    return np.stack(
+        [
+            np.convolve(padded[:, axis], kernel, mode="valid")[: len(values)]
+            for axis in (0, 1)
+        ],
+        axis=1,
+    )
+
+
+def stabilize_video(clip, smooth=0):
+    """Task command: remove a generated clip's accumulated framing drift.
+
+    Args:
+        clip: The video - a frame list, a frame array or tensor, an
+            AudioVideo, whose soundtrack is carried through untouched, or the
+            path or URL of a video file, which is read with its audio, so a
+            shot an earlier run wrote can be steadied without regenerating it.
+            The argument is deliberately not called "video": the engine loads
+            an argument by that name itself, as bare frames, which would strip
+            the soundtrack off before this ever saw it
+        smooth: 0 locks the framing to the first frame, which is what a shot
+            generated from a pinned keyframe wants. A window in frames instead
+            removes only the wander faster than that window, so a slow
+            deliberate move survives and the drift around it does not
+    Returns:
+        The stabilized clip, cropped to the region every frame covers and
+        resized back to its original size - an AudioVideo when one came in
+    """
+    if isinstance(clip, str):
+        clip = load_audio_video(clip)
+
+    frames = frames_as_pil_list(clip)
+    if len(frames) < 2:
+        return clip
+
+    width, height = frames[0].size
+    gray = [np.asarray(f.convert("L"), dtype=np.float32) for f in frames]
+    window = np.outer(np.hanning(height), np.hanning(width))
+
+    trajectory = np.zeros((len(frames), 2), dtype=np.float32)
+    for index in range(1, len(frames)):
+        dx, dy = _pair_shift(gray[index - 1], gray[index], window)
+        trajectory[index] = trajectory[index - 1] + (dx, dy)
+
+    target = _moving_average(trajectory, smooth) if smooth > 1 else 0.0
+    correction = np.rint(trajectory - target).astype(np.int32)
+    logger.debug(
+        f"Stabilizing {len(frames)} frames, peak drift "
+        f"{np.abs(correction).max()}px of {max(width, height)}"
+    )
+    if not correction.any():
+        return clip
+
+    # Shifting frames back uncovers their edges, so keep only the rectangle
+    # every frame still covers, then put that back at the original size.
+    shifts = -correction
+    left = int(max(0, shifts[:, 0].max()))
+    right = int(width + min(0, shifts[:, 0].min()))
+    top = int(max(0, shifts[:, 1].max()))
+    bottom = int(height + min(0, shifts[:, 1].min()))
+
+    held = []
+    for frame, (sx, sy) in zip(frames, shifts):
+        moved = np.roll(np.asarray(frame), (int(sy), int(sx)), axis=(0, 1))
+        held.append(
+            Image.fromarray(moved[top:bottom, left:right]).resize(
+                (width, height), Image.LANCZOS
+            )
+        )
+
+    if isinstance(clip, AudioVideo):
+        # Same frames, one for one, so every shot boundary still holds
+        return AudioVideo(
+            held,
+            clip.audio,
+            clip.sample_rate,
+            fps=clip.fps,
+            shots=carried_shots(clip),
+        )
+    return held