app.py

# The code in this file is almost entirely written by LLMs and is much, much, much messier than it needs to be (at this point it's not clear to what extent it is even human-modifiable). We'd hope to improve this for any future local gradio release(s).

import tempfile
import os
import json
import time
import secrets
import logging
from pathlib import Path
from typing import Tuple, Any
from functools import partial

os.environ['HF_HUB_DISABLE_PROGRESS_BARS'] = '1'
logging.getLogger("huggingface_hub").setLevel(logging.ERROR)

import warnings

# Suppress torchaudio TorchCodec parameter warnings
warnings.filterwarnings('ignore', message='.*encoding.*parameter is not fully supported by TorchCodec')
warnings.filterwarnings('ignore', message='.*bits_per_sample.*parameter is not directly supported by TorchCodec')
warnings.filterwarnings('ignore', message='.* is not used by TorchCodec AudioEncoder. Format is determined by the file extension.')

import boto3
import gradio as gr
import numpy as np
import torch
import torchaudio
from botocore.exceptions import ClientError
from huggingface_hub import snapshot_download
import spaces

# R2 — upload generated audio to CDN so URLs never expire
_R2_ENDPOINT = os.getenv('R2_ENDPOINT')
_R2_ACCESS_KEY_ID = os.getenv('R2_ACCESS_KEY_ID')
_R2_SECRET_ACCESS_KEY = os.getenv('R2_SECRET_ACCESS_KEY')
_R2_BUCKET_NAME = os.getenv('R2_BUCKET_NAME')
_R2_PUBLIC_URL = os.getenv('R2_PUBLIC_URL', 'https://cdn.koevoice.com')

_r2_client = None

def _get_r2():
    global _r2_client
    if _r2_client is None and _R2_ENDPOINT:
        _r2_client = boto3.client(
            's3',
            endpoint_url=_R2_ENDPOINT,
            aws_access_key_id=_R2_ACCESS_KEY_ID,
            aws_secret_access_key=_R2_SECRET_ACCESS_KEY,
            region_name='auto',
        )
    return _r2_client

def upload_to_r2(local_path: Path) -> str | None:
    print(f"[R2] upload_to_r2 called: {local_path}")
    print(f"[R2] R2_ENDPOINT set: {bool(_R2_ENDPOINT)}, R2_BUCKET_NAME: {_R2_BUCKET_NAME!r}")
    r2 = _get_r2()
    if r2 is None:
        print(f"[R2] Client is None — R2_ENDPOINT not set, skipping upload")
        return None
    ext = local_path.suffix.lstrip('.')
    key = f"temp_audios/{local_path.name}"
    content_types = {'wav': 'audio/wav', 'flac': 'audio/flac', 'mp3': 'audio/mpeg'}
    content_type = content_types.get(ext, 'audio/wav')
    print(f"[R2] Uploading {local_path} → bucket={_R2_BUCKET_NAME!r} key={key!r} content_type={content_type}")
    try:
        r2.upload_file(
            str(local_path),
            _R2_BUCKET_NAME,
            key,
            ExtraArgs={'ContentType': content_type},
        )
        cdn_url = f"{_R2_PUBLIC_URL}/{key}"
        print(f"[R2] Upload succeeded: {cdn_url}")
        return cdn_url
    except Exception as e:
        print(f"[R2] Upload failed: {type(e).__name__}: {e}")
        return None

from inference import (
    load_model_from_hf,
    load_fish_ae_from_hf,
    load_pca_state_from_hf,
    load_audio,
    ae_reconstruct,
    sample_pipeline,
    find_flattening_point,
    ae_decode,
    get_text_input_ids_and_mask,
    get_speaker_latent_and_mask,
)
from samplers import sample_euler_cfg_any, GuidanceMode

import tarfile

# --------------------------------------------------------------------
### Configuration
MODEL_DTYPE = torch.bfloat16

# Maximum number of text prompts per batch generation
MAX_TEXTS = 20  # Maximum number of text prompts per batch

FISH_AE_DTYPE = torch.float32
# FISH_AE_DTYPE = torch.bfloat16 # MAYBE SLIGHTLY WORSE QUALITY, IF YOU HAVE ROOM, MAYBE USE FLOAT32

USE_16_BIT_WAV = True  # Save WAV files as 16-bit PCM instead of 32-bit float

# Audio Prompt Library for Custom Audio Panel (included in repo)
AUDIO_PROMPT_FOLDER = Path("./prompt_audio")


# If not on Zero GPU, compile fish_ae encoder/decoder on initialization
COMPILE_FISH_IF_NOT_ON_ZERO_GPU = True

# Silentcipher watermarking configuration
USE_SILENTCIPHER = True  # Enable/disable audio watermarking
SILENTCIPHER_MESSAGE = [91, 57, 81, 60, 83]  # Watermark message (list of integers)
SILENTCIPHER_SDR = 47  # Message SDR in dB (higher = less perceptible but less robust)

# Get HF token from environment for private model access
HF_TOKEN = os.environ.get("HF_TOKEN", None)
# --------------------------------------------------------------------

# Check if running on Zero GPU (compile incompatible with Zero GPU)
IS_ZEROGPU = os.environ.get("SPACES_ZERO_GPU") is not None


# print("FISH_AE_DTYPE:", FISH_AE_DTYPE)
# print("IS_ZEROGPU:", IS_ZEROGPU)
# if IS_ZEROGPU:
#     print("Running on Zero GPU - model compilation disabled")
# else:
#     print("Not on Zero GPU - model compilation available")

def _safe_members(tf, prefix):
    if not prefix.endswith('/'):
        prefix += '/'
    for m in tf.getmembers():
        if not m.name.startswith(prefix):
            continue
        p = Path(m.name)
        if any(part == '..' for part in p.parts) or p.is_absolute():
            continue
        yield m

def ensure_tar_tree(repo_id: str, root: str, *, token: str | None = None, max_workers: int = 4):
    os.environ.setdefault('HF_HUB_ENABLE_HF_TRANSFER', '1')
    from huggingface_hub import snapshot_download
    base = Path(snapshot_download(repo_id=repo_id, repo_type='dataset',
        allow_patterns=[f'{root}.tar', 'index.jsonl', 'README.md', 'LICENSE'],
        resume_download=True, token=token, max_workers=max_workers))
    root_dir = base / root
    if root_dir.exists():
        return root_dir
    tar_path = base / f'{root}.tar'
    if not tar_path.exists():
        raise FileNotFoundError(f'Expected {tar_path} in snapshot')
    with tarfile.open(tar_path, 'r') as tf:
        tf.extractall(base, members=_safe_members(tf, root))
    return root_dir


EARS_PATH = ensure_tar_tree(repo_id="jordand/echo-embeddings-ears-tar", root="EARS", token=HF_TOKEN)
VCTK_PATH = ensure_tar_tree(repo_id="jordand/echo-embeddings-vctk-tar", root="VCTK", token=HF_TOKEN)
EXPRESSO_PATH = ensure_tar_tree(repo_id="jordand/echo-embeddings-expresso-tar", root="Expresso", token=HF_TOKEN)


from huggingface_hub import snapshot_download

HF_CUSTOM_PATH = Path(snapshot_download(
    repo_id="jordand/echo-embeddings-custom",
    repo_type="dataset",
    allow_patterns=[
        "HF-Custom/**/speaker_latent.safetensors",
        "HF-Custom/**/metadata.json",
        "HF-Custom/**/audio.mp3",
    ],
    token=HF_TOKEN,
) + "/HF-Custom")

TEMP_AUDIO_DIR = Path('./temp_gradio_audio')
TEMP_AUDIO_DIR.mkdir(parents=True, exist_ok=True)

AUDIO_CACHE_DIR = Path('./audio_cache')
AUDIO_CACHE_DIR.mkdir(parents=True, exist_ok=True)

# Helper function to download audio from URL with caching
def download_audio_from_url(url: str) -> str:
    """Download audio from URL and cache it. Returns path to cached file."""
    import hashlib
    import urllib.request

    # Create cache key from URL
    url_hash = hashlib.md5(url.encode()).hexdigest()
    cache_path = AUDIO_CACHE_DIR / f"{url_hash}.wav"

    # Return cached file if exists
    if cache_path.exists():
        return str(cache_path)

    # Download to cache with proper headers to avoid 403
    try:
        req = urllib.request.Request(
            url,
            headers={
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
            }
        )
        with urllib.request.urlopen(req) as response:
            with open(cache_path, 'wb') as f:
                f.write(response.read())
        return str(cache_path)
    except Exception as e:
        raise ValueError(f"Failed to download audio from URL: {e}")

# Helper functions for unique filenames and cleanup
def make_stem(prefix: str, user_id: str | None = None) -> str:
    """Create unique filename stem: prefix__user__timestamp_random or prefix__timestamp_random if no user_id."""
    ts = int(time.time() * 1000)
    rand = secrets.token_hex(4)
    if user_id:
        return f"{prefix}__{user_id}__{ts}_{rand}"
    return f"{prefix}__{ts}_{rand}"

def cleanup_temp_audio(dir_: Path, user_id: str | None, max_age_sec: int = 60 * 5):
    """Remove old files globally and all previous files for this user."""
    now = time.time()
    
    # 1) Global TTL: remove any file older than max_age_sec
    for p in dir_.glob("*"):
        try:
            if p.is_file() and (now - p.stat().st_mtime) > max_age_sec:
                p.unlink(missing_ok=True)
        except Exception:
            pass
    
    # 2) Per-user: remove ALL previous files for this user (we don't need to keep any)
    if user_id:
        for p in dir_.glob(f"*__{user_id}__*"):
            try:
                if p.is_file():
                    p.unlink(missing_ok=True)
            except Exception:
                pass

TEXT_PRESETS_PATH = Path('./text_presets.txt')

SAMPLER_PRESETS_PATH = Path('./sampler_presets.json')

# Load models at module level so ZeroGPU can preload shards before each @spaces.GPU call.
# ZeroGPU runs a CUDA emulation mode outside @spaces.GPU, so .to('cuda') works here and
# transfers are optimized for this startup path (lazy loading inside the decorator is slower).
model = load_model_from_hf(dtype=MODEL_DTYPE, compile=False, token=HF_TOKEN)
fish_ae = load_fish_ae_from_hf(compile=(COMPILE_FISH_IF_NOT_ON_ZERO_GPU and not IS_ZEROGPU), dtype=FISH_AE_DTYPE, token=HF_TOKEN)
pca_state = load_pca_state_from_hf(token=HF_TOKEN)

model_compiled = None  # Separate compiled model for toggling
_model_compiled = False

silentcipher_model = None
if USE_SILENTCIPHER:
    try:
        import silentcipher
        silentcipher_model = silentcipher.get_model(model_type='44.1k', device='cuda')
    except Exception as e:
        print(f"Warning: Failed to load silentcipher model: {e}")
        print("Continuing without watermarking...")

# Note: Speaker cache (including zero speaker) is built fresh inside each @spaces.GPU call
# to avoid tensor state issues across GPU allocations in ZeroGPU environment

def load_models():
    """No-op kept for compatibility; models are now loaded at module level."""
    pass



def compile_model(should_compile):
    """Compile the model for faster inference."""
    global model, model_compiled, _model_compiled
    
    # If on Zero GPU, compilation is not supported
    if IS_ZEROGPU:
        return gr.update(value=False, interactive=False), gr.update(value="⚠️ Compile disabled on Zero GPU", visible=True)
    
    if not should_compile:
        # User unchecked - clear status and allow toggling
        return gr.update(value=False, interactive=True), gr.update(value="", visible=False)
    
    if _model_compiled:
        # Already compiled - just show status
        return gr.update(value=True, interactive=True), gr.update(value="✓ Model already compiled", visible=True)
    
    # Need to compile - disable checkbox temporarily and show status
    return gr.update(value=True, interactive=False), gr.update(value="⏳ Compiling... (1-3 minutes)", visible=True)


