import tempfile import time from pathlib import Path from typing import Optional, Tuple import spaces import gradio as gr import numpy as np import soundfile as sf import torch from dia.model import Dia # Load Nari model and config print("[AOS] Loading Nari model...") try: # Use the function from inference.py model = Dia.from_pretrained("nari-labs/Dia-1.6B-0626", compute_dtype="float16") print(f"[AOS] Model loaded successfully. Device: {model.device}") except Exception as e: print(f"[AOS] Error loading Nari model: {e}") raise # ────────────────────────────────────────────── # AOS Post-DSP: pure tensor math on output PCM # ────────────────────────────────────────────── def aos_process( pcm: np.ndarray, volume_gain: float, pitch_flattening: float, sample_rate: int, ) -> np.ndarray: """ Apply AOS mathematical clamping to the output PCM tensor. Args: pcm: float32 array in [-1.0, 1.0] volume_gain: 0.0-1.0 — peak amplitude clamp pitch_flattening: 0.0-0.99 — F0 variance attenuation sample_rate: audio sample rate (Hz) Returns: Processed float32 array """ t0 = time.perf_counter() # ---- 1. Volume Clamp (element-wise tensor clamp) ---- if volume_gain < 1.0: print(f"[AOS] Volume clamp: gain={volume_gain:.3f}, input range=[{pcm.min():.4f}, {pcm.max():.4f}]") # The clamp: multiply by gain, then hard clip to [-1.0, 1.0] pcm = np.clip(pcm * volume_gain, -1.0, 1.0) print(f"[AOS] output range=[{pcm.min():.4f}, {pcm.max():.4f}]") # ---- 2. Pitch Flattening (STFT-based spectral smoothing) ---- if pitch_flattening > 0.0: print(f"[AOS] Pitch flatten: flattening_index={pitch_flattening:.3f}, " f"PCM shape={pcm.shape}, sr={sample_rate}") # Convert to torch tensor for STFT device = torch.device("cpu") pcm_t = torch.from_numpy(pcm).float().to(device) n_fft = min(2048, len(pcm_t) // 8 or 512) hop_length = n_fft // 4 win_length = n_fft window = torch.hann_window(win_length).to(device) # Compute STFT: [channel, freq, time, 2] (real, imag) stft = torch.stft( pcm_t.unsqueeze(0), # add batch dim n_fft=n_fft, hop_length=hop_length, win_length=win_length, window=window, return_complex=True, # complex tensor ) # stft shape: [1, freq, time] print(f"[AOS] STFT shape={stft.shape}") # Magnitude and phase mag = stft.abs() # [1, freq, time] phase = stft.angle() # [1, freq, time] # Compute mean magnitude across time (per frequency bin) mag_mean = mag.mean(dim=-1, keepdim=True) # [1, freq, 1] print(f"[AOS] mag_mean range=[{mag_mean.min():.4f}, {mag_mean.max():.4f}]") # Attenuate deviation from mean: # mag_new = mag_mean + (1 - flattening) * (mag - mag_mean) flattening = max(0.0, min(pitch_flattening, 0.99)) mag_new = mag_mean + (1.0 - flattening) * (mag - mag_mean) print(f"[AOS] mag_new range=[{mag_new.min():.4f}, {mag_new.max():.4f}]") # Reconstruct complex spectrogram stft_new = mag_new * torch.exp(1j * phase) # ISTFT back to time domain pcm_out = torch.istft( stft_new, n_fft=n_fft, hop_length=hop_length, win_length=win_length, window=window, length=pcm_t.shape[-1], ) pcm = pcm_out.squeeze(0).numpy() print(f"[AOS] output range=[{pcm.min():.4f}, {pcm.max():.4f}]") elapsed = time.perf_counter() - t0 print(f"[AOS] aos_process took {elapsed*1000:.1f}ms") return pcm @spaces.GPU(duration=100) def run_inference( text_input: str, audio_prompt_input: Optional[Tuple[int, np.ndarray]], transcription_input: Optional[str], max_new_tokens: int, cfg_scale: float, temperature: float, top_p: float, cfg_filter_top_k: int, speed_factor: float, # ── AOS ACQUSTIC PARAMETERS ── volume_gain: float = 1.0, pitch_flattening: float = 0.0, ): """ Runs Nari inference using the globally loaded model and provided inputs. Uses temporary files for text and audio prompt compatibility with inference.generate. AOS extension: added volume_gain (post-DSP clamp) and pitch_flattening (STFT-based F0 variance attenuation) as pure tensor math operations on the output PCM. """ # ── AOS DEBUG: log all inputs ── print(f"[AOS] run_inference called") print(f"[AOS] text_input={text_input[:100]}{'...' if len(text_input)>100 else ''}") print(f"[AOS] max_new_tokens={max_new_tokens}, cfg_scale={cfg_scale}, temperature={temperature}") print(f"[AOS] top_p={top_p}, cfg_filter_top_k={cfg_filter_top_k}, speed_factor={speed_factor}") print(f"[AOS] AOS params: volume_gain={volume_gain}, pitch_flattening={pitch_flattening}") if not text_input or text_input.isspace(): raise gr.Error("Text input cannot be empty.") temp_txt_file_path = None temp_audio_prompt_path = None output_audio = (44100, np.zeros(1, dtype=np.float32)) try: prompt_path_for_generate = None if audio_prompt_input is not None: sr, audio_data = audio_prompt_input duration_sec = len(audio_data) / float(sr) if sr else 0 if duration_sec > 10.0: raise gr.Error("Audio prompt must be 10 seconds or shorter.") if ( audio_data is None or audio_data.size == 0 or audio_data.max() == 0 ): gr.Warning("Audio prompt seems empty or silent, ignoring prompt.") else: with tempfile.NamedTemporaryFile( mode="wb", suffix=".wav", delete=False ) as f_audio: temp_audio_prompt_path = f_audio.name if np.issubdtype(audio_data.dtype, np.integer): max_val = np.iinfo(audio_data.dtype).max audio_data = audio_data.astype(np.float32) / max_val elif not np.issubdtype(audio_data.dtype, np.floating): gr.Warning( f"Unsupported audio prompt dtype {audio_data.dtype}, attempting conversion." ) try: audio_data = audio_data.astype(np.float32) except Exception as conv_e: raise gr.Error( f"Failed to convert audio prompt to float32: {conv_e}" ) if audio_data.ndim > 1: if audio_data.shape[0] == 2: audio_data = np.mean(audio_data, axis=0) elif audio_data.shape[1] == 2: audio_data = np.mean(audio_data, axis=1) else: gr.Warning( f"Audio prompt has unexpected shape {audio_data.shape}, taking first channel." ) audio_data = ( audio_data[0] if audio_data.shape[0] < audio_data.shape[1] else audio_data[:, 0] ) audio_data = np.ascontiguousarray(audio_data) try: sf.write( temp_audio_prompt_path, audio_data, sr, subtype="FLOAT" ) prompt_path_for_generate = temp_audio_prompt_path print( f"[AOS] Created temporary audio prompt file: {temp_audio_prompt_path} (orig sr: {sr})" ) except Exception as write_e: print(f"[AOS] Error writing temporary audio file: {write_e}") raise gr.Error(f"Failed to save audio prompt: {write_e}") # 3. Run Generation start_time = time.time() print(f"[AOS] Starting model.generate()...") with torch.inference_mode(): combined_text = ( text_input.strip() + "\n" + transcription_input.strip() if transcription_input and not transcription_input.isspace() else text_input ) output_audio_np = model.generate( combined_text, max_tokens=max_new_tokens, cfg_scale=cfg_scale, temperature=temperature, top_p=top_p, cfg_filter_top_k=cfg_filter_top_k, use_torch_compile=False, audio_prompt=prompt_path_for_generate, ) gen_time = time.time() - start_time print(f"[AOS] Generation finished in {gen_time:.2f} seconds.") print(f"[AOS] output_audio_np shape={output_audio_np.shape if output_audio_np is not None else 'None'}, " f"dtype={output_audio_np.dtype if output_audio_np is not None else 'N/A'}") # 4. Convert Codes to Audio if output_audio_np is not None: output_sr = 44100 original_len = len(output_audio_np) print(f"[AOS] Pre-speed PCM: len={original_len}, range=[{output_audio_np.min():.4f}, {output_audio_np.max():.4f}]") # --- Apply AOS Post-DSP (volume + pitch) BEFORE speed adjustment --- if volume_gain < 1.0 or pitch_flattening > 0.0: print(f"[AOS] Applying AOS post-DSP: volume_gain={volume_gain}, pitch_flattening={pitch_flattening}") output_audio_np = aos_process(output_audio_np, volume_gain, pitch_flattening, output_sr) print(f"[AOS] Post-DSP PCM: range=[{output_audio_np.min():.4f}, {output_audio_np.max():.4f}]") # --- Slow down audio --- speed_factor = max(0.1, min(speed_factor, 5.0)) target_len = int(original_len / speed_factor) if target_len != original_len and target_len > 0: x_original = np.arange(original_len) x_resampled = np.linspace(0, original_len - 1, target_len) resampled_audio_np = np.interp(x_resampled, x_original, output_audio_np) output_audio = (output_sr, resampled_audio_np.astype(np.float32)) print(f"[AOS] Resampled audio from {original_len} to {target_len} samples for {speed_factor:.2f}x speed.") else: output_audio = (output_sr, output_audio_np) print(f"[AOS] Skipping speed adjustment (factor: {speed_factor:.2f}).") print(f"[AOS] Final audio: shape={output_audio[1].shape}, sr={output_sr}") # Convert to int16 for Gradio if output_audio[1].dtype == np.float32 or output_audio[1].dtype == np.float64: audio_for_gradio = np.clip(output_audio[1], -1.0, 1.0) audio_for_gradio = (audio_for_gradio * 32767).astype(np.int16) output_audio = (output_sr, audio_for_gradio) print("[AOS] Converted to int16 for Gradio output.") else: print("\n[AOS] Generation finished, but no valid tokens were produced.") gr.Warning("Generation produced no output.") except Exception as e: print(f"[AOS] Error during inference: {e}") import traceback traceback.print_exc() raise gr.Error(f"Inference failed: {e}") finally: if temp_txt_file_path and Path(temp_txt_file_path).exists(): try: Path(temp_txt_file_path).unlink() except OSError as e: print(f"[AOS] Warning: Error deleting temp file {temp_txt_file_path}: {e}") if temp_audio_prompt_path and Path(temp_audio_prompt_path).exists(): try: Path(temp_audio_prompt_path).unlink() except OSError as e: print(f"[AOS] Warning: Error deleting temp audio file {temp_audio_prompt_path}: {e}") return output_audio # --- Create Gradio Interface --- css = """ #col-container {max-width: 90%; margin-left: auto; margin-right: auto;} """ default_text = "[S1] Dia is an open weights text to dialogue model. \n[S2] You get full control over scripts and voices. \n[S1] Wow. Amazing. (laughs) \n[S2] Try it now on Git hub or Hugging Face." example_txt_path = Path("./example.txt") if example_txt_path.exists(): try: default_text = example_txt_path.read_text(encoding="utf-8").strip() if not default_text: default_text = "Example text file was empty." except Exception as e: print(f"[AOS] Warning: Could not read example.txt: {e}") with gr.Blocks(css=css) as demo: gr.Markdown("# Nari Text-to-Speech Synthesis — AOS Enhanced") with gr.Row(equal_height=False): with gr.Column(scale=1): text_input = gr.Textbox( label="Input Text", placeholder="Enter text here...", value=default_text, lines=5, ) audio_prompt_input = gr.Audio( label="Audio Prompt (\u2264 10 s, Optional)", show_label=True, sources=["upload", "microphone"], type="numpy", ) transcription_input = gr.Textbox( label="Audio Prompt Transcription (Optional)", placeholder="Enter transcription of your audio prompt here...", lines=3, ) # ── Generation Parameters ── with gr.Accordion("Generation Parameters", open=False): max_new_tokens = gr.Slider( label="Max New Tokens (Audio Length)", minimum=860, maximum=3072, value=model.config.decoder_config.max_position_embeddings, step=50, info="Controls the maximum length of the generated audio.", ) cfg_scale = gr.Slider( label="CFG Scale (Guidance Strength)", minimum=1.0, maximum=5.0, value=3.0, step=0.1, ) temperature = gr.Slider( label="Temperature (Randomness)", minimum=1.0, maximum=2.5, value=1.8, step=0.05, ) top_p = gr.Slider( label="Top P (Nucleus Sampling)", minimum=0.70, maximum=1.0, value=0.95, step=0.01, ) cfg_filter_top_k = gr.Slider( label="CFG Filter Top K", minimum=15, maximum=100, value=45, step=1, ) speed_factor_slider = gr.Slider( label="Speed Factor", minimum=0.5, maximum=2.0, value=1.0, step=0.02, info="Adjusts the speed of the generated audio (1.0 = original speed). AOS Saint range: 0.70-0.77.", ) # ── AOS Acoustic Controls ── with gr.Accordion("AOS Acoustic Controls (Post-DSP)", open=False): gr.Markdown( "Pure tensor math applied to the output PCM. " "These are hard mathematical bounds, not LLM suggestions." ) volume_gain_slider = gr.Slider( label="Volume Gain (Amax)", minimum=0.50, maximum=1.0, value=1.0, step=0.01, info="Peak amplitude clamp. 1.0 = no change, 0.90 = -1dBFS ceiling (Saint default). " "Every sample is hard-clipped to [-gain, gain].", ) pitch_flattening_slider = gr.Slider( label="Pitch Flattening Index", minimum=0.0, maximum=0.99, value=0.0, step=0.01, info="Attenuates F0 variance via STFT spectral smoothing. " "0.0 = no change, 0.95 = 95% variance reduction (Saint default). " "Applied as: mag_new = mean + (1-index)*(mag - mean).", ) run_button = gr.Button("Generate Audio", variant="primary") with gr.Column(scale=1): audio_output = gr.Audio( label="Generated Audio", type="numpy", autoplay=False, ) # Link button click to function (11 inputs) run_button.click( fn=run_inference, inputs=[ text_input, audio_prompt_input, transcription_input, max_new_tokens, cfg_scale, temperature, top_p, cfg_filter_top_k, speed_factor_slider, volume_gain_slider, pitch_flattening_slider, ], outputs=[audio_output], api_name="generate_audio", ) # Examples example_prompt_path = "./example_prompt.mp3" examples_list = [ [ "[S1] Oh fire! Oh my goodness! What's the procedure? ...", None, "", 3072, 3.0, 1.8, 0.95, 45, 1.0, 1.0, # volume_gain 0.0, # pitch_flattening ], [ "[S1] Open weights text to dialogue model. ...", example_prompt_path if Path(example_prompt_path).exists() else None, "", 3072, 3.0, 1.8, 0.95, 45, 1.0, 0.90, # volume_gain — Saint clamp 0.95, # pitch_flattening — Saint flattening ], [ "[S1] This is a calm measured voice. [S2] Speaking slowly and deliberately.", None, "", 860, 3.0, 1.8, 0.95, 45, 0.75, 0.90, # volume_gain 0.95, # pitch_flattening ], [ "[S1] I am very excited about this! (laughs) Can you believe it?", None, "", 860, 3.0, 1.8, 0.95, 45, 0.70, 0.85, # volume_gain — tight clamp for high-arousal text 0.90, # pitch_flattening ], ] if examples_list: gr.Examples( examples=examples_list, inputs=[ text_input, audio_prompt_input, transcription_input, max_new_tokens, cfg_scale, temperature, top_p, cfg_filter_top_k, speed_factor_slider, volume_gain_slider, pitch_flattening_slider, ], outputs=[audio_output], fn=run_inference, cache_examples=False, label="Examples (Click to Run)", ) # --- Launch --- if __name__ == "__main__": print("[AOS] Launching Gradio interface...") demo.launch()