Spaces:

Alstears
/

chatterbox-id-clone-api

Running

Update app.py

a6bcea0 verified 6 months ago

12.2 kB

import os
# Paksa PyTorch murni berjalan di CPU sebelum library lain masuk
os.environ["CUDA_VISIBLE_DEVICES"] = ""

import re
import gc
import pickle
import inspect
import tempfile
import traceback
from pathlib import Path
from threading import Lock

import requests
import torch
import torchaudio as ta
import gradio as gr

# Optimasi Core CPU shared Hugging Face agar fokus dan stabil
torch.set_num_threads(2)
torch.set_num_interop_threads(2)
torch.set_grad_enabled(False)
torch.cuda.is_available = lambda: False  # noqa: E731

# Override global torch load untuk keamanan jalur CPU
_original_torch_load = torch.load
def _torch_load_cpu(*args, **kwargs):
    kwargs["map_location"] = torch.device("cpu")
    return _original_torch_load(*args, **kwargs)
torch.load = _torch_load_cpu

# =====================================================================
# CONFIG & PATH MANAGEMENT (SWEET SPOT PARAMETERS)
# =====================================================================
MODEL_REPO = "grandhigh/Chatterbox-TTS-Indonesian"
CHECKPOINT_FILENAME = "t3_cfg.safetensors"
DEVICE = "cpu"

MAX_TOTAL_CHARS = int(os.getenv("MAX_TOTAL_CHARS", "2400"))
# Set ke 280 agar CPU ringan melakukan sampling awal dan tidak stuck di 0%
MAX_CHARS_PER_CHUNK = int(os.getenv("MAX_CHARS_PER_CHUNK", "280"))
MAX_CHUNKS = int(os.getenv("MAX_CHUNKS", "10"))
PAUSE_SECONDS = float(os.getenv("PAUSE_SECONDS", "0.0"))
DOWNLOAD_TIMEOUT = int(os.getenv("DOWNLOAD_TIMEOUT", "90"))

ROOT_DIR = Path(__file__).parent
DEFAULT_SPEAKER_PATH = ROOT_DIR / "default_speaker.wav"

# FILE BOBOT PERMANEN: Lokasi penyimpanan biner karakteristik suara kustom kamu
EMBEDDING_CACHE_PATH = ROOT_DIR / "speaker_embedding.pth"

# Global RAM Cache untuk menyimpan koordinat suara siap pakai
CACHED_EMBEDDINGS = {}
_model = None
_model_lock = Lock()


def inisialisasi_bobot_default_permanen(model_instance):
    """
    Fungsi untuk melatih/mengekstrak karakteristik file default_speaker.wav
    HANYA SEKALI saja saat startup. Hasilnya disimpan permanen ke disk (.pth).
    """
    global CACHED_EMBEDDINGS
    try:
        # Jika file biner .pth sudah ada, langsung muat ke memori RAM secara instan
        if EMBEDDING_CACHE_PATH.exists():
            print("[INIT] Menemukan file bobot permanen speaker_embedding.pth. Memuat ke RAM...")
            with open(EMBEDDING_CACHE_PATH, "rb") as f:
                CACHED_EMBEDDINGS[str(DEFAULT_SPEAKER_PATH)] = pickle.load(f)
            print("[INIT] Bobot kustom bawaan siap digunakan instan!")
            return

        # Jika belum ada file .pth tapi ada file .wav asli, lakukan ekstraksi awal
        if DEFAULT_SPEAKER_PATH.exists():
            print("[INIT] Membuat file bobot baru. Mengekstrak default_speaker.wav...")
            if hasattr(model_instance, "extract_conditioning"):
                embedding = model_instance.extract_conditioning(str(DEFAULT_SPEAKER_PATH))
            elif hasattr(model_instance, "get_speaker_embedding"):
                embedding = model_instance.get_speaker_embedding(str(DEFAULT_SPEAKER_PATH))
            else:
                embedding = str(DEFAULT_SPEAKER_PATH)

            # Simpan secara fisik ke disk tingkat root agar permanen
            if not isinstance(embedding, str):
                with open(EMBEDDING_CACHE_PATH, "wb") as f:
                    pickle.dump(embedding, f)
                print("[INIT] File speaker_embedding.pth berhasil disimpan secara permanen!")
            
            CACHED_EMBEDDINGS[str(DEFAULT_SPEAKER_PATH)] = embedding
    except Exception as e:
        print(f"[WARN] Gagal inisialisasi pembentukan bobot permanen awal: {e}")