def do_compile():
    """Actually perform the compilation by creating a separate compiled model."""
    global model, model_compiled, _model_compiled
    
    # Skip if on Zero GPU
    if IS_ZEROGPU:
        return gr.update(value="⚠️ Compile disabled on Zero GPU", visible=True), gr.update(interactive=False)
    
    if _model_compiled:
        return gr.update(value="", visible=False), gr.update(interactive=True)
    
    try:
        # print("Compiling model... This will take 1-3 minutes on first run.")
        # print("Creating a separate compiled model for toggling...")
        
        # Create a compiled version of the model
        model_compiled = torch.compile(model)
        model_compiled.get_kv_cache = torch.compile(model.get_kv_cache)
        model_compiled.get_kv_cache_from_precomputed_speaker_state = torch.compile(model.get_kv_cache_from_precomputed_speaker_state)
        
        _model_compiled = True
        # print("Compilation complete! You can now toggle between compiled/uncompiled.")
        return gr.update(value="", visible=False), gr.update(interactive=True)
    except Exception as e:
        print(f"Compilation failed: {str(e)}")
        return gr.update(value=f"✗ Compilation failed: {str(e)}", visible=True), gr.update(interactive=True)


def save_audio_with_format(audio_tensor: torch.Tensor, base_path: Path, filename: str, sample_rate: int, audio_format: str) -> Path:
    """Save audio in specified format, fallback to WAV if MP3 encoding fails."""
    if audio_format == "mp3":
        try:
            output_path = base_path / f"{filename}.mp3"
            # Try to save as MP3
            torchaudio.save(
                str(output_path),
                audio_tensor,
                sample_rate,
                format="mp3",
                encoding="mp3",
                bits_per_sample=None
            )
            # print(f"Successfully saved as MP3: {output_path}")
            return output_path
        except Exception as e:
            print(f"MP3 encoding failed: {e}, falling back to WAV")
            # Fallback to WAV
            output_path = base_path / f"{filename}.wav"
            if USE_16_BIT_WAV:
                torchaudio.save(str(output_path), audio_tensor, sample_rate, encoding="PCM_S", bits_per_sample=16)
            else:
                torchaudio.save(str(output_path), audio_tensor, sample_rate)
            return output_path
    else:
        # Save as WAV
        output_path = base_path / f"{filename}.wav"
        if USE_16_BIT_WAV:
            torchaudio.save(str(output_path), audio_tensor, sample_rate, encoding="PCM_S", bits_per_sample=16)
        else:
            torchaudio.save(str(output_path), audio_tensor, sample_rate)
        return output_path


@spaces.GPU
def generate_audio(
    texts_input: str,
    speaker_st_path: str,
    speaker_audio_path: str,
    # Sampling parameters
    num_steps: int,
    rng_seed: int,
    cfg_mode: str,
    cfg_scale_text: float,
    cfg_scale_speaker: float,
    cfg_min_t: float,
    cfg_max_t: float,
    truncation_factor: float,
    rescale_k: float,
    rescale_sigma: float,
    speaker_k_enable: bool,
    speaker_k_scale: float,
    speaker_k_min_t: float,
    speaker_k_max_layers: int,
    apg_eta_text: float,
    apg_eta_speaker: float,
    apg_momentum_text: float,
    apg_momentum_speaker: float,
    apg_norm_text: str,
    apg_norm_speaker: str,
    reconstruct_first_30_seconds: bool,
    use_custom_shapes: bool,
    max_text_byte_length: str,
    max_speaker_latent_length: str,
    sample_latent_len: str,
    audio_format: str,
    use_compile: bool,
    show_original_audio: bool,
    session_id: str,
):
    """Batch audio generation with speaker cache reuse for the Advanced View."""

    # Choose which model to use based on compile setting
    global model, model_compiled
    active_model = model_compiled if (use_compile and model_compiled is not None) else model

    if use_compile and model_compiled is None:
        print("Warning: Compile requested but model not yet compiled. Using uncompiled model.")

    # Cleanup old temp files globally and remove ALL previous files for this user
    cleanup_temp_audio(TEMP_AUDIO_DIR, session_id)

    # Extract texts from input (one per line)
    texts = []
    if texts_input and texts_input.strip():
        for line in texts_input.strip().split('\n'):
            line = line.strip()
            if line:
                texts.append(line)
    texts = texts[:MAX_TEXTS]

    if not texts:
        # No valid texts - return all hidden
        empty_results = [gr.update()] * (1 + MAX_TEXTS * 2 + 7)  # section + labels + audios + other outputs
        yield tuple(empty_results)
        return

    # Check if speaker is provided (now optional for zero conditioning)
    use_zero_speaker = not speaker_audio_path or speaker_audio_path == ""
    if use_zero_speaker:
        speaker_audio_path = None

    start_time = time.time()
    
    # Parse parameters (most are already numeric from gr.Number)
    num_steps_int = min(max(int(num_steps), 1), 80)  # Clamp to [1, 80]
    rng_seed_int = int(rng_seed) if rng_seed is not None else 0
    cfg_scale_text_val = float(cfg_scale_text)
    cfg_min_t_val = float(cfg_min_t)
    cfg_max_t_val = float(cfg_max_t)
    truncation_factor_val = float(truncation_factor)
    rescale_k_val = float(rescale_k) if rescale_k != 1.0 else None  # 1.0 means "off"
    rescale_sigma_val = float(rescale_sigma)
    
    # Determine guidance mode
    if cfg_mode == "independent":
        guidance_mode = GuidanceMode.INDEPENDENT
        cfg_scale_speaker_val = float(cfg_scale_speaker) if cfg_scale_speaker is not None else None
        apg_eta_text_val = None
        apg_eta_speaker_val = None
        apg_momentum_text_val = None
        apg_momentum_speaker_val = None
        apg_norm_text_val = None
        apg_norm_speaker_val = None
    elif cfg_mode == "alternating":
        guidance_mode = GuidanceMode.ALTERNATING
        cfg_scale_speaker_val = float(cfg_scale_speaker) if cfg_scale_speaker is not None else None
        apg_eta_text_val = None
        apg_eta_speaker_val = None
        apg_momentum_text_val = None
        apg_momentum_speaker_val = None
        apg_norm_text_val = None
        apg_norm_speaker_val = None
    elif cfg_mode == "apg-independent":
        guidance_mode = GuidanceMode.APG
        cfg_scale_speaker_val = float(cfg_scale_speaker) if cfg_scale_speaker is not None else None
        apg_eta_text_val = float(apg_eta_text) if apg_eta_text is not None else None
        apg_eta_speaker_val = float(apg_eta_speaker) if apg_eta_speaker is not None else None
        apg_momentum_text_val = float(apg_momentum_text) if apg_momentum_text is not None else None
        apg_momentum_speaker_val = float(apg_momentum_speaker) if apg_momentum_speaker is not None else None
        # APG norm: parse and validate (must be positive to be meaningful)
        apg_norm_text_val = float(apg_norm_text) if apg_norm_text.strip() else None
        if apg_norm_text_val is not None and apg_norm_text_val <= 0:
            apg_norm_text_val = None
        apg_norm_speaker_val = float(apg_norm_speaker) if apg_norm_speaker.strip() else None
        if apg_norm_speaker_val is not None and apg_norm_speaker_val <= 0:
            apg_norm_speaker_val = None
    else:  # "joint-unconditional"
        guidance_mode = GuidanceMode.JOINT
        # For unconditional, speaker scale must be None
        cfg_scale_speaker_val = None
        apg_eta_text_val = None
        apg_eta_speaker_val = None
        apg_momentum_text_val = None
        apg_momentum_speaker_val = None
        apg_norm_text_val = None
        apg_norm_speaker_val = None
    
    # Parse speaker K scale parameters (available for all modes)
    if speaker_k_enable:
        speaker_k_scale_val = float(speaker_k_scale) if speaker_k_scale is not None else None
        speaker_k_min_t_val = float(speaker_k_min_t) if speaker_k_min_t is not None else None
        speaker_k_max_layers_val = int(speaker_k_max_layers) if speaker_k_max_layers is not None else None
    else:
        speaker_k_scale_val = None
        speaker_k_min_t_val = None
        speaker_k_max_layers_val = None

    # Parse custom shapes if enabled
    if use_custom_shapes:
        # Allow blank/empty values for first two fields (will use None)
        pad_to_max_text_seq_len = int(max_text_byte_length) if max_text_byte_length.strip() else None
        pad_to_max_speaker_latent_len = int(max_speaker_latent_length) if max_speaker_latent_length.strip() else None
        sample_latent_len_val = int(sample_latent_len) if sample_latent_len.strip() else 640
    else:
        pad_to_max_text_seq_len = 768
        pad_to_max_speaker_latent_len = 2560
        sample_latent_len_val = 640
    
    # Build speaker cache fresh inside @spaces.GPU (avoids tensor state issues)
    zero_speaker_cache = get_zero_speaker_cache(active_model, device='cuda')

    if speaker_audio_path is not None:
        speaker_audio = load_audio(speaker_audio_path).cuda()
        speaker_cache = build_speaker_cache(
            active_model, fish_ae, pca_state, speaker_audio,
            pad_to_max_speaker_latent_len
        )
    else:
        speaker_audio = None
        speaker_cache = zero_speaker_cache
    
    # Compute speaker-level outputs ONCE (original/reconstruction)
    recon_output_path = None
    original_output_path = None

    if reconstruct_first_30_seconds and speaker_audio_path:
        audio_recon = ae_reconstruct(
            fish_ae=fish_ae,
            pca_state=pca_state,
            audio=torch.nn.functional.pad(
                speaker_audio[..., :2048 * 640],
                (0, max(0, 2048 * 640 - speaker_audio.shape[-1]))
            )[None],
        )[..., :speaker_audio.shape[-1]]

        recon_stem = make_stem("speaker_recon", session_id)
        recon_output_path = save_audio_with_format(
            audio_recon.cpu()[0],
            TEMP_AUDIO_DIR,
            recon_stem,
            44100,
            audio_format
        )

    if show_original_audio and speaker_audio_path:
        original_stem = make_stem("original_audio", session_id)
        original_output_path = save_audio_with_format(
            speaker_audio.cpu(),
            TEMP_AUDIO_DIR,
            original_stem,
            44100,
            audio_format
        )

    show_reference_section = (show_original_audio or reconstruct_first_30_seconds) and speaker_audio_path is not None

    # Initialize output arrays
    result_labels = [gr.update(visible=False)] * MAX_TEXTS
    result_audios = [gr.update(visible=False)] * MAX_TEXTS

    # Generate audio for each text (reusing speaker cache)
    for i, text in enumerate(texts):
        text_start = time.time()

        # Build KV cache from speaker cache + fresh text encoding
        text_input_ids, text_mask = get_text_input_ids_and_mask(
            [text], pad_to_max_text_seq_len, device='cuda'
        )
        kv_cache, combined_text_mask, combined_speaker_mask, text_mask_single, text_mask_uncond = build_kv_cache_from_speaker_cache(
            active_model, speaker_cache, zero_speaker_cache,
            text_input_ids, text_mask, guidance_mode=cfg_mode
        )

        # Sample latent
        torch.manual_seed(rng_seed_int + i)  # Different seed per text
        x_t = torch.randn((1, sample_latent_len_val, 80), device='cuda', dtype=torch.float32)
        if truncation_factor_val is not None:
            x_t = x_t * truncation_factor_val

        # Diffusion loop with prebuilt KV cache
        t_schedule = torch.linspace(1., 0., num_steps_int + 1, device='cuda') * 0.999

        # Clone KV cache only if we'll modify it (speaker KV scaling)
        if speaker_k_scale_val is not None and speaker_k_min_t_val is not None:
            kv_cache = [(k.clone(), v.clone()) for k, v in kv_cache]
            # Apply speaker KV scaling
            for layer_i in range(min(speaker_k_max_layers_val or 24, len(kv_cache))):
                text_len = text_input_ids.shape[-1]
                kv_cache[layer_i][0][:, text_len:] *= speaker_k_scale_val
                kv_cache[layer_i][1][:, text_len:] *= speaker_k_scale_val

        # APG momentum buffers (only for APG mode)
        buf_text = torch.zeros_like(x_t) if cfg_mode == "apg-independent" else None
        buf_speaker = torch.zeros_like(x_t) if cfg_mode == "apg-independent" else None

        for step_i in range(num_steps_int):
            t, t_next = t_schedule[step_i], t_schedule[step_i+1]
            has_cfg = ((t >= cfg_min_t_val) * (t <= cfg_max_t_val)).item()

            if has_cfg:
                if cfg_mode == "independent":
                    # 3× batch: [cond, text_uncond, speaker_uncond]
                    v_cond, v_uncond_text, v_uncond_speaker = active_model(
                        x=torch.cat([x_t, x_t, x_t], dim=0).to(active_model.dtype),
                        t=(torch.ones((3,), device='cuda') * t).to(active_model.dtype),
                        text_input_ids=None,
                        text_mask=combined_text_mask,
                        speaker_latent=None,
                        speaker_mask=combined_speaker_mask,
                        kv_cache=kv_cache,
                    ).float().chunk(3, dim=0)
                    v_pred = v_cond + cfg_scale_text_val * (v_cond - v_uncond_text) + cfg_scale_speaker_val * (v_cond - v_uncond_speaker)

                elif cfg_mode == "apg-independent":
                    # 3× batch with APG (Adaptive Projected Guidance)
                    v_cond, v_uncond_text, v_uncond_speaker = active_model(
                        x=torch.cat([x_t, x_t, x_t], dim=0).to(active_model.dtype),
                        t=(torch.ones((3,), device='cuda') * t).to(active_model.dtype),
                        text_input_ids=None,
                        text_mask=combined_text_mask,
                        speaker_latent=None,
                        speaker_mask=combined_speaker_mask,
                        kv_cache=kv_cache,
                    ).float().chunk(3, dim=0)

                    x0_cond = x_t - t * v_cond
                    x0_uncond_text = x_t - t * v_uncond_text
                    x0_uncond_speaker = x_t - t * v_uncond_speaker

                    diff_text = x0_cond - x0_uncond_text
                    diff_speaker = x0_cond - x0_uncond_speaker

                    buf_text = diff_text + apg_momentum_text_val * buf_text
                    diff_text = buf_text

                    buf_speaker = diff_speaker + apg_momentum_speaker_val * buf_speaker
                    diff_speaker = buf_speaker

                    # Norm clipping
                    if apg_norm_text_val is not None:
                        nt = torch.sqrt((diff_text * diff_text).sum(dim=tuple(range(1, diff_text.dim())), keepdim=True) + 1e-12)
                        s = torch.minimum(torch.ones_like(nt), apg_norm_text_val / nt)
                        diff_text = diff_text * s
                    if apg_norm_speaker_val is not None:
                        ns = torch.sqrt((diff_speaker * diff_speaker).sum(dim=tuple(range(1, diff_speaker.dim())), keepdim=True) + 1e-12)
                        s = torch.minimum(torch.ones_like(ns), apg_norm_speaker_val / ns)
                        diff_speaker = diff_speaker * s

                    # Orthogonal projection
                    c_norm = torch.sqrt((x0_cond * x0_cond).sum(dim=tuple(range(1, x0_cond.dim())), keepdim=True) + 1e-12)
                    c_hat = x0_cond / c_norm

                    par_text = (diff_text * c_hat).sum(dim=tuple(range(1, diff_text.dim())), keepdim=True) * c_hat
                    ort_text = diff_text - par_text
                    upd_text = ort_text + apg_eta_text_val * par_text

                    par_speaker = (diff_speaker * c_hat).sum(dim=tuple(range(1, diff_speaker.dim())), keepdim=True) * c_hat
                    ort_speaker = diff_speaker - par_speaker
                    upd_speaker = ort_speaker + apg_eta_speaker_val * par_speaker

                    x0_pred = x0_cond + cfg_scale_text_val * upd_text + cfg_scale_speaker_val * upd_speaker
                    v_pred = (x_t - x0_pred) / t

                elif cfg_mode == "alternating":
                    # 2× batch: alternates between text and speaker uncond per step
                    v_cond, v_uncond = active_model(
                        x=torch.cat([x_t, x_t], dim=0).to(active_model.dtype),
                        t=(torch.ones((2,), device='cuda') * t).to(active_model.dtype),
                        text_input_ids=None,
                        text_mask=torch.cat([text_mask_single, text_mask_uncond if step_i % 2 == 0 else text_mask_single], dim=0),
                        speaker_latent=None,
                        speaker_mask=torch.cat([speaker_cache.speaker_mask, speaker_cache.speaker_mask if step_i % 2 == 0 else zero_speaker_cache.speaker_mask], dim=0),
                        kv_cache=kv_cache,
                    ).float().chunk(2, dim=0)
                    scale = cfg_scale_text_val if step_i % 2 == 0 else cfg_scale_speaker_val
                    v_pred = v_cond + scale * (v_cond - v_uncond)

                else:  # JOINT
                    # 2× batch: [cond, joint_uncond]
                    v_cond, v_uncond = active_model(
                        x=torch.cat([x_t, x_t], dim=0).to(active_model.dtype),
                        t=(torch.ones((2,), device='cuda') * t).to(active_model.dtype),
                        text_input_ids=None,
                        text_mask=combined_text_mask,
                        speaker_latent=None,
                        speaker_mask=combined_speaker_mask,
                        kv_cache=kv_cache,
                    ).float().chunk(2, dim=0)
                    v_pred = v_cond + cfg_scale_text_val * (v_cond - v_uncond)
            else:
                # No CFG - use conditional only
                kv_cache_single = [(kv_cache[j][0][:1], kv_cache[j][1][:1]) for j in range(len(kv_cache))]
                v_pred = active_model(
                    x=x_t.to(active_model.dtype),
                    t=(torch.ones((1,), device='cuda') * t).to(active_model.dtype),
                    text_input_ids=None,
                    text_mask=text_mask_single,
                    speaker_latent=None,
                    speaker_mask=speaker_cache.speaker_mask,
                    kv_cache=kv_cache_single,
                ).float()

            # Temporal rescaling if enabled
            if rescale_k_val is not None and rescale_sigma_val is not None and t < 1:
                snr = (1 - t) ** 2 / (t ** 2)
                ratio = (snr * rescale_sigma_val ** 2 + 1) / (snr * rescale_sigma_val ** 2 / rescale_k_val + 1)
                v_pred = 1 / (1 - t) * (ratio * ((1 - t) * v_pred + x_t) - x_t)

            # Revert speaker KV scaling at the threshold
            if speaker_k_scale_val is not None and speaker_k_min_t_val is not None and t_next < speaker_k_min_t_val and t >= speaker_k_min_t_val:
                for layer_i in range(min(speaker_k_max_layers_val or 24, len(kv_cache))):
                    text_len = text_input_ids.shape[-1]
                    kv_cache[layer_i][0][:, text_len:] /= speaker_k_scale_val
                    kv_cache[layer_i][1][:, text_len:] /= speaker_k_scale_val

            x_t = x_t + v_pred * (t_next - t)

        # Decode latent to audio
        flattening_point = find_flattening_point(x_t[0])
        audio_out = ae_decode(fish_ae, pca_state, x_t)
        audio_out = audio_out[..., :flattening_point * 2048]

        # Apply watermark
        audio_to_save = audio_out[0].cpu()
        if USE_SILENTCIPHER and silentcipher_model is not None:
            try:
                audio_numpy = audio_to_save.squeeze(0).numpy()
                encoded_audio, sdr = silentcipher_model.encode_wav(
                    audio_numpy, 44100, SILENTCIPHER_MESSAGE, message_sdr=SILENTCIPHER_SDR
                )
                audio_to_save = torch.tensor(encoded_audio).unsqueeze(0)
            except Exception as e:
                print(f"Warning: Watermarking failed for text {i+1}: {e}")

        # Save audio
        stem = make_stem(f"generated_{i+1}", session_id)
        output_path = save_audio_with_format(
            audio_to_save, TEMP_AUDIO_DIR, stem, 44100, audio_format
        )

        text_time = time.time() - text_start
        _cdn = upload_to_r2(output_path)
        result_labels[i] = gr.update(value=f"**Text {i+1}** — ⏱️ {text_time:.1f}s", visible=True)
        result_audios[i] = gr.update(value=_cdn if _cdn else str(output_path), visible=True)

        # Yield progressive results (section, labels, audios, other outputs)
        generation_time = time.time() - start_time
        time_str = f"⏱️ Total: {generation_time:.1f}s | Generated {i+1}/{len(texts)} texts"
        text_display = f"**Generated {i+1}/{len(texts)} texts**"

        yield (
            gr.update(),  # generated_section
            *result_labels,
            *result_audios,
            gr.update(value=text_display, visible=True),  # text_prompt_display
            gr.update(value=str(original_output_path) if original_output_path else None, visible=True),  # original_audio
            gr.update(value=time_str, visible=True),  # generation_time_display
            gr.update(value=str(recon_output_path) if recon_output_path else None, visible=True),  # reference_audio
            gr.update(visible=(show_original_audio and speaker_audio_path is not None)),  # original_accordion
            gr.update(visible=(reconstruct_first_30_seconds and speaker_audio_path is not None)),  # reference_accordion
            gr.update(visible=show_reference_section)  # reference_audio_header
        )


