tts-dgx-spark-faster-qwen3-tts/config/extract_embeddings.py

82 lines
2.7 KiB
Python

import os
import json
import torch
import sys
# Append the app directory to import faster_qwen3_tts
sys.path.append("/app/examples")
sys.path.insert(0, "/app")
from faster_qwen3_tts.model import FasterQwen3TTS
def main():
config_file = "/config/voices.json"
if not os.path.exists(config_file):
print(f"Error: {config_file} not found.")
return
with open(config_file, "r") as f:
voices = json.load(f)
# Need to load the model
# Read model path from QWEN_TTS_MODEL or default
model_path = os.environ.get("QWEN_TTS_MODEL", "Qwen/Qwen3-TTS-12Hz-1.7B-Base")
print(f"Loading model {model_path} for embedding extraction...")
tts_model = FasterQwen3TTS.from_pretrained(
model_path,
device="cuda",
dtype=torch.bfloat16,
)
print("Model loaded successfully.")
updates_made = False
for voice_id, entry in voices.items():
spk_emb_path = entry.get("speaker_embeddings") or entry.get("speaker embeddings")
# If there's already a valid path and it exists, skip
if spk_emb_path and os.path.exists(spk_emb_path):
print(f"Skipping {voice_id}, embedding already exists at {spk_emb_path}")
continue
ref_audio = entry.get("ref_audio")
ref_text = entry.get("ref_text", "")
if not ref_audio or not os.path.exists(ref_audio):
print(f"Skipping {voice_id}, ref_audio not found: {ref_audio}")
continue
print(f"Extracting embeddings for {voice_id}...")
try:
prompt_items = tts_model.model.create_voice_clone_prompt(
ref_audio, [ref_text]
)
vcp = tts_model.model._prompt_items_to_voice_clone_prompt(prompt_items)
# Save the .pt file in /config/speakers
pt_path = f"/config/speakers/{voice_id}.pt"
torch.save(vcp, pt_path)
print(f"Saved {pt_path}")
# Update voices.json entry
entry["speaker_embeddings"] = pt_path
# Remove the legacy "speaker embeddings" with space if it exists
if "speaker embeddings" in entry:
del entry["speaker embeddings"]
updates_made = True
except Exception as e:
print(f"Error extracting embedding for {voice_id}: {e}")
if updates_made:
with open(config_file, "w") as f:
json.dump(voices, f, indent=2, ensure_ascii=False)
print("Updated voices.json with new speaker_embeddings paths.")
else:
print("No new embeddings extracted.")
if __name__ == "__main__":
main()