# 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 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()