@spaces.GPU
def generate_batch_simple(
    texts_input: str,
    speaker_audio_path: str,
    preset_name: str,
    rng_seed: int,
    num_steps: int,
    speaker_kv_enable: bool,
    speaker_kv_scale: float,
    session_id: str,
):
    """Batch audio generation with speaker cache reuse for the Simple View."""

    # Use compiled model if available, otherwise uncompiled
    global model, model_compiled
    active_model = model_compiled if model_compiled is not None else model

    # Cleanup old temp files
    cleanup_temp_audio(TEMP_AUDIO_DIR, session_id)

    # Extract texts from input (one per line)
    texts = []
    if texts_input and texts_input.strip():
        for line in texts_input.strip().split('\n'):
            line = line.strip()
            if line:
                texts.append(line)
    texts = texts[:MAX_TEXTS]

    if not texts:
        # No valid texts - return all hidden
        empty_results = []
        for i in range(MAX_TEXTS):
            empty_results.append(gr.update(visible=False))  # labels
        for i in range(MAX_TEXTS):
            empty_results.append(gr.update(visible=False))  # audios
        yield tuple(empty_results)
        return

    start_time = time.time()
    
    # Check if speaker is provided
    use_zero_speaker = not speaker_audio_path or speaker_audio_path == ""
    if use_zero_speaker:
        speaker_audio_path = None
    
    start_time = time.time()
    
    # Load preset values
    presets = load_sampler_presets()
    preset = presets.get(preset_name, {})
    
    # Helper to convert string values to float
    def to_float(val, default):
        try:
            return float(val) if val is not None else default
        except (ValueError, TypeError):
            return default
    
    # Apply preset values (or use defaults)
    num_steps_int = min(max(int(num_steps), 1), 80)
    rng_seed_int = int(rng_seed) if rng_seed is not None else 0
    cfg_scale_text_val = to_float(preset.get("cfg_scale_text"), 3.0)
    cfg_scale_speaker_val = to_float(preset.get("cfg_scale_speaker"), 8.0)
    cfg_min_t_val = to_float(preset.get("cfg_min_t"), 0.5)
    cfg_max_t_val = to_float(preset.get("cfg_max_t"), 1.0)
    truncation_factor_val = to_float(preset.get("truncation_factor"), 1.0)
    rescale_k_raw = to_float(preset.get("rescale_k"), 1.0)
    rescale_k_val = rescale_k_raw if rescale_k_raw != 1.0 else None  # 1.0 means off
    rescale_sigma_val = to_float(preset.get("rescale_sigma"), 3.0)
    guidance_mode = GuidanceMode.INDEPENDENT  # Simple view always uses independent
    
    # Speaker KV parameters (user override takes precedence)
    if speaker_kv_enable:
        speaker_k_scale_val = float(speaker_kv_scale) if speaker_kv_scale else 1.5
        speaker_k_min_t_val = 0.9
        speaker_k_max_layers_val = 24
    else:
        speaker_k_scale_val = None
        speaker_k_min_t_val = None
        speaker_k_max_layers_val = None
    
    # Default shapes
    pad_to_max_text_seq_len = 768
    pad_to_max_speaker_latent_len = 2560
    sample_latent_len_val = 640
    
    # Create sample function with parameters
    sample_fn = partial(
        sample_euler_cfg_any,
        num_steps=num_steps_int,
        guidance_mode=guidance_mode,
        cfg_scale_text=cfg_scale_text_val,
        cfg_scale_speaker=cfg_scale_speaker_val,
        cfg_min_t=cfg_min_t_val,
        cfg_max_t=cfg_max_t_val,
        truncation_factor=truncation_factor_val,
        rescale_k=rescale_k_val,
        rescale_sigma=rescale_sigma_val,
        speaker_k_scale=speaker_k_scale_val,
        speaker_k_min_t=speaker_k_min_t_val,
        speaker_k_max_layers=speaker_k_max_layers_val,
        apg_eta_text=None,
        apg_eta_speaker=None,
        apg_momentum_text=None,
        apg_momentum_speaker=None,
        apg_norm_text=None,
        apg_norm_speaker=None,
        block_size=sample_latent_len_val
    )
    
    # Load speaker audio once (reused across all texts)
    if speaker_audio_path is not None:
        speaker_audio = load_audio(speaker_audio_path).cuda()
    else:
        speaker_audio = None

    # Create sampler function
    sample_fn = partial(
        sample_euler_cfg_any,
        num_steps=num_steps_int,
        guidance_mode=GuidanceMode.INDEPENDENT,
        cfg_scale_text=cfg_scale_text_val,
        cfg_scale_speaker=cfg_scale_speaker_val,
        cfg_min_t=cfg_min_t_val,
        cfg_max_t=cfg_max_t_val,
        truncation_factor=truncation_factor_val,
        rescale_k=rescale_k_val,
        rescale_sigma=rescale_sigma_val,
        speaker_k_scale=speaker_k_scale_val,
        speaker_k_min_t=speaker_k_min_t_val,
        speaker_k_max_layers=speaker_k_max_layers_val,
        apg_eta_text=None,
        apg_eta_speaker=None,
        apg_momentum_text=None,
        apg_momentum_speaker=None,
        apg_norm_text=None,
        apg_norm_speaker=None,
        block_size=sample_latent_len_val
    )

    # Initialize output arrays
    result_labels = [gr.update(visible=False)] * MAX_TEXTS
    result_audios = [gr.update(visible=False)] * MAX_TEXTS

    # Generate audio for each text (reusing speaker audio - encoded fresh inside sample_pipeline)
    for i, text in enumerate(texts):
        text_start = time.time()

        # Use original sample_pipeline with speaker audio reuse
        audio_out = sample_pipeline(
            model=active_model,
            fish_ae=fish_ae,
            pca_state=pca_state,
            sample_fn=sample_fn,
            text_prompt=text,
            speaker_audio=speaker_audio,
            rng_seed=rng_seed_int + i,
            pad_to_max_text_seq_len=pad_to_max_text_seq_len,
            pad_to_max_speaker_latent_len=pad_to_max_speaker_latent_len,
        )

        # Apply watermark
        audio_to_save = audio_out[0].cpu()
        if USE_SILENTCIPHER and silentcipher_model is not None:
            try:
                audio_numpy = audio_to_save.squeeze(0).numpy()
                encoded_audio, sdr = silentcipher_model.encode_wav(
                    audio_numpy, 44100, SILENTCIPHER_MESSAGE, message_sdr=SILENTCIPHER_SDR
                )
                audio_to_save = torch.tensor(encoded_audio).unsqueeze(0)
            except Exception as e:
                print(f"Warning: Watermarking failed for text {i+1}: {e}")

        # Save audio
        stem = make_stem(f"generated_simple_{i+1}", session_id)
        output_path = save_audio_with_format(
            audio_to_save, TEMP_AUDIO_DIR, stem, 44100, "wav"
        )

        text_time = time.time() - text_start
        _cdn = upload_to_r2(output_path)
        result_labels[i] = gr.update(value=f"**Text {i+1}** — ⏱️ {text_time:.1f}s", visible=True)
        result_audios[i] = gr.update(value=_cdn if _cdn else str(output_path), visible=True)

        # Yield progressive results
        yield tuple(result_labels + result_audios)


