bithuman 1.0.2__py3-none-any.whl

bithuman/__init__.py

@@ -0,0 +1,13 @@
+from ._version import __version__
+from .api import AudioChunk, VideoControl, VideoFrame
+from .runtime import Bithuman
+from .runtime_async import AsyncBithuman
+
+__all__ = [
+    "__version__",
+    "Bithuman",
+    "AsyncBithuman",
+    "AudioChunk",
+    "VideoControl",
+    "VideoFrame",
+]

bithuman/_version.py

@@ -0,0 +1 @@
+__version__ = "1.0.2"

bithuman/api.py

@@ -0,0 +1,164 @@
+"""API for Bithuman Runtime."""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Hashable, List, Optional
+
+import numpy as np
+from dataclasses_json import dataclass_json
+
+INT16_MAX = 2**15 - 1
+
+
+@dataclass
+class AudioChunk:
+    """Data class to store audio data.
+
+    Attributes:
+        data (np.ndarray): Audio data in int16 format with 1 channel, shape (n_samples,).
+        sample_rate (int): Sample rate of the audio data.
+        last_chunk (bool): Whether this is the last chunk of the speech.
+    """
+
+    data: np.ndarray
+    sample_rate: int
+    last_chunk: bool = True
+
+    def __post_init__(self) -> None:
+        if isinstance(self.data, bytes):
+            self.data = np.frombuffer(self.data, dtype=np.int16)
+
+        if self.data.dtype in [np.float32, np.float64]:
+            self.data = (self.data * INT16_MAX).astype(np.int16)
+
+        if not isinstance(self.data, np.ndarray) or self.data.dtype != np.int16:
+            raise ValueError("data must be in int16 numpy array format")
+
+    @property
+    def bytes(self) -> bytes:
+        return self.data.tobytes()
+
+    @property
+    def array(self) -> np.ndarray:
+        return self.data
+
+    @property
+    def duration(self) -> float:
+        return len(self.data) / self.sample_rate
+
+    @classmethod
+    def from_bytes(
+        cls,
+        audio_bytes: bytes,
+        sample_rate: int,
+        last_chunk: bool = True,
+    ) -> "AudioChunk":
+        return cls(
+            data=np.frombuffer(audio_bytes, dtype=np.int16),
+            sample_rate=sample_rate,
+            last_chunk=last_chunk,
+        )
+
+
+@dataclass
+class VideoControl:
+    """Dataclass for video control information."""
+
+    audio: Optional[AudioChunk] = None
+    text: Optional[str] = None
+    target_video: Optional[str] = None
+    action: Optional[str | List[str]] = None
+    emotion_preds: Optional[List["EmotionPrediction"]] = None
+    message_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    end_of_speech: bool = False  # mark the end of the speech
+    force_action: bool = False  # force the action to be played
+    stop_on_user_speech: Optional[bool] = None  # Override stop_on_user_speech from video config
+    stop_on_agent_speech: Optional[bool] = None  # Override stop_on_agent_speech from video config
+
+    @property
+    def is_idle(self) -> bool:
+        """Check if the audio or control is idle."""
+        return (
+            self.audio is None
+            and self.target_video is None
+            and not self.action
+            and not self.emotion_preds
+        )
+
+    @property
+    def is_speaking(self) -> bool:
+        """Check if the audio or control is speaking."""
+        return self.audio is not None
+
+    @classmethod
+    def from_audio(
+        cls, audio: bytes | np.ndarray, sample_rate: int, last_chunk: bool = True
+    ) -> "VideoControl":
+        if isinstance(audio, bytes):
+            audio_chunk = AudioChunk.from_bytes(audio, sample_rate, last_chunk)
+        elif isinstance(audio, np.ndarray):
+            audio_chunk = AudioChunk(
+                data=audio, sample_rate=sample_rate, last_chunk=last_chunk
+            )
+        else:
+            raise ValueError("audio must be bytes or numpy array")
+        return cls(audio=audio_chunk)
+
+
+@dataclass
+class VideoFrame:
+    """
+    Dataclass for frame information.
+
+    Attributes:
+        bgr_image: image data in uint8 BGR format, shape (H, W, 3).
+        audio_chunk: audio chunk if the frame is talking,
+            chunked to the duration of this video frame based on FPS.
+        frame_index: frame index generated for this video control.
+        source_message_id: message id of the video control.
+    """
+
+    bgr_image: Optional[np.ndarray] = None
+    audio_chunk: Optional[AudioChunk] = None
+    frame_index: Optional[int] = None
+    source_message_id: Optional[Hashable] = None
+    end_of_speech: bool = False  # mark the end of the speech
+
+    @property
+    def has_image(self) -> bool:
+        return self.bgr_image is not None
+
+    @property
+    def rgb_image(self) -> Optional[np.ndarray]:
+        if self.bgr_image is None:
+            return None
+        return self.bgr_image[:, :, ::-1]
+
+
+class Emotion(str, Enum):
+    """Enumeration representing different emotions."""
+
+    ANGER = "anger"
+    DISGUST = "disgust"
+    FEAR = "fear"
+    JOY = "joy"
+    NEUTRAL = "neutral"
+    SADNESS = "sadness"
+    SURPRISE = "surprise"
+
+
+@dataclass_json
+@dataclass
+class EmotionPrediction:
+    """Dataclass for emotion prediction.
+
+    Attributes:
+        emotion: The emotion.
+        score: The score for the emotion.
+    """
+
+    emotion: Emotion
+    score: float