def get_model():
    global _model
    if _model is None:
        with _model_lock:
            if _model is None:
                print("[INIT] Loading model on CPU...")
                m = ChatterboxTTS.from_pretrained(device=DEVICE)

                ckpt_path = hf_hub_download(
                    repo_id=MODEL_REPO,
                    filename=CHECKPOINT_FILENAME
                )
                t3_state = load_file(ckpt_path, device="cpu")
                m.t3.load_state_dict(t3_state)

                if hasattr(m, "eval"):
                    m.eval()

                _model = m
                print("[INIT] Model ready.")
                
                # JALANKAN PROSES TRAINING/EKSTRAKSI SEKALI SAAT STARTUP
                inisialisasi_bobot_default_permanen(_model)
    return _model


def _resolve_audio_input(audio_file, audio_url: str):
    if isinstance(audio_file, str) and audio_file.strip():
        return audio_file
    if isinstance(audio_file, dict):
        p = audio_file.get("path")
        if p:
            return p
    if audio_url and audio_url.strip():
        try:
            r = requests.get(audio_url.strip(), timeout=DOWNLOAD_TIMEOUT)
            r.raise_for_status()
            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
            tmp.write(r.content)
            tmp.close()
            return tmp.name
        except Exception:
            pass
            
    if DEFAULT_SPEAKER_PATH.exists():
        return str(DEFAULT_SPEAKER_PATH)
        
    return None


def _split_text_safely(text: str, max_chars: int = MAX_CHARS_PER_CHUNK):
    text = re.sub(r"\s+", " ", (text or "").strip())
    if not text:
        return []
    sentences = re.split(r"(?<=[.!?])\s+", text)
    chunks = []
    current = ""
    for s in sentences:
        s = s.strip()
        if not s:
            continue
        parts = [s] if len(s) <= max_chars else re.split(r"(?<=[,;:])\s+", s)
        for p in parts:
            p = p.strip()
            if not p:
                continue
            if len(p) > max_chars:
                words = p.split()
                tmp = ""
                for w in words:
                    cand = f"{tmp} {w}".strip() if tmp else w
                    if len(cand) <= max_chars:
                        tmp = cand
                    else:
                        if tmp:
                            chunks.append(tmp)
                        tmp = w
                if tmp:
                    chunks.append(tmp)
                continue
            candidate = f"{current} {p}".strip() if current else p
            if len(candidate) <= max_chars:
                current = candidate
            else:
                if current:
                    chunks.append(current)
                current = p
    if current:
        chunks.append(current)
    return chunks