# UI Helper Functions

def load_speaker_metadata(speaker_id):
    """Load metadata for a speaker from any of their voice folders."""
    if not EARS_PATH.exists():
        return None
    
    # Find any subfolder for this speaker and load its metadata
    for subdir in EARS_PATH.iterdir():
        if subdir.is_dir() and subdir.name.startswith(f"{speaker_id}_"):
            metadata_path = subdir / "metadata.json"
            if metadata_path.exists():
                try:
                    with open(metadata_path, 'r') as f:
                        data = json.load(f)
                        return data.get("speaker_metadata", {})
                except Exception:
                    continue
    return None


def get_speakers():
    """Get list of unique speakers with their metadata."""
    if not EARS_PATH.exists():
        return []
    
    speakers_dict = {}
    for subdir in sorted(EARS_PATH.iterdir()):
        if subdir.is_dir():
            # Extract speaker ID (pXXX)
            name = subdir.name
            if name.startswith('p') and '_' in name:
                speaker_id = name.split('_')[0]
                if speaker_id not in speakers_dict:
                    speakers_dict[speaker_id] = None
    
    # Load metadata for each speaker
    speakers_with_metadata = []
    for speaker_id in sorted(speakers_dict.keys()):
        metadata = load_speaker_metadata(speaker_id)
        if metadata:
            speakers_with_metadata.append({
                'id': speaker_id,
                'gender': metadata.get('gender', 'unknown'),
                'age': metadata.get('age', 'unknown'),
                'ethnicity': metadata.get('ethnicity', 'unknown'),
                'native_language': metadata.get('native language', 'unknown'),
            })
        else:
            speakers_with_metadata.append({
                'id': speaker_id,
                'gender': 'unknown',
                'age': 'unknown',
                'ethnicity': 'unknown',
                'native_language': 'unknown',
            })
    
    return speakers_with_metadata


def get_speakers_table(search_query=""):
    """Get speakers as table data for Gradio, optionally filtered by search query."""
    speakers = get_speakers()
    result = []
    for s in speakers:
        # Abbreviate gender
        gender = s['gender']
        if gender.lower() == 'male':
            gender = 'M'
        elif gender.lower() == 'female':
            gender = 'F'
        else:
            gender = gender[0].upper() if gender else '?'
        
        # Apply search filter if provided
        if search_query:
            search_lower = search_query.lower()
            searchable_text = f"{s['id']} {gender} {s['age']} {s['ethnicity']} {s['native_language']}".lower()
            if search_lower not in searchable_text:
                continue
        
        result.append([s['id'], gender, s['age'], s['ethnicity'], s['native_language']])
    return result


def get_audio_length_from_metadata(voice_dir):
    """Get audio length from metadata.json file."""
    metadata_path = voice_dir / "metadata.json"
    if metadata_path.exists():
        try:
            with open(metadata_path, 'r') as f:
                data = json.load(f)
                length = data.get("audio_length_seconds", 0)
                return f"{length:.1f}s"
        except Exception:
            return "N/A"
    return "N/A"


def get_freeform_table(speaker_id):
    """Get freeform table for a speaker (single row if exists)."""
    if not EARS_PATH.exists() or not speaker_id:
        return []
    
    freeform_dir = EARS_PATH / f"{speaker_id}_freeform"
    if freeform_dir.exists():
        audio_path = freeform_dir / "audio.mp3"
        st_path = freeform_dir / "speaker_latent.safetensors"
        if audio_path.exists() and st_path.exists():
            audio_length = get_audio_length_from_metadata(freeform_dir)
            return [["Freeform", audio_length]]
    return []


def get_emotions_for_speaker(speaker_id):
    """Get list of emotions with audio lengths available for a given speaker (excluding _joint_)."""
    if not EARS_PATH.exists() or not speaker_id:
        return []
    emotions = []
    for subdir in sorted(EARS_PATH.iterdir()):
        if subdir.is_dir():
            name = subdir.name
            # Match pattern: p{speaker_id}_emo_{emotion} (but not _emo_joint_)
            if name.startswith(f"{speaker_id}_emo_") and "_joint_" not in name:
                # Extract emotion part
                parts = name.split('_emo_')
                if len(parts) == 2:
                    emotion = parts[1]
                    # Verify files exist
                    audio_path = subdir / "audio.mp3"
                    st_path = subdir / "speaker_latent.safetensors"
                    if audio_path.exists() and st_path.exists():
                        audio_length = get_audio_length_from_metadata(subdir)
                        emotions.append((emotion, audio_length))
    return emotions


def get_emotions_table(speaker_id):
    """Get emotions table for a speaker with audio lengths."""
    if not speaker_id:
        return []
    emotions = get_emotions_for_speaker(speaker_id)
    return [[emotion, length] for emotion, length in emotions]


# VCTK Helper Functions

def get_vctk_speakers():
    """Get list of VCTK speakers with their metadata."""
    if not VCTK_PATH.exists():
        return []
    
    speakers_with_metadata = []
    for subdir in sorted(VCTK_PATH.iterdir()):
        if subdir.is_dir() and subdir.name.startswith('p'):
            speaker_id = subdir.name
            audio_path = subdir / "audio.mp3"
            st_path = subdir / "speaker_latent.safetensors"
            metadata_path = subdir / "metadata.json"
            
            if audio_path.exists() and st_path.exists() and metadata_path.exists():
                try:
                    with open(metadata_path, 'r') as f:
                        data = json.load(f)
                        speaker_info = data.get("speaker_info", {})
                        audio_length = data.get("total_audio_length_seconds", 0)
                        
                        speakers_with_metadata.append({
                            'id': speaker_info.get('id', speaker_id),
                            'gender': speaker_info.get('gender', 'unknown'),
                            'age': speaker_info.get('age', 'unknown'),
                            'details': speaker_info.get('details', 'unknown'),
                            'audio_length': f"{audio_length:.1f}s"
                        })
                except Exception:
                    continue
    
    return speakers_with_metadata


def get_vctk_speakers_table(search_query=""):
    """Get VCTK speakers as table data for Gradio, optionally filtered by search query."""
    speakers = get_vctk_speakers()
    result = []
    for s in speakers:
        # Abbreviate gender
        gender = s['gender']
        if gender.lower() == 'male' or gender == 'M':
            gender = 'M'
        elif gender.lower() == 'female' or gender == 'F':
            gender = 'F'
        else:
            gender = gender[0].upper() if gender else '?'
        
        # Apply search filter if provided
        if search_query:
            search_lower = search_query.lower()
            searchable_text = f"{s['id']} {gender} {s['age']} {s['details']} {s['audio_length']}".lower()
            if search_lower not in searchable_text:
                continue
        
        result.append([s['id'], gender, s['age'], s['details'], s['audio_length']])
    return result


def load_text_presets():
    """Load text presets from file with category and word count."""
    if TEXT_PRESETS_PATH.exists():
        with open(TEXT_PRESETS_PATH, 'r', encoding='utf-8') as f:
            lines = [line.strip() for line in f if line.strip()]
        
        result = []
        for line in lines:
            # Split on first " | " to separate category from text
            if " | " in line:
                parts = line.split(" | ", 1)
                category = parts[0]
                text = parts[1]
            else:
                # Fallback if no category
                category = "Uncategorized"
                text = line
            
            # Calculate word count
            word_count = len(text.split())
            
            result.append([category, str(word_count), text])
        
        return result
    return []


def search_speakers(search_query):
    """Filter speakers table based on search query."""
    filtered_data = get_speakers_table(search_query)
    return gr.update(value=filtered_data)


def select_speaker_from_table(evt: gr.SelectData, table_data):
    """Handle speaker selection - populate freeform and emotions tables."""
    if evt.value and table_data is not None:
        # evt.index is a tuple/list (row, col), we need the row to get the speaker ID
        if isinstance(evt.index, (tuple, list)) and len(evt.index) >= 2:
            row_index = evt.index[0]
        else:
            row_index = evt.index
        
        # Use the actual displayed (filtered) table data (pandas DataFrame)
        if isinstance(row_index, int) and row_index < len(table_data):
            speaker_row = table_data.iloc[row_index]
            speaker_id = speaker_row.iloc[0]  # First column is the ID
            
            # Format selection display - clean and simple
            gender_full = "Male" if speaker_row.iloc[1] == "M" else "Female" if speaker_row.iloc[1] == "F" else speaker_row.iloc[1]
            selection_text = f"Selected Speaker: {speaker_id}\n{gender_full} • {speaker_row.iloc[2]} • {speaker_row.iloc[3]}"
            
            # Get freeform and emotions data
            freeform_data = get_freeform_table(speaker_id)
            emotions_data = get_emotions_table(speaker_id)
            
            return (
                gr.update(value=selection_text, visible=True),  # Show speaker selection
                gr.update(value=freeform_data, visible=True),  # Update freeform table
                gr.update(value=emotions_data, visible=True),  # Update emotions table
                gr.update(value=speaker_id),  # Store speaker ID
                gr.update(value=None),  # Clear audio preview
                gr.update(value=""),  # Clear safetensors path
                gr.update(value=""),   # Clear audio path
                gr.update(value="", visible=False)  # Clear voice selection display
            )
    return (
        gr.update(value="", visible=False),
        gr.update(value=[], visible=True),
        gr.update(value=[], visible=True),
        gr.update(value=""),
        gr.update(value=None),
        gr.update(value=""),
        gr.update(value=""),
        gr.update(value="", visible=False)
    )