bithuman/audio/__init__.py

@@ -0,0 +1,19 @@
+from .audio import get_mel_chunks
+from .utils import (
+    AudioStreamBatcher,
+    float32_to_int16,
+    int16_to_float32,
+    load_audio,
+    resample,
+    write_video_with_audio,
+)
+
+__all__ = [
+    "get_mel_chunks",
+    "load_audio",
+    "resample",
+    "write_video_with_audio",
+    "AudioStreamBatcher",
+    "float32_to_int16",
+    "int16_to_float32",
+]

bithuman/audio/audio.py

@@ -0,0 +1,396 @@
+from __future__ import annotations
+
+import librosa
+import librosa.filters
+import numba
+import numpy as np
+from scipy import signal
+from scipy.io import wavfile
+
+from .hparams import hparams as hp
+
+
+def get_mel_chunks(
+    audio_np: np.ndarray, fps: int, mel_step_size: int = 16
+) -> list[np.ndarray]:
+    """
+    Get mel chunks from audio.
+
+    Args:
+        audio_np: The audio data in numpy array, `np.float32` in 16kHz with 1 channel.
+        fps: The frame per second of the video.
+        mel_step_size: The step size for mel spectrogram chunks
+
+    Returns:
+        The mel chunks.
+    """
+    # Convert the audio waveform to a mel spectrogram
+
+    mel_chunks = []
+    if audio_np is None or len(audio_np) == 0:
+        return mel_chunks
+
+    mel = melspectrogram(audio_np)
+
+    # Check for NaN values in the mel spectrogram, which could indicate issues with audio processing
+    if np.isnan(mel.reshape(-1)).sum() > 0:
+        raise ValueError(
+            "Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again"
+        )
+
+    # Split the mel spectrogram into smaller chunks suitable for processing
+    mel_idx_multiplier = 80.0 / fps  # Calculate index multiplier based on frame rate
+    i = 0
+    while True:
+        start_idx = int(i * mel_idx_multiplier)
+        if start_idx + mel_step_size > len(mel[0]):  # Handle last chunk
+            mel_chunks.append(mel[:, len(mel[0]) - mel_step_size :])
+            break
+        mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size])
+        i += 1
+
+    return mel_chunks
+
+
+def save_wav(wav, path, sr):
+    """
+    Save a wav file.
+
+    Args:
+        wav (np.ndarray): The wav file to save, as a numpy array.
+        path (str): The path to save the wav file to.
+        sr (int): The sample rate to use when saving the wav file.
+    """
+    wav *= 32767 / max(0.01, np.max(np.abs(wav)))
+    wavfile.write(path, sr, wav.astype(np.int16))
+
+
+def preemphasis(wav, k, preemphasize=True):
+    """
+    Apply preemphasis to a wav file.
+
+    Args:
+        wav (np.ndarray): The wav file to apply preemphasis to, as a numpy array.
+        k (float): The preemphasis coefficient.
+        preemphasize (bool, optional): Whether to apply preemphasis. Defaults to True.
+
+    Returns:
+        np.ndarray: The wav file with preemphasis applied.
+    """
+    if preemphasize:
+        return signal.lfilter([1, -k], [1], wav)
+    return wav
+
+
+def inv_preemphasis(wav, k, inv_preemphasize=True):
+    """
+    Apply inverse preemphasis to a wav file.
+
+    Args:
+        wav (np.ndarray): The wav file to apply inverse preemphasis to, as a numpy array.
+        k (float): The preemphasis coefficient.
+        inv_preemphasize (bool, optional): Whether to apply inverse preemphasis. Defaults to True.
+
+    Returns:
+        np.ndarray: The wav file with inverse preemphasis applied.
+    """
+    if inv_preemphasize:
+        return signal.lfilter([1], [1, -k], wav)
+    return wav
+
+
+def get_hop_size():
+    """
+    Get the hop size.
+
+    Returns:
+        int: The hop size.
+    """
+    hop_size = hp.hop_size
+    if hop_size is None:
+        assert hp.frame_shift_ms is not None
+        hop_size = int(hp.frame_shift_ms / 1000 * hp.sample_rate)
+    return hop_size
+
+
+def linearspectrogram(wav):
+    """
+    Compute the linear spectrogram of a wav file.
+
+    Args:
+        wav (np.ndarray): The wav file to compute the spectrogram of, as a numpy array.
+
+    Returns:
+        np.ndarray: The linear spectrogram of the wav file.
+    """
+    D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize))
+    S = _amp_to_db(np.abs(D)) - hp.ref_level_db
+
+    if hp.signal_normalization:
+        return _normalize(S)
+    return S
+
+
+def melspectrogram(wav):
+    """
+    Compute the mel spectrogram of a wav file.
+
+    Args:
+        wav (np.ndarray): The wav file to compute the spectrogram of, as a numpy array.
+
+    Returns:
+        np.ndarray: The mel spectrogram of the wav file.
+    """
+    D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize))
+    S = _amp_to_db(_linear_to_mel(np.abs(D))) - hp.ref_level_db
+
+    if hp.signal_normalization:
+        return _normalize(S)
+    return S
+
+
+def _lws_processor():
+    # Third-Party Libraries
+    import lws
+
+    return lws.lws(hp.n_fft, get_hop_size(), fftsize=hp.win_size, mode="speech")
+
+
+def _stft(y):
+    """
+    Compute the short-time Fourier transform of a signal.
+
+    Args:
+        y (np.ndarray): The signal to compute the transform of, as a numpy array.
+
+    Returns:
+        np.ndarray: The short-time Fourier transform of the signal.
+    """
+    if hp.use_lws:
+        return _lws_processor(hp).stft(y).T
+    else:
+        return librosa.stft(
+            y=y,
+            n_fft=hp.n_fft,
+            hop_length=get_hop_size(),
+            win_length=hp.win_size,
+        )
+
+
+##########################################################
+# Those are only correct when using lws!!! (This was messing with Wavenet quality for a long time!)
+def num_frames(length, fsize, fshift):
+    """
+    Compute number of time frames of spectrogram.
+
+    Args:
+        length: length of signal
+        fsize: frame size
+        fshift: frame shift
+
+    """
+    pad = fsize - fshift
+    if length % fshift == 0:
+        M = (length + pad * 2 - fsize) // fshift + 1
+    else:
+        M = (length + pad * 2 - fsize) // fshift + 2
+    return M
+
+
+def pad_lr(x, fsize, fshift):
+    """
+    Compute left and right padding.
+
+    Args:
+        x: input signal
+        fsize: frame size
+        fshift: frame shift
+    """
+    M = num_frames(len(x), fsize, fshift)
+    pad = fsize - fshift
+    T = len(x) + 2 * pad
+    r = (M - 1) * fshift + fsize - T
+    return pad, pad + r
+
+
+##########################################################
+# Librosa correct padding
+def librosa_pad_lr(x, fsize, fshift):
+    """
+    Compute left and right padding for librosa.
+
+    Args:
+        x: input signal
+        fsize: frame size
+        fshift: frame shift.
+
+    """
+    return 0, (x.shape[0] // fshift + 1) * fshift - x.shape[0]
+
+
+# Conversions
+_mel_basis = None
+_mel_basis_cache = None
+_output_cache = None
+
+
+@numba.jit(nopython=True, cache=True)
+def _dot_mel(mel_basis, spectogram):
+    return np.dot(mel_basis, spectogram)
+
+
+def _linear_to_mel_numba(spectogram):
+    """
+    Convert a linear spectrogram to a mel spectrogram using optimized matrix multiplication.
+
+    Args:
+        spectogram (np.ndarray): The linear spectrogram to convert, shape around [401, 22]
+
+    Returns:
+        np.ndarray: The mel spectrogram, shape [80, 22]
+    """
+    global _mel_basis
+    if _mel_basis is None:
+        _mel_basis = _build_mel_basis()
+        _mel_basis = np.ascontiguousarray(_mel_basis, dtype=np.float32)
+
+    spectogram = np.ascontiguousarray(spectogram, dtype=np.float32)
+    return _dot_mel(_mel_basis, spectogram)
+
+
+@numba.jit(nopython=True, cache=True, parallel=False, fastmath=True)
+def _dot_mel_hybrid(mel_basis, spectogram, out):
+    """Hybrid approach using vectorized operations with controlled CPU usage."""
+    # Process in smaller chunks to reduce CPU pressure
+    chunk_size = 16  # Adjust this value to balance speed/CPU usage
+
+    for j in range(0, spectogram.shape[1], chunk_size):
+        end_j = min(j + chunk_size, spectogram.shape[1])
+        # Use numpy's dot for each chunk
+        out[:, j:end_j] = mel_basis.dot(np.ascontiguousarray(spectogram[:, j:end_j]))
+
+    return out
+
+
+def _linear_to_mel_hybrid(spectogram):
+    """Hybrid implementation balancing speed and CPU usage."""
+    global _mel_basis_cache, _output_cache
+
+    if _mel_basis_cache is None:
+        _mel_basis_cache = _build_mel_basis()
+        _mel_basis_cache = np.ascontiguousarray(_mel_basis_cache, dtype=np.float32)
+
+    spectogram = np.ascontiguousarray(spectogram, dtype=np.float32)
+
+    # Reuse output array
+    if _output_cache is None or _output_cache.shape != (
+        _mel_basis_cache.shape[0],
+        spectogram.shape[1],
+    ):
+        _output_cache = np.empty(
+            (_mel_basis_cache.shape[0], spectogram.shape[1]), dtype=np.float32
+        )
+
+    return _dot_mel_hybrid(_mel_basis_cache, spectogram, _output_cache)
+
+
+_linear_to_mel = _linear_to_mel_hybrid
+
+
+def _build_mel_basis():
+    """
+    Build the mel basis for converting linear spectrograms to mel spectrograms.
+
+    Returns:
+        np.ndarray: The mel basis.
+    """
+    assert hp.fmax <= hp.sample_rate // 2
+    return librosa.filters.mel(
+        sr=hp.sample_rate,
+        n_fft=hp.n_fft,
+        n_mels=hp.num_mels,
+        fmin=hp.fmin,
+        fmax=hp.fmax,
+    )
+
+
+def _amp_to_db(x):
+    """
+    Convert amplitude to decibels.
+
+    Args:
+        x (np.ndarray): The amplitude to convert, as a numpy array.
+
+    Returns:
+        np.ndarray: The decibels.
+    """
+    min_level = np.float32(np.exp(hp.min_level_db / 20 * np.log(10)))
+    return np.float32(20) * np.log10(np.maximum(min_level, x)).astype(np.float32)
+
+
+def _db_to_amp(x):
+    return np.power(10.0, (x) * 0.05)
+
+
+def _normalize(S):
+    """
+    Normalize a spectrogram.
+
+    Args:
+        S (np.ndarray): The spectrogram to normalize, as a numpy array.
+
+    Returns:
+        np.ndarray: The normalized spectrogram.
+    """
+    if hp.allow_clipping_in_normalization:
+        if hp.symmetric_mels:
+            return np.clip(
+                (2 * hp.max_abs_value) * ((S - hp.min_level_db) / (-hp.min_level_db))
+                - hp.max_abs_value,
+                -hp.max_abs_value,
+                hp.max_abs_value,
+            )
+        else:
+            return np.clip(
+                hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db)),
+                0,
+                hp.max_abs_value,
+            )
+
+    assert S.max() <= 0 and S.min() - hp.min_level_db >= 0
+    if hp.symmetric_mels:
+        return (2 * hp.max_abs_value) * (
+            (S - hp.min_level_db) / (-hp.min_level_db)
+        ) - hp.max_abs_value
+    else:
+        return hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db))
+
+
+def _denormalize(D):
+    """
+    Denormalize a spectrogram.
+
+    Args:
+        D: spectrogram to denormalize
+
+    Returns:
+        denormalized spectrogram
+    """
+    if hp.allow_clipping_in_normalization:
+        if hp.symmetric_mels:
+            return (
+                (np.clip(D, -hp.max_abs_value, hp.max_abs_value) + hp.max_abs_value)
+                * -hp.min_level_db
+                / (2 * hp.max_abs_value)
+            ) + hp.min_level_db
+        else:
+            return (
+                np.clip(D, 0, hp.max_abs_value) * -hp.min_level_db / hp.max_abs_value
+            ) + hp.min_level_db
+
+    if hp.symmetric_mels:
+        return (
+            (D + hp.max_abs_value) * -hp.min_level_db / (2 * hp.max_abs_value)
+        ) + hp.min_level_db
+    else:
+        return (D * -hp.min_level_db / hp.max_abs_value) + hp.min_level_db

