valenti/spark-pocket-tts
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
d4004e513e995b8175077283420625a6ae937ca0
spark-pocket-tts/pocket_tts/models/tts_model.py
T

720 lines
30 KiB
Python
Raw Blame History

import copy
import logging
import os
import queue
import statistics
import threading
import time
from functools import lru_cache
from pathlib import Path
import safetensors
import torch
from torch import nn
from torch.nn import functional as F
from typing_extensions import Self
from pocket_tts.conditioners.base import TokenizedText
from pocket_tts.data.audio import audio_read
from pocket_tts.data.audio_utils import convert_audio
from pocket_tts.default_parameters import (
DEFAULT_EOS_THRESHOLD,
DEFAULT_LSD_DECODE_STEPS,
DEFAULT_NOISE_CLAMP,
DEFAULT_TEMPERATURE,
DEFAULT_VARIANT,
)
from pocket_tts.models.flow_lm import FlowLMModel
from pocket_tts.models.mimi import MimiModel
from pocket_tts.modules import mimi_transformer
from pocket_tts.modules.dummy_quantizer import DummyQuantizer
from pocket_tts.modules.seanet import SEANetDecoder, SEANetEncoder
from pocket_tts.modules.stateful_module import increment_steps, init_states
from pocket_tts.utils.config import Config, load_config
from pocket_tts.utils.utils import (
PREDEFINED_VOICES,
display_execution_time,
download_if_necessary,
load_predefined_voice,
size_of_dict,
)
from pocket_tts.utils.weights_loading import get_flow_lm_state_dict, get_mimi_state_dict
torch.set_num_threads(1)
logger = logging.getLogger(__name__)
KEYS_TO_MODIFY = [
"mimi.decoder.model.0.conv.conv.bias",
"mimi.decoder.model.0.conv.conv.weight",
"mimi.decoder.model.11.conv.conv.bias",
"mimi.decoder.model.11.conv.conv.weight",
"mimi.decoder.model.2.convtr.convtr.bias",
"mimi.decoder.model.2.convtr.convtr.weight",
"mimi.decoder.model.3.block.1.conv.conv.bias",
"mimi.decoder.model.3.block.1.conv.conv.weight",
"mimi.decoder.model.3.block.3.conv.conv.bias",
"mimi.decoder.model.3.block.3.conv.conv.weight",
"mimi.decoder.model.5.convtr.convtr.bias",
"mimi.decoder.model.5.convtr.convtr.weight",
"mimi.decoder.model.6.block.1.conv.conv.bias",
"mimi.decoder.model.6.block.1.conv.conv.weight",
"mimi.decoder.model.6.block.3.conv.conv.bias",
"mimi.decoder.model.6.block.3.conv.conv.weight",
"mimi.decoder.model.8.convtr.convtr.bias",
"mimi.decoder.model.8.convtr.convtr.weight",
"mimi.decoder.model.9.block.1.conv.conv.bias",
"mimi.decoder.model.9.block.1.conv.conv.weight",
"mimi.decoder.model.9.block.3.conv.conv.bias",
"mimi.decoder.model.9.block.3.conv.conv.weight",
"mimi.downsample.conv.conv.conv.weight",
"mimi.encoder.model.0.conv.conv.bias",
"mimi.encoder.model.0.conv.conv.weight",
"mimi.encoder.model.1.block.1.conv.conv.bias",
"mimi.encoder.model.1.block.1.conv.conv.weight",
"mimi.encoder.model.1.block.3.conv.conv.bias",
"mimi.encoder.model.1.block.3.conv.conv.weight",
"mimi.encoder.model.11.conv.conv.bias",
"mimi.encoder.model.11.conv.conv.weight",
"mimi.encoder.model.3.conv.conv.bias",
"mimi.encoder.model.3.conv.conv.weight",
"mimi.encoder.model.4.block.1.conv.conv.bias",
"mimi.encoder.model.4.block.1.conv.conv.weight",
"mimi.encoder.model.4.block.3.conv.conv.bias",
"mimi.encoder.model.4.block.3.conv.conv.weight",
"mimi.encoder.model.6.conv.conv.bias",
"mimi.encoder.model.6.conv.conv.weight",
"mimi.encoder.model.7.block.1.conv.conv.bias",
"mimi.encoder.model.7.block.1.conv.conv.weight",
"mimi.encoder.model.7.block.3.conv.conv.bias",
"mimi.encoder.model.7.block.3.conv.conv.weight",
"mimi.encoder.model.9.conv.conv.bias",
"mimi.encoder.model.9.conv.conv.weight",
"mimi.upsample.convtr.convtr.convtr.weight",
]
KEY_WITH_REMOVED_ZERO = [
"mimi.encoder_transformer.transformer.layers.0.self_attn.in_projs.0.weight",
"mimi.encoder_transformer.transformer.layers.0.self_attn.out_projs.0.weight",
"mimi.encoder_transformer.transformer.layers.1.self_attn.in_projs.0.weight",
"mimi.encoder_transformer.transformer.layers.1.self_attn.out_projs.0.weight",
"mimi.decoder_transformer.transformer.layers.0.self_attn.in_projs.0.weight",
"mimi.decoder_transformer.transformer.layers.0.self_attn.out_projs.0.weight",
"mimi.decoder_transformer.transformer.layers.1.self_attn.in_projs.0.weight",
"mimi.decoder_transformer.transformer.layers.1.self_attn.out_projs.0.weight",
]
def adapt_state_dict_of_TTSModel(state_dict: dict) -> dict:
# should be removed before release
new_state_dict = {}
for key, value in state_dict.items():
if key == "mimi.quantizer.input_proj.weight":
continue
if key in KEYS_TO_MODIFY:
new_key = key.replace(".conv.conv.", ".conv.").replace(".convtr.convtr.", ".convtr.")
new_state_dict[new_key] = value
elif key in KEY_WITH_REMOVED_ZERO:
new_key = key.replace("s.0.weight", ".weight")
new_state_dict[new_key] = value
else:
new_state_dict[key] = value
return new_state_dict
class TTSModel(nn.Module):
def __init__(
self,
flow_lm: FlowLMModel,
temp: float,
lsd_decode_steps: int,
noise_clamp: float | None,
eos_threshold,
config: Config,
):
super().__init__()
self.flow_lm = flow_lm
self.temp = temp
self.lsd_decode_steps = lsd_decode_steps
self.noise_clamp = noise_clamp
self.eos_threshold = eos_threshold
self.config = config
self.has_voice_cloning = True
@property
def device(self) -> str:
return next(self.parameters()).device.type
@property
def sample_rate(self) -> int:
return self.config.mimi.sample_rate
@classmethod
def _from_pydantic_config(
cls, config: Config, temp, lsd_decode_steps, noise_clamp: float | None, eos_threshold
) -> Self:
flow_lm = FlowLMModel.from_pydantic_config(
config.flow_lm, latent_dim=config.mimi.quantizer.dimension
)
tts_model = cls(flow_lm, temp, lsd_decode_steps, noise_clamp, eos_threshold, config)
return tts_model
@classmethod
def _from_pydantic_config_with_weights(
cls, config: Config, temp, lsd_decode_steps, noise_clamp: float | None, eos_threshold
) -> Self:
tts_model = cls._from_pydantic_config(
config, temp, lsd_decode_steps, noise_clamp, eos_threshold
)
tts_model.flow_lm.speaker_proj_weight = torch.nn.Parameter(
torch.zeros((1024, 512), dtype=torch.float32)
)
if config.flow_lm.weights_path is not None:
if config.mimi.weights_path is None:
raise ValueError(
"If you specify flow_lm.weights_path you should specify mimi.weights_path"
)
logger.info(f"Loading FlowLM weights from {config.flow_lm.weights_path}")
state_dict_flowlm = get_flow_lm_state_dict(
download_if_necessary(config.flow_lm.weights_path)
)
tts_model.flow_lm.load_state_dict(state_dict_flowlm, strict=True)
# safetensors.torch.save_file(tts_model.state_dict(), "7442637a.safetensors")
# Create mimi config directly from the provided config using model_dump
mimi_config = config.mimi.model_dump()
# Build mimi model from config
encoder = SEANetEncoder(**mimi_config["seanet"])
decoder = SEANetDecoder(**mimi_config["seanet"])
encoder_transformer = mimi_transformer.ProjectedTransformer(**mimi_config["transformer"])
decoder_transformer = mimi_transformer.ProjectedTransformer(**mimi_config["transformer"])
quantizer = DummyQuantizer(**mimi_config["quantizer"])
tts_model.mimi = MimiModel(
encoder,
decoder,
quantizer,
channels=mimi_config["channels"],
sample_rate=mimi_config["sample_rate"],
frame_rate=mimi_config["frame_rate"],
encoder_frame_rate=mimi_config["sample_rate"] / encoder.hop_length,
encoder_transformer=encoder_transformer,
decoder_transformer=decoder_transformer,
).to(device="cpu")
# Load mimi weights from the config safetensors file with complete mapping for strict loading
if config.mimi.weights_path is not None:
if config.flow_lm.weights_path is None:
raise ValueError(
"If you specify mimi.weights_path you should specify flow_lm.weights_path"
)
logger.info(f"Loading Mimi weights from {config.mimi.weights_path}")
mimi_state = get_mimi_state_dict(download_if_necessary(config.mimi.weights_path))
tts_model.mimi.load_state_dict(mimi_state, strict=True)
tts_model.mimi.eval()
# tts_model.to(dtype=torch.float32)
# uncomment to save the weights
# tts_model = tts_model.to(dtype=torch.bfloat16)
# safetensors.torch.save_file(tts_model.state_dict(), "tts_b6369a24.safetensors")
if config.weights_path is not None:
logger.info(f"Loading TTSModel weights from {config.weights_path}")
try:
weights_file = download_if_necessary(config.weights_path)
except Exception:
tts_model.has_voice_cloning = False
weights_file = download_if_necessary(config.weights_path_without_voice_cloning)
state_dict = safetensors.torch.load_file(weights_file)
state_dict = adapt_state_dict_of_TTSModel(state_dict)
tts_model.load_state_dict(state_dict, strict=True)
if config.flow_lm.weights_path is None and config.weights_path is None:
logger.warning(
"No weights_path specified for FlowLM or TTSModel, model is uninitialized!"
)
size_in_mb = size_of_dict(tts_model.state_dict()) // 1e6
logging.info(f"TTS Model loaded successfully. Its size is {size_in_mb} MB")
return tts_model
def load_model(
variant: str = DEFAULT_VARIANT,
temp: float | int = DEFAULT_TEMPERATURE,
lsd_decode_steps: int = DEFAULT_LSD_DECODE_STEPS,
noise_clamp: float | int | None = DEFAULT_NOISE_CLAMP,
eos_threshold: float = DEFAULT_EOS_THRESHOLD,
) -> Self:
"""Load a pre-trained TTS model with specified configuration.
This class method loads a complete TTS model including the flow language model
and Mimi compression model from pre-trained weights. The model is initialized
with the specified generation parameters and ready for inference.
Args:
variant: Model variant identifier corresponding to a config file name
(e.g., '610b0b2c'). Must match a YAML file in the config directory.
temp: Sampling temperature for generation. Higher values produce more
diverse but potentially lower quality output.
lsd_decode_steps: Number of steps for Lagrangian Self Distillation
decoding. More steps can improve quality but increase computation.
noise_clamp: Maximum value for noise sampling. If None, no clamping
is applied. Helps prevent extreme values in generation.
eos_threshold: Threshold for end-of-sequence detection. Higher values
make the model more likely to continue generating.
Returns:
TTSModel: Fully initialized model with loaded weights on cpu, ready for
text-to-speech generation.
Raises:
FileNotFoundError: If the specified config file or model weights
are not found.
ValueError: If the configuration is invalid or incompatible.
"""
config = load_config(Path(__file__).parents[1] / f"config/{variant}.yaml")
tts_model = TTSModel._from_pydantic_config_with_weights(
config, temp, lsd_decode_steps, noise_clamp, eos_threshold
)
return tts_model
def _run_flow_lm_and_increment_step(
self,
model_state: dict,
text_tokens: torch.Tensor | None = None,
backbone_input_latents: torch.Tensor | None = None,
audio_conditioning: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""First one is the backbone output, second one is the audio decoding output."""
if text_tokens is None:
text_tokens = torch.zeros((1, 0), dtype=torch.int64, device=self.flow_lm.device)
if backbone_input_latents is None:
backbone_input_latents = torch.empty(
(1, 0, self.flow_lm.ldim), dtype=self.flow_lm.dtype, device=self.flow_lm.device
)
if audio_conditioning is None:
audio_conditioning = torch.empty(
(1, 0, self.flow_lm.dim), dtype=self.flow_lm.dtype, device=self.flow_lm.device
)
output = self._run_flow_lm(
text_tokens=text_tokens,
backbone_input_latents=backbone_input_latents,
model_state=model_state,
audio_conditioning=audio_conditioning,
)
increment_by = (
text_tokens.shape[1] + backbone_input_latents.shape[1] + audio_conditioning.shape[1]
)
increment_steps(self.flow_lm, model_state, increment=increment_by)
return output
def _run_flow_lm(
self,
model_state: dict,
text_tokens: torch.Tensor,
backbone_input_latents: torch.Tensor,
audio_conditioning: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
text_embeddings = self.flow_lm.conditioner(TokenizedText(text_tokens))
text_embeddings = torch.cat([text_embeddings, audio_conditioning], dim=1)
output_embeddings, is_eos = self.flow_lm._sample_next_latent(
backbone_input_latents,
text_embeddings,
model_state=model_state,
lsd_decode_steps=self.lsd_decode_steps,
temp=self.temp,
noise_clamp=self.noise_clamp,
eos_threshold=self.eos_threshold,
)
return output_embeddings[:, None, :], is_eos
def _encode_audio(self, audio: torch.Tensor) -> torch.Tensor:
encoded = self.mimi.encode_to_latent(audio)
latents = encoded.transpose(-1, -2).to(torch.float32)
conditioning = F.linear(latents, self.flow_lm.speaker_proj_weight)
return conditioning
def _decode_audio_worker(self, latents_queue: queue.Queue, result_queue: queue.Queue):
"""Worker thread function for decoding audio latents from queue with immediate streaming."""
try:
audio_chunks = []
mimi_state = init_states(self.mimi, batch_size=1, sequence_length=1000)
while True:
latent = latents_queue.get()
if latent is None:
break
mimi_decoding_input = latent * self.flow_lm.emb_std + self.flow_lm.emb_mean
transposed = mimi_decoding_input.transpose(-1, -2)
quantized = self.mimi.quantizer(transposed)
t = time.monotonic()
audio_frame = self.mimi.decode_from_latent(quantized, mimi_state)
increment_steps(self.mimi, mimi_state, increment=16)
audio_frame_duration = audio_frame.shape[2] / self.config.mimi.sample_rate
# We could log the timings here.
logger.debug(
" " * 30 + "Decoded %d ms of audio with mimi in %d ms",
int(audio_frame_duration * 1000),
int((time.monotonic() - t) * 1000),
)
audio_chunks.append(audio_frame)
result_queue.put(("chunk", audio_frame))
latents_queue.task_done()
# Signal completion
result_queue.put(("done", None))
except Exception as e:
# Put error in result queue
result_queue.put(("error", e))
@torch.no_grad
def generate_audio(
self,
model_state: dict,
text_to_generate: str,
frames_after_eos: int | None = None,
copy_state: bool = True,
) -> torch.Tensor:
"""Generate complete audio tensor from text input.
This method generates the full audio output for the given text prompt
and returns it as a single tensor. It internally uses the streaming
generation method but collects all chunks before returning.
This method is NOT thread-safe; separate model instances should be used
for concurrent generation.
Args:
model_state: Model state dictionary containing hidden states and
positional information. Can be obtained from get_state_for_audio_prompt()
or init_states(). The state may be modified during generation.
text_to_generate: Input text to convert to speech. The text will be
automatically formatted (capitalization, punctuation) for optimal
generation quality.
frames_after_eos: Number of additional frames to generate after
detecting end-of-sequence. If None, automatically determined
based on text length (1-3 frames).
copy_state: Whether to create a deep copy of the model state before
generation. If True, preserves the original state for reuse.
If False, modifies the input state in-place. Defaults to True.
Returns:
torch.Tensor: Generated audio tensor with shape [channels, samples]
at the model's sample rate (typically 24kHz). The audio is
normalized and ready for playback or saving.
You can get the sample rate from the `sample_rate` attribute.
Raises:
ValueError: If text_to_generate is empty or invalid.
RuntimeError: If generation fails due to model errors.
"""
audio_chunks = []
for chunk in self.generate_audio_stream(
model_state=model_state,
text_to_generate=text_to_generate,
frames_after_eos=frames_after_eos,
copy_state=copy_state,
):
audio_chunks.append(chunk)
return torch.cat(audio_chunks, dim=0)
@torch.no_grad
def generate_audio_stream(
self,
model_state: dict,
text_to_generate: str,
frames_after_eos: int | None = None,
copy_state: bool = True,
):
"""Generate audio streaming chunks from text input.
This method generates audio from text and yields chunks as they become
available, enabling real-time playback or processing. It uses multithreading
to parallelize generation and decoding for optimal performance.
This method is NOT thread-safe; separate model instances should be used
for concurrent generation.
Args:
model_state: Model state dictionary containing hidden states and
positional information. Can be obtained from get_state_for_audio_prompt()
or init_states(). The state may be modified during generation.
text_to_generate: Input text to convert to speech. The text will be
automatically formatted (capitalization, punctuation) for optimal
generation quality.
frames_after_eos: Number of additional frames to generate after
detecting end-of-sequence. If None, automatically determined
based on text length (1-3 frames). Defaults to None.
copy_state: Whether to create a deep copy of the model state before
generation. If True, preserves the original state for reuse.
If False, modifies the input state in-place. Defaults to True.
Yields:
torch.Tensor: Audio chunks with shape [samples] at the model's
sample rate (typically 24kHz). Chunks are yielded as soon as
they are decoded, enabling real-time streaming.
Raises:
ValueError: If text_to_generate is empty or invalid.
RuntimeError: If generation fails due to model errors or threading issues.
Note:
This method uses multithreading to parallelize latent generation
and audio decoding. Generation performance is logged including
real-time factor (RTF) metrics.
"""
text_to_generate, frames_after_eos_guess = prepare_text_prompt(text_to_generate)
if frames_after_eos is None:
frames_after_eos = frames_after_eos_guess
if copy_state:
model_state = copy.deepcopy(model_state)
# Set up multithreaded generation and decoding
latents_queue = queue.Queue()
result_queue = queue.Queue()
# Start decoder worker thread
decoder_thread = threading.Thread(
target=self._decode_audio_worker, args=(latents_queue, result_queue), daemon=True
)
logger.info("starting timer now!")
t_generating = time.monotonic()
decoder_thread.start()
# Generate latents and add them to queue (decoder processes them in parallel)
self._generate(
model_state=model_state,
text_to_generate=text_to_generate,
frames_after_eos=frames_after_eos,
latents_queue=latents_queue,
result_queue=result_queue,
)
# Stream audio chunks as they become available
total_generated_samples = 0
while True:
result = result_queue.get()
if result[0] == "chunk":
# Audio chunk available immediately for streaming/playback
audio_chunk = result[1]
total_generated_samples += audio_chunk.shape[-1]
yield audio_chunk[0, 0] # Remove batch, channel
elif result[0] == "done":
# Generation complete
break
elif result[0] == "error":
# Wait for decoder thread to finish cleanly before propagating error
with display_execution_time("Waiting for mimi decoder to finish"):
decoder_thread.join()
# Propagate error
raise result[1]
# Wait for decoder thread to finish cleanly
with display_execution_time("Waiting for mimi decoder to finish"):
decoder_thread.join()
# Print timing information
duration_generated_audio = int(
total_generated_samples * 1000 / self.config.mimi.sample_rate
)
generation_time = int((time.monotonic() - t_generating) * 1000)
real_time_factor = duration_generated_audio / generation_time
logger.info(
"Generated: %d ms of audio in %d ms so %.2fx faster than real-time",
duration_generated_audio,
generation_time,
real_time_factor,
)
@torch.no_grad
def _generate(
self,
model_state: dict,
text_to_generate: str,
frames_after_eos: int,
latents_queue: queue.Queue,
result_queue: queue.Queue,
):
gen_len_sec = len(text_to_generate.split()) * 1 + 2.0
max_gen_len = int(gen_len_sec * 12.5)
prepared = self.flow_lm.conditioner.prepare(text_to_generate)
with display_execution_time("Prompting text"):
self._run_flow_lm_and_increment_step(
model_state=model_state, text_tokens=prepared.tokens
)
def run_generation():
try:
self._autoregressive_generation(
model_state, max_gen_len, frames_after_eos, latents_queue
)
except Exception as e:
logger.error(f"Error in autoregressive generation: {e}")
# Signal decoder to stop by putting None (completion sentinel)
if latents_queue is not None:
latents_queue.put(None)
# Report error to main thread
if result_queue is not None:
result_queue.put(("error", e))
generation_thread = threading.Thread(target=run_generation, daemon=True)
generation_thread.start()
def _autoregressive_generation(
self, model_state: dict, max_gen_len: int, frames_after_eos: int, latents_queue: queue.Queue
):
backbone_input = torch.full(
(1, 1, self.flow_lm.ldim),
fill_value=float("NaN"),
device=next(iter(self.flow_lm.parameters())).device,
dtype=self.flow_lm.dtype,
)
steps_times = []
eos_step = None
for generation_step in range(max_gen_len):
with display_execution_time("Generating latent", print_output=False) as timer:
next_latent, is_eos = self._run_flow_lm_and_increment_step(
model_state=model_state, backbone_input_latents=backbone_input
)
if is_eos.item() and eos_step is None:
eos_step = generation_step
if eos_step is not None and generation_step >= eos_step + frames_after_eos:
break
# Add generated latent to queue for immediate decoding
latents_queue.put(next_latent)
backbone_input = next_latent
steps_times.append(timer.elapsed_time_ms)
else:
if os.environ.get("KPOCKET_TTS_ERROR_WITHOUT_EOS", "0") == "1":
raise RuntimeError("Generation reached maximum length without EOS!")
logger.warning(
"Maximum generation length reached without EOS, this very often indicates an error."
)
# Add sentinel value to signal end of generation
latents_queue.put(None)
logger.info("Average generation step time: %d ms", int(statistics.mean(steps_times)))
@lru_cache(maxsize=2)
def _cached_get_state_for_audio_prompt(
self, audio_conditioning: Path | str | torch.Tensor, truncate: bool = False
) -> dict:
return self.get_state_for_audio_prompt(audio_conditioning, truncate)
@torch.no_grad
def get_state_for_audio_prompt(
self, audio_conditioning: Path | str | torch.Tensor, truncate: bool = False
) -> dict:
"""Create model state conditioned on audio prompt for continuation.
This method processes an audio prompt and creates a model state that
captures the acoustic characteristics (speaker voice, style, prosody)
for use in subsequent text-to-speech generation. The resulting state
enables voice cloning and audio continuation with speaker consistency.
Args:
audio_conditioning: Audio prompt to condition on. Can be:
- Path: Local file path to audio file
- str: URL to download audio file from
- torch.Tensor: Pre-loaded audio tensor with shape [channels, samples]
truncate: Whether to truncate long audio prompts to 30 seconds.
Helps prevent memory issues with very long inputs. Defaults to False.
Returns:
dict: Model state dictionary containing hidden states and positional
information conditioned on the audio prompt. This state can be
passed to `generate_audio()` or `generate_audio_stream()` for
voice-consistent generation.
Raises:
FileNotFoundError: If audio file path doesn't exist.
ValueError: If audio tensor is invalid or empty.
RuntimeError: If audio processing or encoding fails.
Note:
- Audio is automatically resampled to the model's sample rate (24kHz)
- The audio is encoded using the Mimi compression model and projected
to the flow model's latent space
- Processing time is logged for performance monitoring
- The state preserves speaker characteristics for voice cloning
"""
if isinstance(audio_conditioning, str) and audio_conditioning in PREDEFINED_VOICES:
# We get the audio conditioning directly from the safetensors file.
prompt = load_predefined_voice(audio_conditioning)
else:
if not self.has_voice_cloning and isinstance(audio_conditioning, (str, Path)):
raise ValueError(
f"We could not download the weights for the model with voice cloning, "
f"but you're trying to use voice cloning. "
f"Without voice cloning, you can use our catalog of voices {list(PREDEFINED_VOICES)}. "
f"If you want access to the model with voice cloning, go to "
f"https://huggingface.co/kyutai/pocket-tts and accept the terms, "
f"then make sure you're logged in locally with `uvx hf auth login`."
)
if isinstance(audio_conditioning, str):
audio_conditioning = download_if_necessary(audio_conditioning)
if isinstance(audio_conditioning, Path):
audio, conditioning_sample_rate = audio_read(audio_conditioning)
if truncate:
max_samples = int(30 * conditioning_sample_rate) # 30 seconds of audio
if audio.shape[-1] > max_samples:
audio = audio[..., :max_samples]
logger.info(f"Audio truncated to first 30 seconds ({max_samples} samples)")
audio_conditioning = convert_audio(
audio, conditioning_sample_rate, self.config.mimi.sample_rate, 1
)
with display_execution_time("Encoding audio prompt"):
prompt = self._encode_audio(audio_conditioning.unsqueeze(0).to(self.device))
# import safetensors.torch
# safetensors.torch.save_file({"audio_prompt": prompt}, "embeddings/eponine.safetensors")
model_state = init_states(self.flow_lm, batch_size=1, sequence_length=1000)
with display_execution_time("Prompting audio"):
self._run_flow_lm_and_increment_step(model_state=model_state, audio_conditioning=prompt)
return model_state
def prepare_text_prompt(text: str) -> tuple[str, int]:
text = text.strip()
if text == "":
raise ValueError("Text prompt cannot be empty")
number_of_words = len(text.split())
if number_of_words <= 4:
frames_after_eos_guess = 3
else:
frames_after_eos_guess = 1
# Make sure it starts with an uppercase letter
if not text[0].isupper():
text = text[0].upper() + text[1:]
# Let's make sure it ends with some kind of punctuation
# If it ends with a letter or digit, we add a period.
if text[-1].isalnum():
text = text + "."
# The model does not perform well when there are very few tokens, so
# we can add empty spaces at the beginning to increase the token count.
if len(text.split()) < 5:
text = " " * 8 + text
return text, frames_after_eos_guess
Reference in New Issue View Git Blame Copy Permalink