def select_freeform_from_table(evt: gr.SelectData, speaker_id: str):
    """Handle freeform selection from table - load freeform voice files."""
    if speaker_id:
        voice_name = f"{speaker_id}_freeform"
        voice_dir = EARS_PATH / voice_name
        audio_path = str(voice_dir / "audio.mp3")
        st_path = str(voice_dir / "speaker_latent.safetensors")
        
        if voice_dir.exists():
            # Format freeform display
            freeform_display = f"Selected: Freeform\n{speaker_id}_freeform"
            
            return (
                gr.update(value=freeform_display, visible=True),  # Show freeform selection
                gr.update(value=audio_path),  # Update audio player
                gr.update(value=st_path),      # Update safetensors path
                gr.update(value=audio_path)    # Update audio path for reconstruction
            )
    return gr.update(value="", visible=False), gr.update(value=None), gr.update(value=""), gr.update(value="")


def select_emotion_from_table(evt: gr.SelectData, speaker_id: str):
    """Handle emotion selection - load voice files."""
    if evt.value and speaker_id:
        # evt.index is (row, col) - get the row to extract emotion from first column
        if isinstance(evt.index, (tuple, list)) and len(evt.index) >= 2:
            row_index = evt.index[0]
        else:
            row_index = 0
        
        # Get emotions data and extract the emotion name from first column
        emotions_data = get_emotions_table(speaker_id)
        if isinstance(row_index, int) and row_index < len(emotions_data):
            emotion = emotions_data[row_index][0]  # First column is emotion name
            
            voice_name = f"{speaker_id}_emo_{emotion}"
            voice_dir = EARS_PATH / voice_name
            audio_path = str(voice_dir / "audio.mp3")
            st_path = str(voice_dir / "speaker_latent.safetensors")
            
            if voice_dir.exists():
                # Format emotion display - clean and simple
                emotion_display = f"Selected Emotion: {emotion.title()}\n{speaker_id}_emo_{emotion}"
                
                return (
                    gr.update(value=emotion_display, visible=True),  # Show emotion selection
                    gr.update(value=audio_path),  # Update audio player
                    gr.update(value=st_path),      # Update safetensors path
                    gr.update(value=audio_path)    # Update audio path for reconstruction
                )
    return gr.update(value="", visible=False), gr.update(value=None), gr.update(value=""), gr.update(value="")


def select_vctk_speaker_from_table(evt: gr.SelectData, table_data):
    """Handle VCTK speaker selection - load voice files directly."""
    if evt.value and table_data is not None:
        # evt.index is a tuple/list (row, col), we need the row to get the speaker ID
        if isinstance(evt.index, (tuple, list)) and len(evt.index) >= 2:
            row_index = evt.index[0]
        else:
            row_index = evt.index
        
        # Use the actual displayed (filtered) table data (pandas DataFrame)
        if isinstance(row_index, int) and row_index < len(table_data):
            speaker_row = table_data.iloc[row_index]
            speaker_id = speaker_row.iloc[0]  # First column is the ID
            
            # Load voice files from VCTK
            voice_dir = VCTK_PATH / speaker_id
            audio_path = str(voice_dir / "audio.mp3")
            st_path = str(voice_dir / "speaker_latent.safetensors")
            
            if voice_dir.exists():
                # Format selection display
                gender_full = "Male" if speaker_row.iloc[1] == "M" else "Female" if speaker_row.iloc[1] == "F" else speaker_row.iloc[1]
                selection_text = f"Selected Speaker: {speaker_id}\n{gender_full} • {speaker_row.iloc[2]} • {speaker_row.iloc[3]}"
                
                return (
                    gr.update(value=selection_text, visible=True),  # Show speaker selection
                    gr.update(value=speaker_id),  # Store speaker ID
                    gr.update(value=audio_path),  # Update audio player
                    gr.update(value=st_path),      # Update safetensors path
                    gr.update(value=audio_path)    # Update audio path for reconstruction
                )
    return (
        gr.update(value="", visible=False),
        gr.update(value=""),
        gr.update(value=None),
        gr.update(value=""),
        gr.update(value="")
    )


def search_vctk_speakers(search_query):
    """Filter VCTK speakers table based on search query."""
    filtered_data = get_vctk_speakers_table(search_query)
    return gr.update(value=filtered_data)


# Expresso Helper Functions

def get_expresso_speakers():
    """Get list of all Expresso speakers with their metadata."""
    if not EXPRESSO_PATH.exists():
        return []
    
    speakers_with_metadata = []
    for subdir in sorted(EXPRESSO_PATH.iterdir()):
        if subdir.is_dir() and subdir.name.startswith('expresso_'):
            speaker_id = subdir.name
            audio_path = subdir / "audio.mp3"
            st_path = subdir / "speaker_latent.safetensors"
            metadata_path = subdir / "metadata.json"
            
            if audio_path.exists() and st_path.exists() and metadata_path.exists():
                try:
                    with open(metadata_path, 'r') as f:
                        data = json.load(f)
                        audio_length = data.get("audio_length_seconds", 0)
                        
                        speakers_with_metadata.append({
                            'id': speaker_id,
                            'type': data.get('type', 'unknown'),
                            'speakers': data.get('speakers', 'unknown'),
                            'style': data.get('style', 'unknown'),
                            'audio_length': f"{audio_length:.1f}s"
                        })
                except Exception:
                    continue
    
    return speakers_with_metadata


def get_expresso_speakers_table(search_query=""):
    """Get Expresso speakers as table data for Gradio, optionally filtered by search query."""
    speakers = get_expresso_speakers()
    result = []
    for s in speakers:
        # Apply search filter if provided
        if search_query:
            search_lower = search_query.lower()
            # Search in all fields
            if not any(search_lower in str(v).lower() for v in [s['id'], s['type'], s['speakers'], s['style']]):
                continue
        
        result.append([
            s['id'],
            s['type'],
            s['speakers'],
            s['style'],
            s['audio_length']
        ])
    
    return result


def select_expresso_speaker_from_table(evt: gr.SelectData, table_data):
    """Handle Expresso speaker selection - load voice files directly."""
    if evt.value and table_data is not None:
        # evt.index is a tuple/list (row, col), we need the row to get the speaker ID
        if isinstance(evt.index, (tuple, list)) and len(evt.index) >= 2:
            row_index = evt.index[0]
        else:
            row_index = evt.index
        
        # Use the actual displayed (filtered) table data (pandas DataFrame)
        if isinstance(row_index, int) and row_index < len(table_data):
            speaker_row = table_data.iloc[row_index]
            speaker_id = speaker_row.iloc[0]  # First column is the ID
            
            # Load voice files from Expresso
            voice_dir = EXPRESSO_PATH / speaker_id
            audio_path = str(voice_dir / "audio.mp3")
            st_path = str(voice_dir / "speaker_latent.safetensors")
            
            if voice_dir.exists():
                # Format selection display
                selection_text = f"Selected Voice: {speaker_id}\nType: {speaker_row.iloc[1]} • Speakers: {speaker_row.iloc[2]} • Style: {speaker_row.iloc[3]}"
                
                return (
                    gr.update(value=selection_text, visible=True),  # Show speaker selection
                    gr.update(value=speaker_id),  # Store speaker ID
                    gr.update(value=audio_path),  # Update audio player
                    gr.update(value=st_path),      # Update safetensors path
                    gr.update(value=audio_path)    # Update audio path for reconstruction
                )
    return (
        gr.update(value="", visible=False),
        gr.update(value=""),
        gr.update(value=None),
        gr.update(value=""),
        gr.update(value="")
    )


def search_expresso_speakers(search_query):
    """Filter Expresso speakers table based on search query."""
    filtered_data = get_expresso_speakers_table(search_query)
    return gr.update(value=filtered_data)


# HF-Custom Helper Functions

def get_hf_custom_speakers():
    """Get list of all HF-Custom speakers with their metadata."""
    if not HF_CUSTOM_PATH.exists():
        return []
    
    speakers_with_metadata = []
    for subdir in sorted(HF_CUSTOM_PATH.iterdir()):
        if subdir.is_dir():
            speaker_name = subdir.name
            audio_path = subdir / "audio.mp3"
            st_path = subdir / "speaker_latent.safetensors"
            metadata_path = subdir / "metadata.json"
            
            if audio_path.exists() and st_path.exists() and metadata_path.exists():
                try:
                    with open(metadata_path, 'r') as f:
                        data = json.load(f)
                        audio_length = data.get("audio_duration_seconds", 0)
                        
                        speakers_with_metadata.append({
                            'name': data.get('speaker_name', speaker_name),
                            'dataset': data.get('dataset_name', ''),
                            'description': data.get('speaker_description', ''),
                            'audio_length': f"{audio_length:.1f}s"
                        })
                except Exception:
                    continue
    
    return speakers_with_metadata


def get_hf_custom_speakers_table(search_query=""):
    """Get HF-Custom speakers as table data for Gradio, optionally filtered by search query."""
    speakers = get_hf_custom_speakers()
    result = []
    for s in speakers:
        # Apply search filter if provided
        if search_query:
            search_lower = search_query.lower()
            # Search in all fields
            if not any(search_lower in str(v).lower() for v in [s['name'], s['dataset'], s['description']]):
                continue
        
        result.append([
            s['name'],
            s['dataset'],
            s['description'],
            s['audio_length']
        ])
    
    return result


def select_hf_custom_speaker_from_table(evt: gr.SelectData, table_data):
    """Handle HF-Custom speaker selection - load voice files directly."""
    if evt.value and table_data is not None:
        # evt.index is a tuple/list (row, col), we need the row to get the speaker name
        if isinstance(evt.index, (tuple, list)) and len(evt.index) >= 2:
            row_index = evt.index[0]
        else:
            row_index = evt.index
        
        # Use the actual displayed (filtered) table data (pandas DataFrame)
        if isinstance(row_index, int) and row_index < len(table_data):
            speaker_row = table_data.iloc[row_index]
            speaker_name = speaker_row.iloc[0]  # First column is the name
            
            # Load voice files from HF-Custom
            voice_dir = HF_CUSTOM_PATH / speaker_name
            audio_path = str(voice_dir / "audio.mp3")
            st_path = str(voice_dir / "speaker_latent.safetensors")
            
            if voice_dir.exists():
                # Format selection display
                dataset_info = f" • {speaker_row.iloc[1]}" if speaker_row.iloc[1] else ""
                selection_text = f"Selected Voice: {speaker_name}{dataset_info}\n{speaker_row.iloc[2]}"
                
                return (
                    gr.update(value=selection_text, visible=True),  # Show speaker selection
                    gr.update(value=speaker_name),  # Store speaker name
                    gr.update(value=audio_path),  # Update audio player
                    gr.update(value=st_path),      # Update safetensors path
                    gr.update(value=audio_path)    # Update audio path for reconstruction
                )
    return (
        gr.update(value="", visible=False),
        gr.update(value=""),
        gr.update(value=None),
        gr.update(value=""),
        gr.update(value="")
    )


def search_hf_custom_speakers(search_query):
    """Filter HF-Custom speakers table based on search query."""
    filtered_data = get_hf_custom_speakers_table(search_query)
    return gr.update(value=filtered_data)


# Audio Prompt Library functions
AUDIO_EXTS = {".wav", ".mp3", ".m4a", ".ogg", ".flac", ".webm", ".aac", ".opus"}

def get_audio_prompt_files():
    """Get list of audio files from the audio prompt folder."""
    if AUDIO_PROMPT_FOLDER is None or not AUDIO_PROMPT_FOLDER.exists():
        return []
    
    files = sorted([
        f.name for f in AUDIO_PROMPT_FOLDER.iterdir() 
        if f.is_file() and f.suffix.lower() in AUDIO_EXTS
    ], key=str.lower)
    
    return [[file] for file in files]


def select_audio_prompt_file(evt: gr.SelectData):
    """Handle audio prompt file selection from table."""
    if evt.value and AUDIO_PROMPT_FOLDER is not None:
        file_path = AUDIO_PROMPT_FOLDER / evt.value
        if file_path.exists():
            return gr.update(value=str(file_path))
    return gr.update()