bithuman/audio/hparams.py

@@ -0,0 +1,108 @@
+"""hparams module."""
+
+
+class HParams:
+    """
+    A class for managing hyperparameters.
+
+    This class provides an interface for getting and setting hyperparameters.
+    Hyperparameters are stored in a dictionary.
+    """
+
+    def __init__(self, **kwargs):
+        """
+        Initialize a HParams instance.
+
+        Args:
+            **kwargs: The initial hyperparameters, as keyword arguments.
+        """
+        self.data = {}
+
+        for key, value in kwargs.items():
+            self.data[key] = value
+
+    def __getattr__(self, key):
+        """
+        Get a hyperparameter.
+
+        Args:
+            key (str): The name of the hyperparameter.
+
+        Returns:
+            The value of the hyperparameter.
+
+        Raises:
+            AttributeError: If the hyperparameter does not exist.
+        """
+        if key not in self.data:
+            raise AttributeError("'HParams' object has no attribute %s" % key)
+        return self.data[key]
+
+    def set_hparam(self, key, value):
+        """
+        Set a hyperparameter.
+
+        Args:
+            key (str): The name of the hyperparameter.
+            value: The new value of the hyperparameter.
+        """
+        self.data[key] = value
+
+
+# Default hyperparameters
+hparams = HParams(
+    num_mels=80,  # Number of mel-spectrogram channels and local conditioning dimensionality
+    #  network
+    rescale=True,  # Whether to rescale audio prior to preprocessing
+    rescaling_max=0.9,  # Rescaling value
+    # Use LWS (https://github.com/Jonathan-LeRoux/lws) for STFT and phase reconstruction
+    # It"s preferred to set True to use with https://github.com/r9y9/wavenet_vocoder
+    # Does not work if n_ffit is not multiple of hop_size!!
+    use_lws=False,
+    n_fft=800,  # Extra window size is filled with 0 paddings to match this parameter
+    hop_size=200,  # For 16000Hz, 200 = 12.5 ms (0.0125 * sample_rate)
+    win_size=800,  # For 16000Hz, 800 = 50 ms (If None, win_size = n_fft) (0.05 * sample_rate)
+    sample_rate=16000,  # 16000Hz (corresponding to librispeech) (sox --i <filename>)
+    frame_shift_ms=None,  # Can replace hop_size parameter. (Recommended: 12.5)
+    # Mel and Linear spectrograms normalization/scaling and clipping
+    signal_normalization=True,
+    # Whether to normalize mel spectrograms to some predefined range (following below parameters)
+    allow_clipping_in_normalization=True,  # Only relevant if mel_normalization = True
+    symmetric_mels=True,
+    # Whether to scale the data to be symmetric around 0. (Also multiplies the output range by 2,
+    # faster and cleaner convergence)
+    max_abs_value=4.0,
+    # max absolute value of data. If symmetric, data will be [-max, max] else [0, max] (Must not
+    # be too big to avoid gradient explosion,
+    # not too small for fast convergence)
+    # Contribution by @begeekmyfriend
+    # Spectrogram Pre-Emphasis (Lfilter: Reduce spectrogram noise and helps model certitude
+    # levels. Also allows for better G&L phase reconstruction)
+    preemphasize=True,  # whether to apply filter
+    preemphasis=0.97,  # filter coefficient.
+    # Limits
+    min_level_db=-100,
+    ref_level_db=20,
+    fmin=55,
+    # Set this to 55 if your speaker is male! if female, 95 should help taking off noise. (To
+    # test depending on dataset. Pitch info: male~[65, 260], female~[100, 525])
+    fmax=7600,  # To be increased/reduced depending on data.
+    ###################### Our training parameters #################################
+    img_size=96,
+    fps=25,
+    batch_size=16,
+    initial_learning_rate=1e-4,
+    nepochs=200000000000000000,
+    ### ctrl + c, stop whenever eval loss is consistently greater than train loss for ~10 epochs
+    num_workers=16,
+    checkpoint_interval=3000,
+    eval_interval=3000,
+    save_optimizer_state=True,
+    syncnet_wt=0.0,  # is initially zero, will be set automatically to 0.03 later. Leads to faster convergence.
+    syncnet_batch_size=64,
+    syncnet_lr=1e-4,
+    syncnet_eval_interval=10000,
+    syncnet_checkpoint_interval=10000,
+    disc_wt=0.07,
+    disc_initial_learning_rate=1e-4,
+)