valenti/spark-pocket-tts
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add .safetensors save/load to get_state_for_audio_prompt (#69)

Browse Source 
Quality of life improvement to save user voice embedding as a
.safetensors file that can be loaded for subsequent use, achieving the
same efficiency as the predefined voices, all via the existing
get_state_for_audio_prompt function.

---------

Co-authored-by: Václav Volhejn <8401624+vvolhejn@users.noreply.github.com>
This commit is contained in:
thataboy
2026-01-26 10:36:27 +01:00
committed by GitHub
co-authored by Václav Volhejn
parent bfbaf80437
commit edec7f8790
6 changed files with 290 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pocket_tts/main.py
+110
View File
@@ -263,5 +263,115 @@ def generate(
)
# ----------------------------------------------
# export audio to safetensors CLI implementation
# ----------------------------------------------
@cli_app.command()
def export_voice(
audio_path: Annotated[
str, typer.Argument(help="Audio file or directory to convert and export")
],
export_path: Annotated[str, typer.Argument(help="Output file or directory")],
truncate: Annotated[
bool, typer.Option("-tr", "--truncate", help="Truncate long audio")
] = False,
quiet: Annotated[bool, typer.Option("-q", "--quiet", help="Disable logging output")] = False,
variant: Annotated[str, typer.Option(help="Model signature")] = DEFAULT_VARIANT,
lsd_decode_steps: Annotated[
int, typer.Option(help="Number of generation steps")
] = DEFAULT_LSD_DECODE_STEPS,
temperature: Annotated[
float, typer.Option(help="Temperature for generation")
] = DEFAULT_TEMPERATURE,
noise_clamp: Annotated[float, typer.Option(help="Noise clamp value")] = DEFAULT_NOISE_CLAMP,
eos_threshold: Annotated[float, typer.Option(help="EOS threshold")] = DEFAULT_EOS_THRESHOLD,
frames_after_eos: Annotated[
int, typer.Option(help="Number of frames to generate after EOS")
] = DEFAULT_FRAMES_AFTER_EOS,
device: Annotated[str, typer.Option(help="Device to use")] = "cpu",
):
"""Convert and save audio to .safetensors file"""
import re
def url(path):
return path.startswith(("http:", "https:", "hf:"))
def normalize_url(url):
# utils.py expects urls to be xxx:// so normalize them
return re.sub(r"^(http|https|hf)\:\/*(.+)$", r"\1://\2", url)
def likely_file(path):
return not url(path) and not likely_dir(path)
def likely_dir(path):
return not url(path) and (path.endswith(("/", "\\")) or path == ".")
def convert_one(in_path, out_path, join_path):
"""helper convert function"""
voice = in_path.stem
if url(str(in_path)):
in_path = normalize_url(str(in_path))
if join_path:
out_path = out_path / f"{voice}.safetensors"
else:
# ensure output file has correct extension
out_path = out_path.with_suffix(".safetensors")
try:
tts_model.save_audio_prompt(in_path, out_path, truncate)
except Exception as e:
logger.error(f"❌ Unable to export voice '{in_path}': {e}")
return False
logger.info(f"✅ Successfully exported voice '{voice}' to '{out_path}'")
return True
if "cuda" in device:
# Cuda graphs capturing does not play nice with multithreading.
os.environ["NO_CUDA_GRAPH"] = "1"
log_level = logging.ERROR if quiet else logging.INFO
success_count = 0
with enable_logging("pocket_tts", log_level):
tts_model = TTSModel.load_model(
variant, temperature, lsd_decode_steps, noise_clamp, eos_threshold
)
tts_model.to(device)
in_path = Path(audio_path)
out_path = Path(export_path)
if likely_dir(export_path):
# make sure output dir exists
out_path.mkdir(parents=True, exist_ok=True)
if likely_dir(audio_path): # batch convert whole directory
if not in_path.is_dir():
logger.error(f"Input dir '{audio_path}' does not exists")
exit(1)
if not likely_dir(export_path):
# batch convert, output path must be directory, not file
out_path = Path("./")
for path in Path(in_path).iterdir():
if path.is_file() and path.suffix.lower() in [
".wav",
".mp3",
".flac",
".ogg",
".aiff",
]:
if convert_one(path, out_path, True):
success_count += 1
else: # convert single file
if likely_file(audio_path) and not in_path.exists():
logger.error(f"Input file '{in_path}'' does not exists")
exit(1)
if convert_one(in_path, out_path, likely_dir(export_path)):
success_count += 1
if success_count > 0:
logger.info(f"🎉 Successfully exported {success_count} voices.")
if __name__ == "__main__":
cli_app()
pocket_tts/models/tts_model.py
+89 -17
View File
@@ -44,6 +44,15 @@ from pocket_tts.utils.weights_loading import get_flow_lm_state_dict, get_mimi_st
torch.set_num_threads(1)
logger = logging.getLogger(__name__)
VOICE_CLONING_UNSUPPORTED = (
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`."
)
class TTSModel(nn.Module):
def __init__(
@@ -637,9 +646,9 @@ class TTSModel(nn.Module):
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
audio_conditioning: Audio prompt to condition (or .safetensors to load). Can be:
- Path: Local file path to audio file (or .safetensors)
- str: URL to download audio file (or .safetensors) 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.
@@ -662,19 +671,21 @@ class TTSModel(nn.Module):
- 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:
if isinstance(audio_conditioning, (str, Path)) and str(audio_conditioning).endswith(
".safetensors"
):
if isinstance(audio_conditioning, str):
audio_conditioning = download_if_necessary(audio_conditioning)
import safetensors.torch
prompt = safetensors.torch.load_file(audio_conditioning)["audio_prompt"]
elif 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`."
)
raise ValueError(VOICE_CLONING_UNSUPPORTED)
if isinstance(audio_conditioning, str):
audio_conditioning = download_if_necessary(audio_conditioning)
@@ -693,11 +704,6 @@ class TTSModel(nn.Module):
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},
# "/projects/huggingface/pocket-tts/embeddings/cosette.safetensors"
# )
model_state = init_states(self.flow_lm, batch_size=1, sequence_length=1000)
@@ -710,6 +716,72 @@ class TTSModel(nn.Module):
return model_state
@torch.no_grad
def save_audio_prompt(
self,
audio_conditioning: Path | str | torch.Tensor,
export_path: Path | str,
truncate: bool = False,
) -> torch.Tensor:
"""Save audio prompt to .safetensors file
This method processes an audio prompt and exports it to a .safetensors file,
which can be loaded by get_state_for_audio_prompt in subsequent uses
without converting the audio again.
It also takes an already converted audio tensor and exports it as a .safetensors file
Args:
audio_conditioning: Audio to export
- Path: Local file path to audio file
- str: URL to download audio file
- torch.Tensor: Pre-loaded audio tensor with shape [channels, samples]
export_path: Path to output file
truncate: Whether to truncate long audio prompts to 30 seconds.
Returns:
Audio tensor of converted audio
Raises:
FileNotFoundError: If audio file path doesn't exist.
ValueError: If audio tensor export path is invalid or empty.
RuntimeError: If audio processing or encoding fails.
Note:
- Send resulting audio tensor to get_state_for_audio_prompt
in order to get the model state for generation.
"""
if not export_path or not isinstance(export_path, (str, Path)):
raise ValueError("export_path must be of type str or Path")
export_path = Path(export_path).with_suffix(".safetensors")
if not self.has_voice_cloning and isinstance(audio_conditioning, (str, Path)):
raise ValueError(VOICE_CLONING_UNSUPPORTED)
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("Exporting audio prompt"):
prompt = self._encode_audio(audio_conditioning.unsqueeze(0).to(self.device))
import safetensors.torch
safetensors.torch.save_file({"audio_prompt": prompt}, export_path)
return audio_conditioning
def prepare_text_prompt(text: str) -> tuple[str, int]:
text = text.strip()