app.py

import os
import subprocess
import sys
# Fix OMP_NUM_THREADS issue before any imports
os.environ["OMP_NUM_THREADS"] = "4"
# Install dependencies programmatically to avoid conflicts
def setup_dependencies():
    try:
        # Check if already installed
        if os.path.exists('/tmp/deps_installed'):
            return
            
        print("Installing transformers dev version...")
        subprocess.check_call([
            sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-cache-dir",
            "git+https://github.com/huggingface/transformers.git"
        ])
        
        # Mark as installed
        with open('/tmp/deps_installed', 'w') as f:
            f.write('done')
            
    except Exception as e:
        print(f"Dependencies setup error: {e}")
# Run setup
setup_dependencies()

import spaces
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch
import librosa
import gradio as gr
from nemo.collections.tts.models import AudioCodecModel
import os
import sys

# Add the parent directory to sys.path to import kanitts
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from kanitts import Config

# Load configuration
config = Config.default()

# Load KaniTTS model and tokenizer
kani_model_id = config.model.model_name

tokenizer = AutoTokenizer.from_pretrained(
    kani_model_id,
    trust_remote_code=True,
    use_fast=True,
)

model = AutoModelForCausalLM.from_pretrained(
    kani_model_id,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
    trust_remote_code=True,
)
model.eval()

# Load Nemo codec
nemo_model_id = config.audio.nemo_model_name
nemo_codec = AudioCodecModel.from_pretrained(nemo_model_id).eval().cuda()

# Load Whisper for transcription
whisper_turbo_pipe = pipeline(
    "automatic-speech-recognition",
    model="openai/whisper-large-v3-turbo",
    torch_dtype=torch.float16,
    device='cuda',
)

# KaniTTS token IDs from config
tokens = config.tokens
SOH_ID = tokens.start_of_human
EOH_ID = tokens.end_of_human
SOA_ID = tokens.start_of_ai
EOA_ID = tokens.end_of_ai
SOT_ID = tokens.start_of_text
EOT_ID = tokens.end_of_text
SOS_ID = tokens.start_of_speech
EOS_ID = tokens.end_of_speech

def tokenize_audio(waveform, target_sample_rate=22050):
    """
    Tokenize audio using Nemo codec for KaniTTS.
    """
    # Ensure correct sample rate
    if waveform.shape[0] > 1:
        waveform = waveform.mean(dim=0, keepdim=True)  # Convert to mono if stereo
    
    # Resample if needed (simplified - in practice you'd use proper resampling)
    waveform = waveform.to(dtype=torch.float32)
    
    # Ensure we have the right shape: [batch, samples]
    if waveform.dim() == 1:
        waveform = waveform.unsqueeze(0)
    
    waveform = waveform.to(nemo_codec.device)
    
    # Calculate audio length in samples
    audio_len = torch.tensor([waveform.shape[-1]], dtype=torch.int64).to(waveform.device)
    
    # Encode audio to get token codes
    with torch.inference_mode():
        encoded_tokens, _ = nemo_codec.encode(audio=waveform, audio_len=audio_len)
    
    # encoded_tokens shape: [batch, num_codebooks, sequence_length]
    # For nemo-nano-codec: [1, 4, seq_len]
    codes = encoded_tokens[0]  # Remove batch dimension -> [4, seq_len]
    seq_len = codes.shape[1]
    
    # Flatten the 4 codebook levels per frame (KaniTTS uses 4 tokens per frame)
    all_codes = []
    
    for i in range(seq_len):
        # Extract one frame across all 4 codebook levels
        for level in range(4):
            token_id = codes[level, i].item()
            # Add offset for each codebook level
            offset_token = token_id + config.tokens.audio_tokens_start + (level * config.tokens.codebook_size)
            all_codes.append(offset_token)
    
    return all_codes

def redistribute_codes(code_list):
    """
    Decode audio codes back to waveform using Nemo codec.
    """
    if len(code_list) % 4 != 0:
        print(f"Warning: Code list length {len(code_list)} is not divisible by 4")
        return None

    num_frames = len(code_list) // 4
    codebook_size = config.tokens.codebook_size

    # Separate the 4 codebook levels
    level_0 = []
    level_1 = []
    level_2 = []
    level_3 = []

    for i in range(num_frames):
        # Extract each level and remove offsets
        level_0.append((code_list[4*i] - config.tokens.audio_tokens_start) % codebook_size)
        level_1.append((code_list[4*i + 1] - config.tokens.audio_tokens_start - codebook_size) % codebook_size)
        level_2.append((code_list[4*i + 2] - config.tokens.audio_tokens_start - 2*codebook_size) % codebook_size)
        level_3.append((code_list[4*i + 3] - config.tokens.audio_tokens_start - 3*codebook_size) % codebook_size)

    # Convert to tensors in format expected by Nemo: [batch, num_codebooks, sequence_length]
    codes = torch.stack([
        torch.tensor(level_0, dtype=torch.long),
        torch.tensor(level_1, dtype=torch.long),
        torch.tensor(level_2, dtype=torch.long),
        torch.tensor(level_3, dtype=torch.long)
    ]).unsqueeze(0)  # Add batch dimension

    try:
        # Move to codec device
        codes = codes.to(nemo_codec.device)

        # Calculate length
        tokens_len = torch.tensor([codes.shape[-1]], dtype=torch.int64).to(nemo_codec.device)

        # Decode
        with torch.no_grad():
            audio_hat, _ = nemo_codec.decode(tokens=codes, tokens_len=tokens_len)

        return audio_hat.cpu()

    except Exception as e:
        print(f"Error decoding audio: {e}")
        return None
