app.py

import os
import json
import torch
import logging
import commons
import utils
import random
import re
import gradio as gr
from models import SynthesizerTrn
from text import text_to_sequence
from torch import no_grad, LongTensor
from huggingface_hub import hf_hub_download

# --- 1. SETUP & CONFIG ---
logging.getLogger('numba').setLevel(logging.WARNING)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

os.makedirs("pretrained_models", exist_ok=True)
os.makedirs("PICTURE", exist_ok=True)

try:
    hps = utils.get_hparams_from_file(r'config/config.json')
except:
    hps = None

# --- 2. ASSETS & DOWNLOADER ---
BASE_REPO = "Blue-Archive/VITS-MODELS"

def get_category_image(cat_name):
    file_name = f"{cat_name}.PNG"
    local_path = os.path.join("PICTURE", file_name)
    if not os.path.exists(local_path):
        try:
            hf_hub_download(repo_id=BASE_REPO, filename=f"PICTURE/{file_name}", local_dir=".")
        except: return None
    return local_path

def download_assets(char_id, file_name):
    repo_id = "Blue-Archive/Library-VITS"
    local_dir = f"pretrained_models/{char_id}"
    local_path = os.path.join(local_dir, file_name)
    if not os.path.exists(local_path):
        try:
            hf_path = f"VITS-MODELS/pretrained_models/{char_id}/{file_name}"
            hf_hub_download(repo_id=repo_id, filename=hf_path, local_dir=".")
            os.makedirs(local_dir, exist_ok=True)
            if os.path.exists(hf_path): os.replace(hf_path, local_path)
        except: return None
    return local_path

# --- 3. ENGINE FUNCTIONS ---
def tts_execute(char_id, text, lang, speed):
    if not char_id or not text: return "❌ Pilih karakter dulu!", None
    model_pth = download_assets(char_id, f"{char_id}.pth")
    if not model_pth: return "❌ Gagal download model", None
    
    net_g = SynthesizerTrn(len(hps.symbols), hps.data.filter_length // 2 + 1, 
                           hps.train.segment_size // hps.data.hop_length, 
                           n_speakers=hps.data.n_speakers, **hps.model).to(device)
    utils.load_checkpoint(model_pth, net_g, None)
    _ = net_g.eval()

    text = f"[{'JA' if lang == 'Japanese' else 'ZH'}]{text}[{'JA' if lang == 'Japanese' else 'ZH'}]"
    try:
        text_norm, _ = text_to_sequence(text, hps.symbols, hps.data.text_cleaners)
        if hps.data.add_blank: text_norm = commons.intersperse(text_norm, 0)
        stn_tst = LongTensor(text_norm)
        with no_grad():
            audio = net_g.infer(stn_tst.unsqueeze(0).to(device), LongTensor([stn_tst.size(0)]).to(device), 
                                sid=LongTensor([0]).to(device), noise_scale=0.6, 
                                noise_scale_w=0.668, length_scale=1.0/speed)[0][0,0].data.cpu().float().numpy()
        return "✅ Selesai!", (22050, audio)
    except Exception as e: return f"Error: {e}", None

def get_random_jp():
    return random.choice(["こんにちは！", "お元気ですか？", "先生、お疲れ様です！", "大好きだよ！", "また明日ね。"])

# --- 4. CSS STYLING (THEME BLUE WITH ORANGE WARNING) ---
css = """
/* Loading Gradio Jadi Biru Secara Global */
:root { 
    --primary-600: #1299ff !important; 
    --accent-600: #1299ff !important; 
}

/* Spinner & Progress Bar Biru */
.progress-view, .progress-view > div, [role='progressbar'] > div { background-color: #1299ff !important; }
.loading { color: #1299ff !important; }

footer {display: none !important;}
.gradio-container { background-color: #f8fafc !important; }

/* Header & Status */
.ba-header-container { border: 1.5px solid #e1e8f0; border-radius: 12px; padding: 20px; margin-bottom: 12px; background: white; text-align: center; }
.ba-header-container h1 { color: #1299ff !important; font-weight: 800 !important; font-size: 32px !important; margin: 0; }
.status-container { border: 1.5px solid #e1e8f0; border-radius: 12px; padding: 15px; margin-bottom: 15px; background: white; }

/* Dual Scroll Boxes */
.scroll-box { height: 200px !important; overflow-y: auto !important; border: 1px solid #f0f4f8; border-radius: 12px; padding: 10px; background: #fafbfc; margin-bottom: 10px; }
.info-scroll-box { height: 200px !important; overflow-y: auto !important; padding: 5px; }

/* Character Buttons */
.char-btn { background: white !important; border: 1px solid #e2e8f0 !important; border-left: 5px solid #1299ff !important; text-align: left !important; padding: 10px !important; font-size: 13px !important; margin-bottom: 6px !important; width: 100%; color: #4a5568 !important; border-radius: 8px !important; }

/* Picture styling */
.image-display img { object-fit: contain !important; max-height: 150px !important; border-radius: 10px; margin: 0 auto; }

/* Warning Card (Warna Kembali ke Orange Semula) */
.warning-card { background: #fff9f0; border: 2px dashed #f5a623; border-radius: 10px; padding: 12px; margin: 10px 0; text-align: center; }

/* Character Information Grid */
.info-grid { display: grid !important; grid-template-columns: 1fr 1fr !important; gap: 8px !important; width: 100%; }
.info-item { border: 1px solid #f0f4f8; border-radius: 10px; padding: 10px; border-left: 4px solid #1299ff; background: white; box-shadow: 0 2px 4px rgba(0,0,0,0.02); }

/* Buttons & Inputs */
.gen-btn { background: #1299ff !important; color: white !important; font-weight: 700 !important; border-radius: 12px !important; height: 48px !important; width: 100%; border: none !important; }
.jp-btn { background: #f8fafc !important; border: 1px solid #cbd5e1 !important; font-size: 11px !important; border-radius: 8px !important; margin-bottom: 8px; width: 100%; font-weight: 700 !important; }
"""