def switch_dataset(dataset_name):
    """Switch between Custom Audio Panel, EARS, VCTK, Expresso, and HF-Custom datasets."""
    if dataset_name == "Custom Audio Panel":
        # Show Custom Audio Panel only, hide all voicebank UI
        return (
            gr.update(value="", visible=False),  # dataset_license_info
            gr.update(visible=True),  # custom_audio_row
            gr.update(visible=False),  # voicebank_row
            gr.update(visible=False),  # voice_type_column
            gr.update(visible=True),  # ears_column (within voicebank_row)
            gr.update(visible=False),  # vctk_column
            gr.update(visible=False),  # expresso_column
            gr.update(visible=False),  # hf_custom_column
            # Clear selections
            gr.update(value="", visible=False),  # selected_speaker_display
            gr.update(value=[]),  # freeform_table
            gr.update(value=[]),  # emotions_table
            gr.update(value="", visible=False),  # selected_voice_display
            gr.update(value="", visible=False),  # vctk_speaker_display
            gr.update(value="", visible=False),  # expresso_speaker_display
            gr.update(value="", visible=False),  # hf_custom_speaker_display
            gr.update(value=""),  # selected_speaker_state
            gr.update(value=None),  # audio_preview
            gr.update(value=""),  # speaker_st_path_state
            gr.update(value="")  # speaker_audio_path_state
        )
    elif dataset_name == "EARS":
        # Show EARS UI, hide others, show Voice Type column
        license_text = "**EARS Dataset License:** Creative Commons Attribution 4.0 International ([CC-BY-4.0](https://creativecommons.org/licenses/by/4.0/))"
        return (
            gr.update(value=license_text, visible=True),  # dataset_license_info
            gr.update(visible=False),  # custom_audio_row
            gr.update(visible=True),  # voicebank_row
            gr.update(visible=True),  # voice_type_column (show for EARS)
            gr.update(visible=True),  # ears_column
            gr.update(visible=False),  # vctk_column
            gr.update(visible=False),  # expresso_column
            gr.update(visible=False),  # hf_custom_column
            gr.update(value=""),  # selected_speaker_display
            gr.update(value=[], visible=True),  # freeform_table
            gr.update(value=[], visible=True),  # emotions_table
            gr.update(value="", visible=False),  # selected_voice_display
            gr.update(value="", visible=False),  # vctk_speaker_display
            gr.update(value="", visible=False),  # expresso_speaker_display
            gr.update(value="", visible=False),  # hf_custom_speaker_display
            gr.update(value=""),  # selected_speaker_state
            gr.update(value=None),  # audio_preview
            gr.update(value=""),  # speaker_st_path_state
            gr.update(value="")   # speaker_audio_path_state
        )
    elif dataset_name == "VCTK":
        # Show VCTK UI, hide others, hide Voice Type column
        license_text = "**VCTK Dataset License:** Creative Commons Attribution 4.0 International ([CC-BY-4.0](https://creativecommons.org/licenses/by/4.0/))"
        return (
            gr.update(value=license_text, visible=True),  # dataset_license_info
            gr.update(visible=False),  # custom_audio_row
            gr.update(visible=True),  # voicebank_row
            gr.update(visible=False),  # voice_type_column
            gr.update(visible=False),  # ears_column
            gr.update(visible=True),  # vctk_column
            gr.update(visible=False),  # expresso_column
            gr.update(visible=False),  # hf_custom_column (hide for VCTK)
            gr.update(value=""),  # selected_speaker_display
            gr.update(value=[], visible=True),  # freeform_table
            gr.update(value=[], visible=True),  # emotions_table
            gr.update(value="", visible=False),  # selected_voice_display
            gr.update(value="", visible=False),  # vctk_speaker_display
            gr.update(value="", visible=False),  # expresso_speaker_display
            gr.update(value="", visible=False),  # hf_custom_speaker_display
            gr.update(value=""),  # selected_speaker_state
            gr.update(value=None),  # audio_preview
            gr.update(value=""),  # speaker_st_path_state
            gr.update(value="")   # speaker_audio_path_state
        )
    elif dataset_name == "Expresso":
        # Show Expresso UI, hide others, hide Voice Type column
        license_text = "**Expresso Dataset License:** Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International ([CC-BY-NC-SA-4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/))"
        return (
            gr.update(value=license_text, visible=True),  # dataset_license_info
            gr.update(visible=False),  # custom_audio_row
            gr.update(visible=True),  # voicebank_row
            gr.update(visible=False),  # voice_type_column
            gr.update(visible=False),  # ears_column
            gr.update(visible=False),  # vctk_column
            gr.update(visible=True),  # expresso_column
            gr.update(visible=False),  # hf_custom_column (hide for Expresso)
            gr.update(value=""),  # selected_speaker_display
            gr.update(value=[], visible=True),  # freeform_table
            gr.update(value=[], visible=True),  # emotions_table
            gr.update(value="", visible=False),  # selected_voice_display
            gr.update(value="", visible=False),  # vctk_speaker_display
            gr.update(value="", visible=False),  # expresso_speaker_display
            gr.update(value="", visible=False),  # hf_custom_speaker_display
            gr.update(value=""),  # selected_speaker_state
            gr.update(value=None),  # audio_preview
            gr.update(value=""),  # speaker_st_path_state
            gr.update(value="")   # speaker_audio_path_state
        )
    else:  # HF-Custom
        # Show HF-Custom UI, hide others, hide Voice Type column
        license_text = "**HF-Custom Voices:** Available in dataset cache (information in metadata.json per voice). Also view dataset at [jordand/echo-embeddings-custom](https://huggingface.co/datasets/jordand/echo-embeddings-custom)"
        return (
            gr.update(value=license_text, visible=True),  # dataset_license_info
            gr.update(visible=False),  # custom_audio_row
            gr.update(visible=True),  # voicebank_row
            gr.update(visible=False),  # voice_type_column
            gr.update(visible=False),  # ears_column
            gr.update(visible=False),  # vctk_column
            gr.update(visible=False),  # expresso_column
            gr.update(visible=True),  # hf_custom_column
            gr.update(value=""),  # selected_speaker_display
            gr.update(value=[], visible=True),  # freeform_table
            gr.update(value=[], visible=True),  # emotions_table
            gr.update(value="", visible=False),  # selected_voice_display
            gr.update(value="", visible=False),  # vctk_speaker_display
            gr.update(value="", visible=False),  # expresso_speaker_display
            gr.update(value="", visible=False),  # hf_custom_speaker_display
            gr.update(value=""),  # selected_speaker_state
            gr.update(value=None),  # audio_preview
            gr.update(value=""),  # speaker_st_path_state
            gr.update(value="")   # speaker_audio_path_state
        )


def select_text_preset(evt: gr.SelectData):
    """Handle text preset selection - extract text from the row."""
    if evt.value:
        # Get the row index from the selected cell
        if isinstance(evt.index, (tuple, list)) and len(evt.index) >= 2:
            row_index = evt.index[0]
        else:
            row_index = evt.index
        
        # Get all presets and extract the text (column 2) from the selected row
        presets_data = load_text_presets()
        if isinstance(row_index, int) and row_index < len(presets_data):
            text = presets_data[row_index][2]  # Column 2 is the text
            return gr.update(value=text)
    return gr.update()


def update_cfg_visibility(cfg_mode):
    """Update visibility of CFG parameters based on selected mode."""
    if cfg_mode == "joint-unconditional":
        return (
            gr.update(label="Text/Speaker CFG Scale", info="Guidance strength for text and speaker (joint)"),
            gr.update(visible=False),
            gr.update(visible=False)
        )
    elif cfg_mode == "apg-independent":
        return (
            gr.update(label="Text CFG Scale", info="Guidance strength for text"),
            gr.update(visible=True),
            gr.update(visible=True)
        )
    else:  # independent or alternating
        return (
            gr.update(label="Text CFG Scale", info="Guidance strength for text"),
            gr.update(visible=True),
            gr.update(visible=False)
        )


def toggle_speaker_k_fields(enabled):
    """Toggle visibility of speaker K row. Hidden components preserve their values automatically."""
    return gr.update(visible=enabled)


def toggle_custom_shapes_fields(enabled):
    """Toggle visibility of custom shapes row and reset to defaults if disabled."""
    if enabled:
        return gr.update(visible=True)
    else:
        # When disabled, hide the row and reset fields to defaults
        return gr.update(visible=False)


def toggle_mode(mode, speaker_k_enable_val, speaker_kv_simple_val):
    """Toggle between simple and advanced modes and sync speaker KV state."""
    if mode == "Simple Mode":
        # Sync simple checkbox with advanced mode's speaker_k_enable value
        return (
            gr.update(visible=True),  # simple_mode_row (speaker KV checkbox)
            gr.update(visible=False),  # advanced_mode_compile_column
            gr.update(visible=False),  # advanced_mode_column (all other parameters)
            gr.update(value=speaker_k_enable_val),  # sync simple checkbox with advanced
            gr.update(value=speaker_k_enable_val),  # also update speaker_k_enable (keep same)
        )
    else:  # Advanced Mode
        # Sync advanced mode's speaker_k_enable with simple checkbox value
        return (
            gr.update(visible=False),  # simple_mode_row (speaker KV checkbox)
            gr.update(visible=True),  # advanced_mode_compile_column
            gr.update(visible=True),  # advanced_mode_column (all other parameters)
            gr.update(value=speaker_kv_simple_val),  # sync simple checkbox (keep same)
            gr.update(value=speaker_kv_simple_val),  # sync advanced with simple checkbox
        )


def sync_simple_to_advanced(simple_enabled):
    """Sync simple mode speaker KV checkbox to advanced mode controls."""
    if simple_enabled:
        return (
            gr.update(value=True),  # speaker_k_enable
            gr.update(visible=True),  # speaker_k_row
            gr.update(value=1.5),  # speaker_k_scale
            gr.update(value=0.9),  # speaker_k_min_t
            gr.update(value=24),  # speaker_k_max_layers
        )
    else:
        return (
            gr.update(value=False),  # speaker_k_enable
            gr.update(visible=False),  # speaker_k_row
            gr.update(),  # speaker_k_scale (no change)
            gr.update(),  # speaker_k_min_t (no change)
            gr.update(),  # speaker_k_max_layers (no change)
        )


def apply_core_preset(preset_name):
    """Apply core sampling parameters preset."""
    if preset_name == "default":
        return [
            gr.update(value=0),  # rng_seed
            gr.update(value=40),  # num_steps
            gr.update(value="independent"),  # cfg_mode
            gr.update(value="Custom"),  # Set main preset to Custom
        ]
    return [gr.update()] * 4


def apply_cfg_preset(preset_name):
    """Apply CFG guidance preset."""
    presets = {
        "default": (3.0, 5.0, 0.5, 1.0),
        "higher speaker": (3.0, 8.0, 0.5, 1.0),
        "large guidances": (8.0, 8.0, 0.5, 1.0),
    }
    
    if preset_name not in presets:
        return [gr.update()] * 5
    
    text_scale, speaker_scale, min_t, max_t = presets[preset_name]
    
    return [
        gr.update(value=text_scale),  # cfg_scale_text
        gr.update(value=speaker_scale),  # cfg_scale_speaker
        gr.update(value=min_t),  # cfg_min_t
        gr.update(value=max_t),  # cfg_max_t
        gr.update(value="Custom"),  # Set main preset to Custom
    ]


def apply_speaker_kv_preset(preset_name):
    """Apply speaker KV attention control preset."""
    if preset_name == "enable":
        return [
            gr.update(value=True),  # speaker_k_enable
            gr.update(visible=True),  # speaker_k_row
            gr.update(value="Custom"),  # Set main preset to Custom
        ]
    elif preset_name == "off":
        return [
            gr.update(value=False),  # speaker_k_enable
            gr.update(visible=False),  # speaker_k_row
            gr.update(value="Custom"),  # Set main preset to Custom
        ]
    return [gr.update()] * 3


def apply_truncation_preset(preset_name):
    """Apply truncation & temporal rescaling preset."""
    presets = {
        "flat": (0.8, 1.2, 3.0),
        "sharp": (0.9, 0.96, 3.0),
        "baseline(sharp)": (1.0, 1.0, 3.0),
    }
    
    if preset_name == "custom" or preset_name not in presets:
        return [gr.update()] * 4  # Return no changes for custom
    
    truncation, rescale_k, rescale_sigma = presets[preset_name]
    
    return [
        gr.update(value=truncation),
        gr.update(value=rescale_k),
        gr.update(value=rescale_sigma),
        gr.update(value="Custom"),  # Set main preset to Custom
    ]