# =====================================================================
# ENGINE UTAMA (FAST INFERENCE - INSTANT EMBEDDING LOADING)
# =====================================================================
def clone_voice(text: str, audio_file, audio_url: str, progress=gr.Progress(track_tqdm=False)):
    global CACHED_EMBEDDINGS
    try:
        raw_text = (text or "").strip()
        if not raw_text:
            raise gr.Error("Text prompt tidak boleh kosong.")

        prompt_path = _resolve_audio_input(audio_file, audio_url)
        if not prompt_path:
            raise gr.Error("Suara acuan default_speaker.wav tidak ditemukan di root server.")

        chunks = _split_text_safely(raw_text, max_chars=MAX_CHARS_PER_CHUNK)
        chunks = [ch for ch in chunks if re.search(r'[a-zA-Z0-9]', ch)]

        model = get_model()
        sr = getattr(model, "sr", 24000)
        torch.manual_seed(42)

        # Caching dinamis sebagai pengaman cadangan untuk Opsi Custom Voice Baru
        if prompt_path not in CACHED_EMBEDDINGS:
            progress(0.0, desc="Mengekstrak karakteristik gelombang audio kustom baru ke RAM...")
            if hasattr(model, "extract_conditioning") or hasattr(model, "get_speaker_embedding"):
                extract_fn = getattr(model, "extract_conditioning", getattr(model, "get_speaker_embedding", None))
                CACHED_EMBEDDINGS[prompt_path] = extract_fn(prompt_path)
            else:
                CACHED_EMBEDDINGS[prompt_path] = prompt_path
            print(f"[CACHE] Fitur karakteristik kustom untuk {prompt_path} berhasil dikunci di RAM.")

        # LOAD INSTAN: Mengambil data biner kustom tanpa membedah file .wav lagi
        speaker_embedding = CACHED_EMBEDDINGS[prompt_path]

        wav_parts = []
        total = len(chunks)
        
        sig = inspect.signature(model.generate)
        params = sig.parameters

        with torch.inference_mode():
            for i, ch in enumerate(chunks, start=1):
                progress((i - 1) / total, desc=f"Memproses kloning materi chunk {i}/{total}...")
                
                ch = re.sub(r"\s+", " ", ch.strip())
                if not re.search(r"[.!?…]$", ch):
                    ch += "."

                kwargs = {}
                if "audio_prompt_path" in params:
                    kwargs["audio_prompt_path"] = speaker_embedding
                elif "speaker_embedding" in params:
                    kwargs["speaker_embedding"] = speaker_embedding
                elif "prompt" in params and len(params) > 1:
                    list_keys = list(params.keys())
                    if len(list_keys) > 1:
                        kwargs[list_keys[1]] = speaker_embedding
                
                if "temperature" in params:
                    kwargs["temperature"] = 0.05
                if "top_p" in params:
                    kwargs["top_p"] = 0.7
                if "exaggeration" in params:
                    kwargs["exaggeration"] = 0.25
                if "cfg_weight" in params:
                    kwargs["cfg_weight"] = 0.3
                if "max_new_tokens" in params:
                    kwargs["max_new_tokens"] = 350

                try:
                    wav = model.generate(ch, **kwargs)
                except TypeError:
                    if "text" in params:
                        kwargs["text"] = ch
                        wav = model.generate(**kwargs)
                    else:
                        wav = model.generate(ch)

                if wav.dim() == 1:
                    wav = wav.unsqueeze(0)

                # Gabungkan potongan audio secara mulus (Seamless Concatenation)
                wav_parts.append(wav.detach().cpu().clone())

        progress(0.95, desc="Menyambungkan seluruh fragmentasi gelombang secara natural...")
        full_wav = torch.cat(wav_parts, dim=1)

        out_path = tempfile.NamedTemporaryFile(delete=False, suffix=".wav").name
        ta.save(out_path, full_wav, sr)

        del wav_parts
        del full_wav
        gc.collect()

        progress(1.0, desc="Selesai ✅")
        return out_path

    except Exception as e:
        print("[ERROR]", repr(e))
        print(traceback.format_exc())
        raise gr.Error(f"Gagal generate audio: {e}")

# =====================================================================
# INTERFACE DESIGN
# =====================================================================
with gr.Blocks(title="Chatterbox Persistent Weight Engine") as demo:
    gr.Markdown("## EduScanner AI Voice Backend - Pre-computed Weights Edition")
    gr.Markdown("Sistem memuat bobot suara bawaan secara permanen untuk memotong durasi penundaan CPU.")

    text_in = gr.Textbox(label="Teks Rangkuman Materi Kuliah", lines=8, placeholder="Ketik teks di sini...")
    wav_in = gr.Audio(label="Opsi Custom Voice (Kosongkan jika ingin pakai bobot default Mythia Batford)", type="filepath")
    url_in = gr.Textbox(label="Opsi URL File Audio Custom", placeholder="https://domain-kamu.com/audio.wav")
    
    btn = gr.Button("Sintesis Audio Kloning Suara", variant="primary")
    out_audio = gr.Audio(label="Hasil Audio Akhir (WAV)", type="filepath")

    btn.click(
        fn=clone_voice,
        inputs=[text_in, wav_in, url_in],
        outputs=[out_audio],
        api_name="clone_voice"
    )

from chatterbox.tts import ChatterboxTTS
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