@spaces.GPU(duration=30) 
def transcribe_audio(sample_audio_path, progress=gr.Progress()):
    """Transcribe uploaded audio using Whisper."""
    if not sample_audio_path:
        gr.Warning("Please upload an audio file first.")
        return ""

    try:
        progress(0, 'Loading audio...')
        audio_array, sample_rate = librosa.load(sample_audio_path, sr=config.audio.sample_rate)

        # Trim audio to max 15 seconds for transcription
        if len(audio_array) / sample_rate > 15:
            num_samples_to_keep = int(sample_rate * 15)
            audio_array = audio_array[:num_samples_to_keep]

        progress(0.5, 'Transcribing...')
        transcript = whisper_turbo_pipe(audio_array)['text'].strip()
        progress(1, 'Transcription complete!')

        return transcript
    except Exception as e:
        gr.Error(f"Transcription failed: {str(e)}")
        return ""

@spaces.GPU(duration=60)
def infer(sample_audio_path, ref_transcript, target_text, temperature, top_p, repetition_penalty, progress=gr.Progress()):
    if not target_text or not target_text.strip():
        gr.Warning("Please input text to generate audio.")
        return None

    if len(target_text) > 500:
        gr.Warning("Text is too long. Please keep it under 500 characters.")
        target_text = target_text[:500]

    target_text = target_text.strip()

    if sample_audio_path and (not ref_transcript or not ref_transcript.strip()):
        gr.Warning("Please provide a transcript for the reference audio or use the transcribe button.")
        return None
    
    with torch.no_grad():
        if sample_audio_path and ref_transcript:
            progress(0, 'Loading and trimming audio...')
            audio_array, sample_rate = librosa.load(sample_audio_path, sr=config.audio.sample_rate)

            # Trim audio to max 15 seconds
            if len(audio_array) / sample_rate > 15:
                gr.Warning("Trimming audio to first 15secs.")
                num_samples_to_keep = int(sample_rate * 15)
                audio_array = audio_array[:num_samples_to_keep]

            prompt_wav = torch.from_numpy(audio_array).unsqueeze(0)
            prompt_wav = prompt_wav.to(dtype=torch.float32)

            progress(0.4, 'Encoding reference audio...')

            # Encode the prompt wav
            voice_tokens = tokenize_audio(prompt_wav)

            # Use the provided transcript instead of auto-transcribing
            prompt_text = ref_transcript.strip()

            progress(0.6, "Generating audio...")

            # Tokenize target text
            target_text_ids = tokenizer.encode(target_text, add_special_tokens=False)

            # Create complete sentence (reference + target)
            complete_text = prompt_text + " " + target_text
            complete_text_ids = tokenizer.encode(complete_text, add_special_tokens=False)

            # Create prompt: Human says complete sentence, AI provides partial audio + continues
            prompt_ids = (
                [SOH_ID]
                + complete_text_ids  # Full sentence as human input
                + [EOT_ID]
                + [EOH_ID]
                + [SOA_ID]
                + [SOS_ID]
                + voice_tokens        # Audio only for reference part
                # Model should continue generating audio for the target part
            )
        else:
            # No reference audio case
            prompt_ids = []
            progress(0.6, "Generating audio...")

            # Tokenize target text
            target_text_ids = tokenizer.encode(target_text, add_special_tokens=False)

            # Simple generation without reference
            prompt_ids.extend([SOH_ID])
            prompt_ids.extend(target_text_ids)
            prompt_ids.extend([EOT_ID])
            prompt_ids.extend([EOH_ID])
            prompt_ids.extend([SOA_ID])
            prompt_ids.extend([SOS_ID])

        print(f"Prompt length: {len(prompt_ids)} tokens")

        input_ids = torch.tensor([prompt_ids], dtype=torch.int64).cuda()

        # Generate the speech autoregressively
        outputs = model.generate(
            input_ids,
            max_new_tokens=config.model.max_new_tokens,
            eos_token_id=EOS_ID,
            do_sample=True,
            top_p=top_p,
            temperature=temperature,
            repetition_penalty=repetition_penalty,
            pad_token_id=config.tokens.pad_token,
            use_cache=True,
        )
        generated_ids = outputs[0].tolist()
        print(f"Generated {len(generated_ids)} total tokens")

        progress(0.8, "Decoding generated audio...")

        # Since we end our prompt with SOS_ID, the generated tokens should be audio tokens directly
        # We need to find where our input prompt ends and the generated tokens begin
        input_length = len(prompt_ids)
        speech_tokens = generated_ids[input_length:]

        print(f"Input prompt length: {input_length}, generated tokens: {len(speech_tokens)}")

        # Remove end of speech token if present
        if EOS_ID in speech_tokens:
            speech_tokens = speech_tokens[:speech_tokens.index(EOS_ID)]

        if not speech_tokens:
            gr.Error("Audio generation failed: No speech tokens were generated.")
            return None

        # Filter out non-audio tokens
        audio_tokens = [token for token in speech_tokens if token >= config.tokens.audio_tokens_start]
        
        if not audio_tokens:
            gr.Error("Audio generation failed: No valid audio tokens found.")
            return None

        print(f"Decoding {len(audio_tokens)} audio tokens")
        gen_wav_tensor = redistribute_codes(audio_tokens)

        if gen_wav_tensor is None:
            gr.Error("Audio decoding failed.")
            return None

        gen_wav = gen_wav_tensor.squeeze()

        progress(1, 'Synthesized!')
        return (config.audio.sample_rate, gen_wav.numpy())