def apply_apg_preset(preset_name):
    """Apply APG parameters preset."""
    presets = {
        "default": (0.5, 0.5, -0.25, -0.25, "", ""),  # default: -0.25 momentum
        "no momentum": (0.0, 0.0, 0.0, 0.0, "", ""),  # no momentum: 0 momentum
        "norms": (0.5, 0.5, -0.25, -0.25, "7.5", "7.5"),  # norms: default + 7.5 norms
        "no eta": (0.0, 0.0, -0.25, -0.25, "", ""),  # no eta: 0 eta
    }
    
    if preset_name not in presets:
        return [gr.update()] * 7
    
    eta_text, eta_speaker, momentum_text, momentum_speaker, norm_text, norm_speaker = presets[preset_name]
    
    return [
        gr.update(value=eta_text),  # apg_eta_text
        gr.update(value=eta_speaker),  # apg_eta_speaker
        gr.update(value=momentum_text),  # apg_momentum_text
        gr.update(value=momentum_speaker),  # apg_momentum_speaker
        gr.update(value=norm_text),  # apg_norm_text
        gr.update(value=norm_speaker),  # apg_norm_speaker
        gr.update(value="Custom"),  # Set main preset to Custom
    ]


def load_sampler_presets():
    """Load sampler presets from JSON file."""
    if SAMPLER_PRESETS_PATH.exists():
        with open(SAMPLER_PRESETS_PATH, 'r') as f:
            return json.load(f)
    else:
        # Create default presets (will use existing JSON file if it exists)
        default_presets = {
            "Flat (Independent)": {
                "num_steps": "30",
                "cfg_mode": "independent",
                "cfg_scale_text": "3.0",
                "cfg_scale_speaker": "5.0",
                "cfg_min_t": "0.5",
                "cfg_max_t": "1.0",
                "truncation_factor": "0.8",
                "rescale_k": "1.2",
                "rescale_sigma": "3.0"
            },
            "Sharp (Independent)": {
                "num_steps": "30",
                "cfg_mode": "independent",
                "cfg_scale_text": "3.0",
                "cfg_scale_speaker": "5.0",
                "cfg_min_t": "0.5",
                "cfg_max_t": "1.0",
                "truncation_factor": "0.9",
                "rescale_k": "0.96",
                "rescale_sigma": "3.0"
            },
        }
        with open(SAMPLER_PRESETS_PATH, 'w') as f:
            json.dump(default_presets, f, indent=2)
        return default_presets


def apply_sampler_preset(preset_name):
    """Apply a sampler preset to all fields."""
    presets = load_sampler_presets()
    if preset_name == "Custom" or preset_name not in presets:
        return [gr.update()] * 20  # Return no changes for custom
    
    preset = presets[preset_name]
    
    # Determine visibility based on cfg_mode
    cfg_mode_value = preset["cfg_mode"]
    speaker_visible = (cfg_mode_value != "joint-unconditional")
    apg_visible = (cfg_mode_value == "apg-independent")
    
    speaker_k_enabled = preset.get("speaker_k_enable", False)
    
    # Convert string values to numeric where appropriate
    def to_num(val, default):
        try:
            return float(val) if isinstance(val, str) else val
        except (ValueError, TypeError):
            return default
    
    return [
        gr.update(value=int(to_num(preset["num_steps"], 40))),
        gr.update(value=preset["cfg_mode"]),
        gr.update(value=to_num(preset["cfg_scale_text"], 3.0)),
        gr.update(value=to_num(preset["cfg_scale_speaker"], 5.0), visible=speaker_visible),
        gr.update(value=to_num(preset["cfg_min_t"], 0.5)),
        gr.update(value=to_num(preset["cfg_max_t"], 1.0)),
        gr.update(value=to_num(preset["truncation_factor"], 0.8)),
        gr.update(value=to_num(preset["rescale_k"], 1.2)),  # Now numeric
        gr.update(value=to_num(preset["rescale_sigma"], 3.0)),
        gr.update(value=speaker_k_enabled),
        gr.update(visible=speaker_k_enabled),  # speaker_k_row
        gr.update(value=to_num(preset.get("speaker_k_scale", "1.5"), 1.5)),
        gr.update(value=to_num(preset.get("speaker_k_min_t", "0.9"), 0.9)),
        gr.update(value=int(to_num(preset.get("speaker_k_max_layers", "24"), 24))),
        gr.update(value=to_num(preset.get("apg_eta_text", "0.0"), 0.0)),
        gr.update(value=to_num(preset.get("apg_eta_speaker", "0.0"), 0.0)),
        gr.update(value=to_num(preset.get("apg_momentum_text", "0.0"), 0.0)),
        gr.update(value=to_num(preset.get("apg_momentum_speaker", "0.0"), 0.0)),
        gr.update(value=preset.get("apg_norm_text", "")),  # Keep as string (can be empty)
        gr.update(value=preset.get("apg_norm_speaker", "")),  # Keep as string (can be empty)
    ]


# Build Gradio Interface
LINK_CSS = """
.preset-inline { display:flex; align-items:baseline; gap:6px; margin-top:-4px; margin-bottom:-12px; }
.preset-inline .title { font-weight:600; font-size:.95rem; }
.preset-inline .dim   { color:#666; margin:0 4px; }
/* blue, linky */
a.preset-link { color: #0a5bd8; text-decoration: underline; cursor: pointer; font-weight: 400; }
a.preset-link:hover { text-decoration: none; opacity: 0.8; }

/* Dark mode support for preset links */
.dark a.preset-link,
[data-theme="dark"] a.preset-link {
    color: #60a5fa !important;
}
.dark a.preset-link:hover,
[data-theme="dark"] a.preset-link:hover {
    color: #93c5fd !important;
}
.dark .preset-inline .dim,
[data-theme="dark"] .preset-inline .dim {
    color: #9ca3af !important;
}

/* keep proxy buttons in DOM but invisible */
.proxy-btn { position:absolute; width:0; height:0; overflow:hidden; padding:0 !important; margin:0 !important; border:0 !important; opacity:0; pointer-events:none; }

/* Better contrast for parameter group boxes */
.gr-group {
    border: 1px solid #d1d5db !important;
    background: #f3f4f6 !important;
}
.dark .gr-group,
[data-theme="dark"] .gr-group {
    border: 1px solid #4b5563 !important;
    background: #1f2937 !important;
}

/* Highlight generated audio */
.generated-audio-player {
    border: 3px solid #667eea !important;
    border-radius: 12px !important;
    padding: 20px !important;
    background: linear-gradient(135deg, rgba(102, 126, 234, 0.08) 0%, rgba(118, 75, 162, 0.05) 100%) !important;
    box-shadow: 0 4px 12px rgba(102, 126, 234, 0.2) !important;
    margin: 1rem 0 !important;
}
.generated-audio-player > div {
    background: transparent !important;
}

/* Make Parameter Mode selector more prominent */
#component-mode-selector {
    text-align: center;
    padding: 1rem 0;
}
#component-mode-selector label {
    font-size: 1.1rem !important;
    font-weight: 600 !important;
    margin-bottom: 0.5rem !important;
}
#component-mode-selector .wrap {
    justify-content: center !important;
}
#component-mode-selector fieldset {
    border: 2px solid #e5e7eb !important;
    border-radius: 8px !important;
    padding: 1rem !important;
    background: #f9fafb !important;
}
.dark #component-mode-selector fieldset,
[data-theme="dark"] #component-mode-selector fieldset {
    border: 2px solid #4b5563 !important;
    background: #1f2937 !important;
}

/* Stronger section separators */
.section-separator {
    height: 3px !important;
    background: linear-gradient(90deg, transparent 0%, #667eea 20%, #764ba2 80%, transparent 100%) !important;
    border: none !important;
    margin: 2rem 0 !important;
}
.dark .section-separator,
[data-theme="dark"] .section-separator {
    background: linear-gradient(90deg, transparent 0%, #667eea 20%, #764ba2 80%, transparent 100%) !important;
}

/* Section headers styling */
.gradio-container h1,
.gradio-container h2 {
    font-weight: 700 !important;
    margin-top: 1.5rem !important;
    margin-bottom: 1rem !important;
}

/* Highlighted tip box */
.tip-box {
    background: linear-gradient(135deg, #fef3c7 0%, #fde68a 100%) !important;
    border-left: 4px solid #f59e0b !important;
    border-radius: 8px !important;
    padding: 1rem 1.5rem !important;
    margin: 1rem 0 !important;
    box-shadow: 0 2px 4px rgba(245, 158, 11, 0.1) !important;
}
.tip-box strong {
    color: #92400e !important;
}
.dark .tip-box,
[data-theme="dark"] .tip-box {
    background: linear-gradient(135deg, #451a03 0%, #78350f 100%) !important;
    border-left: 4px solid #f59e0b !important;
}
.dark .tip-box strong,
[data-theme="dark"] .tip-box strong {
    color: #fbbf24 !important;
}
"""

JS_CODE = r"""
function () {
  // Get a queryable root, regardless of Shadow DOM
  const appEl = document.querySelector("gradio-app");
  const root  = appEl && appEl.shadowRoot ? appEl.shadowRoot : document;

  function clickHiddenButtonById(id) {
    if (!id) return;
    const host = root.getElementById(id);
    if (!host) return;
    const realBtn = host.querySelector("button, [role='button']") || host;
    realBtn.click();
  }

  // Delegate clicks from any <a class="preset-link" data-fire="...">
  root.addEventListener("click", (ev) => {
    const a = ev.target.closest("a.preset-link");
    if (!a) return;
    ev.preventDefault();
    ev.stopPropagation();
    ev.stopImmediatePropagation();
    clickHiddenButtonById(a.getAttribute("data-fire"));
    return false;
  }, true);
}
"""

def init_session():
    """Initialize session ID for this browser tab/session."""
    return secrets.token_hex(8)

def init_and_compile():
    """Initialize session and trigger compilation on page load."""
    session_id = secrets.token_hex(8)
    
    # Trigger compilation automatically on page load if not on Zero GPU
    # This ensures Simple mode (which defaults compile=True) gets compiled
    if not IS_ZEROGPU:
        # Just call do_compile directly - it will load models and compile
        # Status updates will be visible in Advanced mode, hidden in Simple mode
        status_update, checkbox_update = do_compile()
        return session_id, status_update, checkbox_update
    else:
        # On Zero GPU, don't try to compile
        return session_id, gr.update(), gr.update()

SIMPLE_CSS = """
.simple-container {
    max-width: 1200px;
    margin: 0 auto;
}
.simple-generate-btn {
    font-size: 1.2rem !important;
    padding: 1rem 2rem !important;
}
.simple-output-container {
    min-height: 200px;
    display: flex;
    flex-direction: column;
    justify-content: center;
}
"""


