82 lines
2.7 KiB
Python
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()
|