# --- 5. DATA PREP ---
try:
    if not os.path.exists("pretrained_models/info.json"):
        hf_hub_download(repo_id="Blue-Archive/Library-VITS", filename="VITS-MODELS/pretrained_models/info.json", local_dir=".")
        os.replace("VITS-MODELS/pretrained_models/info.json", "pretrained_models/info.json")
    with open("pretrained_models/info.json", "r", encoding="utf-8") as f:
        MODELS_INFO = json.load(f)
except: MODELS_INFO = {}

CATEGORIES = {}
for m_id, info in MODELS_INFO.items():
    cat = info['title'].split("-")[0].strip()
    if cat == "Overwatch 2": continue
    if cat not in CATEGORIES: CATEGORIES[cat] = []
    CATEGORIES[cat].append((m_id, info))

# --- 6. UI CONSTRUCTION ---
with gr.Blocks(css=css, title="VITS MODELS") as demo:
    with gr.Column(elem_classes="slim-card"):
        gr.HTML("""
            <div class="ba-header-container"><h1>VITS MODELS</h1><p style='color:#64748b; font-size:13px; font-weight:600;'>Voice Library System</p></div>
            <div class="status-container">
                <div style='color:#1299ff; font-weight:800; font-size:16px; margin-bottom:5px;'>System Status</div>
                <div style='font-size:13px;'><span style='color:#4a5568'>Engine :</span> <span style='color:#28a745; font-weight:900;'>READY ✅</span></div>
                <div style='font-size:13px;'><span style='color:#4a5568'>System :</span> <span style='color:#1299ff; font-weight:700;'>Online</span></div>
            </div>
        """)
        
        with gr.Tabs():
            for idx, cat_name in enumerate(sorted(CATEGORIES.keys())):
                with gr.Tab(f"Model {idx}"):
                    gr.Markdown(f"### 📂 {cat_name}")
                    
                    cat_img = get_category_image(cat_name)
                    char_preview = gr.Image(value=cat_img, show_label=False, interactive=False, elem_classes="image-display")
                    
                    sel_char_id = gr.State("")
                    char_display = gr.Markdown("📍 *Silakan pilih karakter...*")
                    
                    # Selection Scroll Box
                    with gr.Column(elem_classes="scroll-box"):
                        for m_id, info in CATEGORIES[cat_name]:
                            name_en = info.get('name_en', m_id)
                            btn = gr.Button(f"👤 {name_en}", elem_classes="char-btn")
                            
                            def update_sel(c_id=m_id, n_en=name_en, cov=info.get('cover')):
                                img = download_assets(c_id, cov)
                                return c_id, f"📍 Selected: **{n_en}**", img
                            
                            btn.click(fn=update_sel, outputs=[sel_char_id, char_display, char_preview])
                    
                    # Warning Card (Warna Kembali ke Orange/Jingga)
                    gr.HTML("""<div class="warning-card"><b style='color:#f5a623;'>🔖 PERINGATAN MINNA 🔖</b><br><span style='font-size:11px; font-weight:600; color:#855d1a;'>Klik character, ketik teks, lalu Generate! ✨</span></div>""")
                    
                    txt_in = gr.TextArea(label="Input Text", placeholder="Masukkan teks...", lines=3)
                    gr.Button("🎲 INPUTS RANDOM TEXT 🎲", elem_classes="jp-btn").click(get_random_jp, outputs=[txt_in])
                    
                    with gr.Row():
                        lang_drop = gr.Dropdown(choices=["Japanese", "Chinese"], label="Lang", value="Japanese")
                        spd = gr.Slider(0.5, 2.0, 1.0, step=0.1, label="Speed Audio")
                        
                    btn_gen = gr.Button("🎐 GENERATE VOICE 🎐", elem_classes="gen-btn")
                    aud_out = gr.Audio(label="Voice Output")
                    
                    # Info Accordion with Scrollable Grid
                    with gr.Accordion("📑 Character Information 📑", open=False):
                        with gr.Column(elem_classes="info-scroll-box"):
                            with gr.Column(elem_classes="info-grid"):
                                for m_id, info in CATEGORIES[cat_name]:
                                    name_en = info.get('name_en', m_id)
                                    gr.HTML(f"""
                                        <div class="info-item">
                                            <div style="font-weight:800; color:#2d3748; font-size:13px;">{name_en}</div>
                                            <div style="color:#1299ff; font-size:10px; font-weight:700;">Character</div>
                                        </div>
                                    """)
                    
                    btn_gen.click(fn=tts_execute, inputs=[sel_char_id, txt_in, lang_drop, spd], outputs=[gr.Textbox(visible=False), aud_out])

        gr.HTML("<div style='text-align:center; margin-top:25px; padding:15px; background:white; border-radius:12px; border-bottom:4px solid #1299ff; color:#94a3b8; font-weight:700; font-size:12px; letter-spacing:2px;'>🌥️ CREATED BY MUTSUMI 🌥️</div>")

if __name__ == "__main__":
    demo.queue().launch()