@spaces.GPU
def generate_unified(
    texts_input: str,
    voice_url: str,
    voice_file: str,
    num_steps: int,
    rng_seed: int,
    cfg_mode: str,
    cfg_scale_text: float,
    cfg_scale_speaker: float,
    cfg_min_t: float,
    cfg_max_t: float,
    truncation_factor: float,
    rescale_k: float,
    rescale_sigma: float,
    speaker_k_enable: bool,
    speaker_k_scale: float,
    speaker_k_min_t: float,
    speaker_k_max_layers: int,
    apg_eta_text: float,
    apg_eta_speaker: float,
    apg_momentum_text: float,
    apg_momentum_speaker: float,
    apg_norm_text: str,
    apg_norm_speaker: str,
    audio_format: str,
    session_id: str,
):
    """Unified generation function for the new simple UI."""

    global model, model_compiled
    active_model = model_compiled if model_compiled is not None else model

    # Cleanup old files
    cleanup_temp_audio(TEMP_AUDIO_DIR, session_id)

    # Parse texts
    texts = []
    if texts_input and texts_input.strip():
        for line in texts_input.strip().split('\n'):
            line = line.strip()
            if line:
                texts.append(line)
    texts = texts[:MAX_TEXTS]

    if not texts:
        return tuple([gr.update(visible=False)] * (1 + MAX_TEXTS * 2))

    # Determine speaker audio path (URL or file)
    speaker_audio_path = None
    if voice_url and voice_url.strip():
        try:
            speaker_audio_path = download_audio_from_url(voice_url.strip())
        except Exception as e:
            print(f"Error downloading URL: {e}")
            speaker_audio_path = None
    elif voice_file:
        speaker_audio_path = voice_file

    # Load speaker audio
    if speaker_audio_path:
        speaker_audio = load_audio(speaker_audio_path).cuda()
    else:
        speaker_audio = None

    # Parse parameters
    num_steps_int = int(num_steps)
    rng_seed_int = int(rng_seed)
    truncation_factor_val = float(truncation_factor) if truncation_factor != 1.0 else None
    rescale_k_val = float(rescale_k) if rescale_k != 0.0 else None
    rescale_sigma_val = float(rescale_sigma) if rescale_sigma != 1.0 else None

    speaker_k_scale_val = float(speaker_k_scale) if speaker_k_enable else None
    speaker_k_min_t_val = float(speaker_k_min_t) if speaker_k_enable else None
    speaker_k_max_layers_val = int(speaker_k_max_layers) if speaker_k_enable else None

    # APG parameters (only for apg-independent mode)
    if cfg_mode == "apg-independent":
        apg_eta_text_val = float(apg_eta_text)
        apg_eta_speaker_val = float(apg_eta_speaker)
        apg_momentum_text_val = float(apg_momentum_text)
        apg_momentum_speaker_val = float(apg_momentum_speaker)
        apg_norm_text_val = float(apg_norm_text) if apg_norm_text.strip() else None
        if apg_norm_text_val is not None and apg_norm_text_val <= 0:
            apg_norm_text_val = None
        apg_norm_speaker_val = float(apg_norm_speaker) if apg_norm_speaker.strip() else None
        if apg_norm_speaker_val is not None and apg_norm_speaker_val <= 0:
            apg_norm_speaker_val = None
    else:
        apg_eta_text_val = None
        apg_eta_speaker_val = None
        apg_momentum_text_val = None
        apg_momentum_speaker_val = None
        apg_norm_text_val = None
        apg_norm_speaker_val = None

    # Map cfg_mode to GuidanceMode
    guidance_mode_map = {
        "independent": GuidanceMode.INDEPENDENT,
        "joint-unconditional": GuidanceMode.JOINT,
        "alternating": GuidanceMode.ALTERNATING,
        "apg-independent": GuidanceMode.APG,
    }
    guidance_mode = guidance_mode_map.get(cfg_mode, GuidanceMode.INDEPENDENT)

    # Create sampler function
    sample_fn = partial(
        sample_euler_cfg_any,
        num_steps=num_steps_int,
        guidance_mode=guidance_mode,
        cfg_scale_text=float(cfg_scale_text),
        cfg_scale_speaker=float(cfg_scale_speaker),
        cfg_min_t=float(cfg_min_t),
        cfg_max_t=float(cfg_max_t),
        truncation_factor=truncation_factor_val,
        rescale_k=rescale_k_val,
        rescale_sigma=rescale_sigma_val,
        speaker_k_scale=speaker_k_scale_val,
        speaker_k_min_t=speaker_k_min_t_val,
        speaker_k_max_layers=speaker_k_max_layers_val,
        apg_eta_text=apg_eta_text_val,
        apg_eta_speaker=apg_eta_speaker_val,
        apg_momentum_text=apg_momentum_text_val,
        apg_momentum_speaker=apg_momentum_speaker_val,
        apg_norm_text=apg_norm_text_val,
        apg_norm_speaker=apg_norm_speaker_val,
        block_size=640
    )

    # Initialize outputs
    result_labels = [gr.update(visible=False)] * MAX_TEXTS
    result_audios = [gr.update(visible=False)] * MAX_TEXTS
    time_display = gr.update(visible=False)

    # Generate each text
    batch_start = time.time()
    for i, text in enumerate(texts):
        text_start = time.time()

        # Generate audio
        audio_out = sample_pipeline(
            model=active_model,
            fish_ae=fish_ae,
            pca_state=pca_state,
            sample_fn=sample_fn,
            text_prompt=text,
            speaker_audio=speaker_audio,
            rng_seed=rng_seed_int + i,
            pad_to_max_text_seq_len=768,
            pad_to_max_speaker_latent_len=2560,
        )

        text_duration = time.time() - text_start

        # Save audio
        audio_stem = make_stem(f"text_{i}", session_id)
        audio_path = save_audio_with_format(
            audio_out.cpu()[0],
            TEMP_AUDIO_DIR,
            audio_stem,
            44100,
            audio_format
        )

        # Update this output slot
        result_labels[i] = gr.update(
            value=f"**Text {i+1}** ({text_duration:.1f}s)\n\n> {text[:100]}{'...' if len(text) > 100 else ''}",
            visible=True
        )
        _cdn = upload_to_r2(audio_path)
        result_audios[i] = gr.update(value=_cdn if _cdn else str(audio_path), visible=True)

        # Yield progressive results
        yield tuple([time_display] + result_labels + result_audios)

    # Final time display
    batch_duration = time.time() - batch_start
    time_display = gr.update(
        value=f"✅ Generated {len(texts)} audio(s) in {batch_duration:.1f}s",
        visible=True
    )

    yield tuple([time_display] + result_labels + result_audios)


with gr.Blocks(title="Echo-TTS", css=LINK_CSS + SIMPLE_CSS, js=JS_CODE) as demo:
    gr.Markdown("# Echo-TTS")
    gr.Markdown("*Jordan Darefsky, 2025. See technical details [here](https://jordandarefsky.com/blog/2025/echo/). All audio outputs are subject to non-commercial use [CC-BY-NC-SA-4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/).*")

    # Session state for per-user file management
    session_id_state = gr.State(None)

    # ==================== UNIFIED UI ====================
    gr.Markdown("# Echo TTS")
    gr.Markdown("Fast, high-quality text-to-speech with voice cloning. Enter text (one per line for batch) and provide a voice reference.")

    with gr.Row():
        # LEFT COLUMN - Inputs
        with gr.Column(scale=1):
            # Voice Input
            gr.Markdown("### Voice Reference")
            voice_url = gr.Textbox(
                label="Voice URL (audio file URL - cached automatically)",
                placeholder="https://example.com/voice.wav",
                interactive=True
            )
            gr.Markdown("**OR**")
            voice_file = gr.Audio(
                sources=["upload", "microphone"],
                type="filepath",
                label="Upload Voice File",
                max_length=600
            )

            gr.Markdown("---")

            # Text Input
            gr.Markdown("### Text Prompts")
            texts_input = gr.Textbox(
                value="[S1] Hello, this is a test of Echo TTS.",
                lines=10,
                max_lines=30,
                label=f"Text Prompts (one per line, max {MAX_TEXTS})",
                placeholder="Enter your text prompts here, one per line...",
                interactive=True
            )

        # RIGHT COLUMN - Settings & Output
        with gr.Column(scale=1):
            # Settings Accordions
            with gr.Accordion("⚙️ Basic Settings", open=True):
                preset_dropdown = gr.Dropdown(
                    choices=list(load_sampler_presets().keys()),
                    value=list(load_sampler_presets().keys())[0],
                    label="Preset",
                    interactive=True
                )
                num_steps = gr.Slider(5, 50, value=10, step=1, label="Steps")
                rng_seed = gr.Number(value=42, label="Random Seed", precision=0)

            with gr.Accordion("🎚️ CFG Settings", open=False):
                cfg_mode = gr.Dropdown(
                    choices=["independent", "joint-unconditional", "alternating", "apg-independent"],
                    value="independent",
                    label="CFG Mode"
                )
                cfg_scale_text = gr.Slider(0, 5, value=3.0, step=0.1, label="Text Guidance")
                cfg_scale_speaker = gr.Slider(0, 5, value=1.0, step=0.1, label="Speaker Guidance")
                cfg_min_t = gr.Slider(0, 1, value=0.0, step=0.05, label="CFG Min T")
                cfg_max_t = gr.Slider(0, 1, value=1.0, step=0.05, label="CFG Max T")

            with gr.Accordion("🔧 Advanced Settings", open=False):
                truncation_factor = gr.Slider(0.5, 1.5, value=1.0, step=0.05, label="Truncation Factor")
                rescale_k = gr.Slider(0, 2, value=0.0, step=0.1, label="Rescale K")
                rescale_sigma = gr.Slider(0, 2, value=1.0, step=0.1, label="Rescale Sigma")

                speaker_k_enable = gr.Checkbox(value=False, label="Enable Speaker KV Scaling")
                with gr.Row(visible=False) as speaker_k_row:
                    speaker_k_scale = gr.Slider(0, 2, value=0.3, step=0.05, label="Speaker KV Scale")
                    speaker_k_min_t = gr.Slider(0, 1, value=0.7, step=0.05, label="Speaker KV Min T")
                    speaker_k_max_layers = gr.Slider(1, 24, value=12, step=1, label="Speaker KV Max Layers")

                # APG settings (only for apg-independent mode)
                with gr.Group(visible=False) as apg_group:
                    gr.Markdown("**APG Settings** (for apg-independent mode)")
                    apg_eta_text = gr.Slider(0, 1, value=0.5, step=0.05, label="APG Eta Text")
                    apg_eta_speaker = gr.Slider(0, 1, value=0.5, step=0.05, label="APG Eta Speaker")
                    apg_momentum_text = gr.Slider(0, 1, value=0.9, step=0.05, label="APG Momentum Text")
                    apg_momentum_speaker = gr.Slider(0, 1, value=0.9, step=0.05, label="APG Momentum Speaker")
                    apg_norm_text = gr.Textbox(value="", label="APG Norm Text (leave empty to disable)", placeholder="e.g. 1.0")
                    apg_norm_speaker = gr.Textbox(value="", label="APG Norm Speaker (leave empty to disable)", placeholder="e.g. 1.0")

                audio_format = gr.Radio(
                    choices=["wav", "flac"],
                    value="wav",
                    label="Output Format"
                )

            gr.Markdown("---")

            # Generate Button
            generate_btn = gr.Button("🎵 Generate Audio", variant="primary", size="lg")

            # Output Section
            gr.Markdown("### Generated Audio")
            generation_time_display = gr.Markdown(visible=False)

            # Pre-allocate output slots (up to MAX_TEXTS)
            output_labels = []
            output_audios = []
            for i in range(MAX_TEXTS):
                output_labels.append(gr.Markdown(visible=False))
                output_audios.append(gr.Textbox(visible=False, interactive=False, label=f"Audio URL {i+1}"))

    # Session state
    session_id_state = gr.State(None)

    # ==================== EVENT HANDLERS ====================

    # Toggle speaker KV row visibility
    speaker_k_enable.change(
        lambda enabled: gr.update(visible=enabled),
        inputs=[speaker_k_enable],
        outputs=[speaker_k_row]
    )

    # Toggle APG group visibility based on cfg_mode
    cfg_mode.change(
        lambda mode: gr.update(visible=(mode == "apg-independent")),
        inputs=[cfg_mode],
        outputs=[apg_group]
    )

    # Apply preset
    def apply_preset(preset_name):
        presets = load_sampler_presets()
        if preset_name in presets:
            p = presets[preset_name]
            return (
                p["num_steps"],
                p["cfg_mode"],
                p.get("cfg_scale_text", 3.0),
                p.get("cfg_scale_speaker", 1.0),
                p.get("cfg_min_t", 0.0),
                p.get("cfg_max_t", 1.0),
                p.get("truncation_factor", 1.0),
                p.get("rescale_k", 0.0),
                p.get("rescale_sigma", 1.0),
            )
        return tuple([gr.update()] * 9)

    preset_dropdown.change(
        apply_preset,
        inputs=[preset_dropdown],
        outputs=[num_steps, cfg_mode, cfg_scale_text, cfg_scale_speaker, cfg_min_t, cfg_max_t, truncation_factor, rescale_k, rescale_sigma]
    )

    # Generate button click
    generate_btn.click(
        generate_unified,
        inputs=[
            texts_input,
            voice_url,
            voice_file,
            num_steps,
            rng_seed,
            cfg_mode,
            cfg_scale_text,
            cfg_scale_speaker,
            cfg_min_t,
            cfg_max_t,
            truncation_factor,
            rescale_k,
            rescale_sigma,
            speaker_k_enable,
            speaker_k_scale,
            speaker_k_min_t,
            speaker_k_max_layers,
            apg_eta_text,
            apg_eta_speaker,
            apg_momentum_text,
            apg_momentum_speaker,
            apg_norm_text,
            apg_norm_speaker,
            audio_format,
            session_id_state,
        ],
        outputs=[generation_time_display] + output_labels + output_audios
    )

    # Initialize session
    demo.load(lambda: secrets.token_hex(8), outputs=[session_id_state])

# Enable queue for better handling of concurrent requests on HF Spaces
demo.queue(max_size=20)
demo.launch()