theme = gr.themes.Glass(
    primary_hue="cyan",
)


with gr.Blocks(theme=theme, title="KaniTTS Zero-Shot Voice Cloning") as app_tts:
    gr.Markdown("# KaniTTS Zero-Shot Voice Cloning")
    gr.Markdown("Upload reference audio, provide its transcript, and enter text to generate speech in the reference voice.")

    ref_audio_input = gr.Audio(label="Reference Audio", type="filepath")

    with gr.Row():
        ref_transcript_input = gr.Textbox(
            label="Reference Audio Transcript",
            lines=3,
            placeholder="Enter what the reference audio says, or use the transcribe button...",
            info="This should match exactly what is said in the reference audio"
        )
        transcribe_btn = gr.Button("Transcribe", variant="secondary", size="sm")

    gen_text_input = gr.Textbox(
        label="Text to Generate",
        lines=10,
        placeholder="Enter the text you want to generate in the reference voice..."
    )

    with gr.Row():
        temperature_slider = gr.Slider(
            minimum=0.0, maximum=2.0, value=1.4, step=0.05,
            label="Temperature",
            info="Higher values make output more random"
        )
        top_p_slider = gr.Slider(
            minimum=0.0, maximum=1.0, value=0.9, step=0.05, 
            label="Top-p",
            info="Nucleus sampling threshold"
        )
        repetition_penalty_slider = gr.Slider(
            minimum=1.0, maximum=1.5, value=1.1, step=0.05, 
            label="Repetition Penalty",
            info="Penalty for repeating tokens"
        )

    generate_btn = gr.Button("Generate Speech", variant="primary")

    audio_output = gr.Audio(label="Generated Audio")

    # Connect transcribe button
    transcribe_btn.click(
        transcribe_audio,
        inputs=[ref_audio_input],
        outputs=[ref_transcript_input],
    )

    # Connect generate button
    generate_btn.click(
        infer,
        inputs=[
            ref_audio_input,
            ref_transcript_input,
            gen_text_input,
            temperature_slider,
            top_p_slider,
            repetition_penalty_slider,
        ],
        outputs=[audio_output],
    )

with gr.Blocks() as app_info:
    gr.Markdown("""
# About KaniTTS

KaniTTS is a conversational text-to-speech model that can perform zero-shot voice cloning.

## How to use:
1. Upload a reference audio file (WAV or MP3, max 15 seconds)
2. Either enter the transcript manually or click "Transcribe" to auto-transcribe
3. Edit the transcript if needed to ensure accuracy
4. Enter the text you want to generate in that voice
5. Adjust generation parameters if needed
6. Click "Generate Speech"

The model will use your provided transcript to understand the reference voice and generate the target text in the same voice.

## Tips:
- Use clear, high-quality reference audio
- Keep reference audio under 15 seconds
- The model works best with conversational speech
- Try different temperature settings for varied results

## Credits:
- KaniTTS model by the KaniTTS team
- Nemo codec by NVIDIA
- Interface adapted from Orpheus TTS demo
""")

with gr.Blocks() as app:
    gr.Markdown(
        """
# KaniTTS Zero-Shot Voice Cloning

This is a web interface for KaniTTS zero-shot voice cloning. Upload reference audio and generate speech in any voice!

"""
    )
    gr.TabbedInterface([app_tts, app_info], ["Voice Cloning", "About"])

if __name__ == "__main__":
    app.launch()