if __name__ == "__main__":
    port = int(os.getenv("PORT", "7860"))
    demo.queue(default_concurrency_limit=1)
    demo.launch(server_name="0.0.0.0", server_port=port, show_error=True)

	import os
	# Paksa PyTorch murni berjalan di CPU sebelum library lain masuk
	os.environ["CUDA_VISIBLE_DEVICES"] = ""

	import re
	import gc
	import pickle
	import inspect
	import tempfile
	import traceback
	from pathlib import Path
	from threading import Lock

	import requests
	import torch
	import torchaudio as ta
	import gradio as gr

	# Optimasi Core CPU shared Hugging Face agar fokus dan stabil
	torch.set_num_threads(2)
	torch.set_num_interop_threads(2)
	torch.set_grad_enabled(False)
	torch.cuda.is_available = lambda: False # noqa: E731

	# Override global torch load untuk keamanan jalur CPU
	_original_torch_load = torch.load
	def _torch_load_cpu(args, *kwargs):
	kwargs["map_location"] = torch.device("cpu")
	return _original_torch_load(args, *kwargs)
	torch.load = _torch_load_cpu

	# =====================================================================
	# CONFIG & PATH MANAGEMENT (SWEET SPOT PARAMETERS)
	# =====================================================================
	MODEL_REPO = "grandhigh/Chatterbox-TTS-Indonesian"
	CHECKPOINT_FILENAME = "t3_cfg.safetensors"
	DEVICE = "cpu"

	MAX_TOTAL_CHARS = int(os.getenv("MAX_TOTAL_CHARS", "2400"))
	# Set ke 280 agar CPU ringan melakukan sampling awal dan tidak stuck di 0%
	MAX_CHARS_PER_CHUNK = int(os.getenv("MAX_CHARS_PER_CHUNK", "280"))
	MAX_CHUNKS = int(os.getenv("MAX_CHUNKS", "10"))
	PAUSE_SECONDS = float(os.getenv("PAUSE_SECONDS", "0.0"))
	DOWNLOAD_TIMEOUT = int(os.getenv("DOWNLOAD_TIMEOUT", "90"))

	ROOT_DIR = Path(__file__).parent
	DEFAULT_SPEAKER_PATH = ROOT_DIR / "default_speaker.wav"

	# FILE BOBOT PERMANEN: Lokasi penyimpanan biner karakteristik suara kustom kamu
	EMBEDDING_CACHE_PATH = ROOT_DIR / "speaker_embedding.pth"

	# Global RAM Cache untuk menyimpan koordinat suara siap pakai
	CACHED_EMBEDDINGS = {}
	_model = None
	_model_lock = Lock()


	def inisialisasi_bobot_default_permanen(model_instance):
	"""
	Fungsi untuk melatih/mengekstrak karakteristik file default_speaker.wav
	HANYA SEKALI saja saat startup. Hasilnya disimpan permanen ke disk (.pth).
	"""
	global CACHED_EMBEDDINGS
	try:
	# Jika file biner .pth sudah ada, langsung muat ke memori RAM secara instan
	if EMBEDDING_CACHE_PATH.exists():
	print("[INIT] Menemukan file bobot permanen speaker_embedding.pth. Memuat ke RAM...")
	with open(EMBEDDING_CACHE_PATH, "rb") as f:
	CACHED_EMBEDDINGS[str(DEFAULT_SPEAKER_PATH)] = pickle.load(f)
	print("[INIT] Bobot kustom bawaan siap digunakan instan!")
	return

	# Jika belum ada file .pth tapi ada file .wav asli, lakukan ekstraksi awal
	if DEFAULT_SPEAKER_PATH.exists():
	print("[INIT] Membuat file bobot baru. Mengekstrak default_speaker.wav...")
	if hasattr(model_instance, "extract_conditioning"):
	embedding = model_instance.extract_conditioning(str(DEFAULT_SPEAKER_PATH))
	elif hasattr(model_instance, "get_speaker_embedding"):
	embedding = model_instance.get_speaker_embedding(str(DEFAULT_SPEAKER_PATH))
	else:
	embedding = str(DEFAULT_SPEAKER_PATH)

	# Simpan secara fisik ke disk tingkat root agar permanen
	if not isinstance(embedding, str):
	with open(EMBEDDING_CACHE_PATH, "wb") as f:
	pickle.dump(embedding, f)
	print("[INIT] File speaker_embedding.pth berhasil disimpan secara permanen!")

	CACHED_EMBEDDINGS[str(DEFAULT_SPEAKER_PATH)] = embedding
	except Exception as e:
	print(f"[WARN] Gagal inisialisasi pembentukan bobot permanen awal: {e}")


	def get_model():
	global _model
	if _model is None:
	with _model_lock:
	if _model is None:
	print("[INIT] Loading model on CPU...")
	m = ChatterboxTTS.from_pretrained(device=DEVICE)

	ckpt_path = hf_hub_download(
	repo_id=MODEL_REPO,
	filename=CHECKPOINT_FILENAME
	)
	t3_state = load_file(ckpt_path, device="cpu")
	m.t3.load_state_dict(t3_state)

	if hasattr(m, "eval"):
	m.eval()

	_model = m
	print("[INIT] Model ready.")

	# JALANKAN PROSES TRAINING/EKSTRAKSI SEKALI SAAT STARTUP
	inisialisasi_bobot_default_permanen(_model)
	return _model


	def _resolve_audio_input(audio_file, audio_url: str):
	if isinstance(audio_file, str) and audio_file.strip():
	return audio_file
	if isinstance(audio_file, dict):
	p = audio_file.get("path")
	if p:
	return p
	if audio_url and audio_url.strip():
	try:
	r = requests.get(audio_url.strip(), timeout=DOWNLOAD_TIMEOUT)
	r.raise_for_status()
	tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
	tmp.write(r.content)
	tmp.close()
	return tmp.name
	except Exception:
	pass

	if DEFAULT_SPEAKER_PATH.exists():
	return str(DEFAULT_SPEAKER_PATH)

	return None


	def _split_text_safely(text: str, max_chars: int = MAX_CHARS_PER_CHUNK):
	text = re.sub(r"\s+", " ", (text or "").strip())
	if not text:
	return []
	sentences = re.split(r"(?<=[.!?])\s+", text)
	chunks = []
	current = ""
	for s in sentences:
	s = s.strip()
	if not s:
	continue
	parts = [s] if len(s) <= max_chars else re.split(r"(?<=[,;:])\s+", s)
	for p in parts:
	p = p.strip()
	if not p:
	continue
	if len(p) > max_chars:
	words = p.split()
	tmp = ""
	for w in words:
	cand = f"{tmp} {w}".strip() if tmp else w
	if len(cand) <= max_chars:
	tmp = cand
	else:
	if tmp:
	chunks.append(tmp)
	tmp = w
	if tmp:
	chunks.append(tmp)
	continue
	candidate = f"{current} {p}".strip() if current else p
	if len(candidate) <= max_chars:
	current = candidate
	else:
	if current:
	chunks.append(current)
	current = p
	if current:
	chunks.append(current)
	return chunks

	# =====================================================================
	# ENGINE UTAMA (FAST INFERENCE - INSTANT EMBEDDING LOADING)
	# =====================================================================
	def clone_voice(text: str, audio_file, audio_url: str, progress=gr.Progress(track_tqdm=False)):
	global CACHED_EMBEDDINGS
	try:
	raw_text = (text or "").strip()
	if not raw_text:
	raise gr.Error("Text prompt tidak boleh kosong.")

	prompt_path = _resolve_audio_input(audio_file, audio_url)
	if not prompt_path:
	raise gr.Error("Suara acuan default_speaker.wav tidak ditemukan di root server.")

	chunks = _split_text_safely(raw_text, max_chars=MAX_CHARS_PER_CHUNK)
	chunks = [ch for ch in chunks if re.search(r'[a-zA-Z0-9]', ch)]

	model = get_model()
	sr = getattr(model, "sr", 24000)
	torch.manual_seed(42)

	# Caching dinamis sebagai pengaman cadangan untuk Opsi Custom Voice Baru
	if prompt_path not in CACHED_EMBEDDINGS:
	progress(0.0, desc="Mengekstrak karakteristik gelombang audio kustom baru ke RAM...")
	if hasattr(model, "extract_conditioning") or hasattr(model, "get_speaker_embedding"):
	extract_fn = getattr(model, "extract_conditioning", getattr(model, "get_speaker_embedding", None))
	CACHED_EMBEDDINGS[prompt_path] = extract_fn(prompt_path)
	else:
	CACHED_EMBEDDINGS[prompt_path] = prompt_path
	print(f"[CACHE] Fitur karakteristik kustom untuk {prompt_path} berhasil dikunci di RAM.")

	# LOAD INSTAN: Mengambil data biner kustom tanpa membedah file .wav lagi
	speaker_embedding = CACHED_EMBEDDINGS[prompt_path]

	wav_parts = []
	total = len(chunks)

	sig = inspect.signature(model.generate)
	params = sig.parameters

	with torch.inference_mode():
	for i, ch in enumerate(chunks, start=1):
	progress((i - 1) / total, desc=f"Memproses kloning materi chunk {i}/{total}...")

	ch = re.sub(r"\s+", " ", ch.strip())
	if not re.search(r"[.!?…]$", ch):
	ch += "."

	kwargs = {}
	if "audio_prompt_path" in params:
	kwargs["audio_prompt_path"] = speaker_embedding
	elif "speaker_embedding" in params:
	kwargs["speaker_embedding"] = speaker_embedding
	elif "prompt" in params and len(params) > 1:
	list_keys = list(params.keys())
	if len(list_keys) > 1:
	kwargs[list_keys[1]] = speaker_embedding

	if "temperature" in params:
	kwargs["temperature"] = 0.05
	if "top_p" in params:
	kwargs["top_p"] = 0.7
	if "exaggeration" in params:
	kwargs["exaggeration"] = 0.25
	if "cfg_weight" in params:
	kwargs["cfg_weight"] = 0.3
	if "max_new_tokens" in params:
	kwargs["max_new_tokens"] = 350

	try:
	wav = model.generate(ch, **kwargs)
	except TypeError:
	if "text" in params:
	kwargs["text"] = ch
	wav = model.generate(**kwargs)
	else:
	wav = model.generate(ch)

	if wav.dim() == 1:
	wav = wav.unsqueeze(0)

	# Gabungkan potongan audio secara mulus (Seamless Concatenation)
	wav_parts.append(wav.detach().cpu().clone())

	progress(0.95, desc="Menyambungkan seluruh fragmentasi gelombang secara natural...")
	full_wav = torch.cat(wav_parts, dim=1)

	out_path = tempfile.NamedTemporaryFile(delete=False, suffix=".wav").name
	ta.save(out_path, full_wav, sr)

	del wav_parts
	del full_wav
	gc.collect()

	progress(1.0, desc="Selesai ✅")
	return out_path

	except Exception as e:
	print("[ERROR]", repr(e))
	print(traceback.format_exc())
	raise gr.Error(f"Gagal generate audio: {e}")

	# =====================================================================
	# INTERFACE DESIGN
	# =====================================================================
	with gr.Blocks(title="Chatterbox Persistent Weight Engine") as demo:
	gr.Markdown("## EduScanner AI Voice Backend - Pre-computed Weights Edition")
	gr.Markdown("Sistem memuat bobot suara bawaan secara permanen untuk memotong durasi penundaan CPU.")

	text_in = gr.Textbox(label="Teks Rangkuman Materi Kuliah", lines=8, placeholder="Ketik teks di sini...")
	wav_in = gr.Audio(label="Opsi Custom Voice (Kosongkan jika ingin pakai bobot default Mythia Batford)", type="filepath")
	url_in = gr.Textbox(label="Opsi URL File Audio Custom", placeholder="https://domain-kamu.com/audio.wav")

	btn = gr.Button("Sintesis Audio Kloning Suara", variant="primary")
	out_audio = gr.Audio(label="Hasil Audio Akhir (WAV)", type="filepath")

	btn.click(
	fn=clone_voice,
	inputs=[text_in, wav_in, url_in],
	outputs=[out_audio],
	api_name="clone_voice"
	)

	from chatterbox.tts import ChatterboxTTS
	from huggingface_hub import hf_hub_download
	from safetensors.torch import load_file

	if __name__ == "__main__":
	port = int(os.getenv("PORT", "7860"))
	demo.queue(default_concurrency_limit=1)
	demo.launch(server_name="0.0.0.0", server_port=port, show_error=True)