app.py

import dash
from dash import dcc, html, Input, Output, State, callback
import dash_bootstrap_components as dbc
import base64
import io
import os
from snac import SNAC
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import google.generativeai as genai
import re
import logging
import numpy as np
from pydub import AudioSegment
from docx import Document
import PyPDF2
from tqdm import tqdm
import soundfile as sf

# Initialize logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load models
print("Loading SNAC model...")
snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
snac_model = snac_model.to(device)

model_name = "canopylabs/orpheus-3b-0.1-ft"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
model.to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name)
print(f"Orpheus model loaded to {device}")

# Available voices and emotive tags
VOICES = ["tara", "leah", "jess", "leo", "dan", "mia", "zac", "zoe"]
EMOTIVE_TAGS = ["<laugh>", "<chuckle>", "<sigh>", "<cough>", "<sniffle>", "<groan>", "<yawn>", "<gasp>"]

# Initialize Dash app
app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])

app.layout = dbc.Container([
    dbc.Row([
        dbc.Col([
            html.H1("Orpheus Text-to-Speech", className="text-center mb-4"),
        ], width=12),
    ]),
    dbc.Row([
        dbc.Col([
            dbc.Input(id="host1-name", placeholder="Enter name of first host", className="mb-2"),
            dbc.Input(id="host2-name", placeholder="Enter name of second host", className="mb-2"),
            dbc.Input(id="podcast-name", placeholder="Enter podcast name", className="mb-2"),
            dbc.Input(id="podcast-topic", placeholder="Enter podcast topic", className="mb-2"),
            dbc.Textarea(id="prompt", placeholder="Enter your text here...", rows=5, className="mb-2"),
            dcc.Upload(
                id='upload-file',
                children=html.Div(['Drag and Drop or ', html.A('Select a File')]),
                style={
                    'width': '100%',
                    'height': '60px',
                    'lineHeight': '60px',
                    'borderWidth': '1px',
                    'borderStyle': 'dashed',
                    'borderRadius': '5px',
                    'textAlign': 'center',
                    'margin': '10px 0'
                },
            ),
            html.Label("Duration (minutes)", className="mt-2"),
            dcc.Slider(id="duration", min=1, max=60, value=5, step=1, marks={1: '1', 30: '30', 60: '60'}, className="mb-2"),
            html.Label("Number of Hosts", className="mt-2"),
            dbc.RadioItems(
                id="num-hosts",
                options=[{"label": i, "value": i} for i in ["1", "2"]],
                value="1",
                inline=True,
                className="mb-2"
            ),
            dbc.Button("Generate Podcast Script", id="generate-script-btn", color="primary", className="mb-2"),
            dbc.Spinner(html.Div(id="script-loading"), color="primary"),
        ], width=6),
        dbc.Col([
            dbc.Textarea(id="script-output", placeholder="Generated script will appear here...", rows=10, className="mb-2"),
            dbc.Button("Clear", id="clear-btn", color="secondary", className="mb-2 d-block"),
            html.Label("Voice 1", className="mt-3"),
            dcc.Dropdown(id="voice1", options=[{"label": v, "value": v} for v in VOICES], value="tara", className="mb-2"),
            html.Label("Voice 2", className="mt-2"),
            dcc.Dropdown(id="voice2", options=[{"label": v, "value": v} for v in VOICES], value="zac", className="mb-2"),
            dbc.Button("Generate Audio", id="generate-audio-btn", color="success", className="mb-2"),
            dbc.Spinner(html.Div(id="audio-loading"), color="primary"),
            html.Div(id="audio-output"),
            dbc.Button("Advanced Settings", id="advanced-settings-toggle", color="info", className="mb-2"),
            dbc.Collapse([
                html.Label("Temperature", className="mt-2"),
                dcc.Slider(id="temperature", min=0.1, max=1.5, value=0.6, step=0.05, marks={0.1: '0.1', 0.8: '0.8', 1.5: '1.5'}, className="mb-2"),
                html.Label("Top P", className="mt-2"),
                dcc.Slider(id="top-p", min=0.1, max=1.0, value=0.9, step=0.05, marks={0.1: '0.1', 0.5: '0.5', 1.0: '1.0'}, className="mb-2"),
                html.Label("Repetition Penalty", className="mt-2"),
                dcc.Slider(id="repetition-penalty", min=1.0, max=2.0, value=1.2, step=0.1, marks={1.0: '1.0', 1.5: '1.5', 2.0: '2.0'}, className="mb-2"),
                html.Label("Max New Tokens", className="mt-2"),
                dcc.Slider(id="max-new-tokens", min=100, max=16384, value=4096, step=100, marks={100: '100', 8192: '8192', 16384: '16384'}, className="mb-2"),
            ], id="advanced-settings", is_open=False),
        ], width=6),
    ]),
    dcc.Store(id='generated-script'),
    dcc.Store(id='generated-audio'),
])

def process_prompt(text, voice, tokenizer, device):
    prompt = f"{voice}: {text}"
    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs["input_ids"].to(device)
    attention_mask = inputs["attention_mask"].to(device)
    return input_ids, attention_mask

def parse_output(generated_ids):
    decoded = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    code_list = [int(code) for code in decoded.split() if code.isdigit()]
    return code_list

def redistribute_codes(code_list, snac_model):
    audio = snac_model.codes_to_audio(torch.tensor(code_list).unsqueeze(0).to(device))
    return audio.cpu().numpy().flatten()

def detect_silence(audio, threshold=0.01, min_silence_len=1000):
    is_silent = np.abs(audio) < threshold
    silent_regions = []
    silent_start = None
    for i, silent in enumerate(is_silent):
        if silent and silent_start is None:
            silent_start = i
        elif not silent and silent_start is not None:
            if i - silent_start >= min_silence_len:
                silent_regions.append((silent_start, i))
            silent_start = None
    if silent_start is not None and len(audio) - silent_start >= min_silence_len:
        silent_regions.append((silent_start, len(audio)))
    return silent_regions

def generate_audio(script_output, voice1, voice2, num_hosts, temperature, top_p, repetition_penalty, max_new_tokens):
    try:
        paragraphs = script_output.split('\n\n')  # Split by double newline
        audio_samples = []
        
        for i, paragraph in tqdm(enumerate(paragraphs), total=len(paragraphs), desc="Generating audio"):
            if not paragraph.strip():
                continue
            
            voice = voice1 if num_hosts == "1" or i % 2 == 0 else voice2
            
            input_ids, attention_mask = process_prompt(paragraph, voice, tokenizer, device)
            
            with torch.no_grad():
                generated_ids = model.generate(
                    input_ids,
                    attention_mask=attention_mask,
                    do_sample=True,
                    temperature=temperature,
                    top_p=top_p,
                    repetition_penalty=repetition_penalty,
                    max_new_tokens=max_new_tokens,
                    num_return_sequences=1,
                    eos_token_id=128258,
                    pad_token_id=128258,
                )
            
            code_list = parse_output(generated_ids)
            
            # Ensure the code list matches the expected input size of the SNAC model
            expected_size = 2048  # This should match the model's expected input size
            if len(code_list) < expected_size:
                code_list = code_list + [0] * (expected_size - len(code_list))
            elif len(code_list) > expected_size:
                code_list = code_list[:expected_size]
            
            # Convert to float tensor to match bias type
            codes_tensor = torch.tensor(code_list, dtype=torch.float32).unsqueeze(0).to(device)
            
            # Reshape the tensor to match the expected input shape
            codes_tensor = codes_tensor.view(1, -1, 2048)  # Adjust these dimensions as needed
            
            # Generate audio
            with torch.no_grad():
                paragraph_audio = snac_model(codes_tensor)
            
            # Handle tuple output
            if isinstance(paragraph_audio, tuple):
                paragraph_audio = paragraph_audio[0]  # Assume the first element is the audio tensor
            
            paragraph_audio = paragraph_audio.cpu().numpy().flatten()
            
            # Normalize audio to [-1, 1] range
            paragraph_audio = paragraph_audio / np.max(np.abs(paragraph_audio))
            
            audio_samples.append(paragraph_audio)
        
        final_audio = np.concatenate(audio_samples)
        
        # Convert to 16-bit PCM
        final_audio = (final_audio * 32767).astype(np.int16)
        
        # Save as WAV file in memory
        buffer = io.BytesIO()
        sf.write(buffer, final_audio, 24000, format='WAV', subtype='PCM_16')
        buffer.seek(0)
        
        return buffer
    except Exception as e:
        logger.error(f"Error generating speech: {str(e)}")
        return None
        
@callback(
    Output("script-output", "value"),
    Output("audio-output", "children"),
    Output("advanced-settings", "is_open"),
    Output("prompt", "value"),
    Output("script-loading", "children"),
    Output("audio-loading", "children"),
    Input("generate-script-btn", "n_clicks"),
    Input("generate-audio-btn", "n_clicks"),
    Input("advanced-settings-toggle", "n_clicks"),
    Input("clear-btn", "n_clicks"),
    State("host1-name", "value"),
    State("host2-name", "value"),
    State("podcast-name", "value"),
    State("podcast-topic", "value"),
    State("prompt", "value"),
    State("upload-file", "contents"),
    State("duration", "value"),
    State("num-hosts", "value"),
    State("script-output", "value"),
    State("voice1", "value"),
    State("voice2", "value"),
    State("temperature", "value"),
    State("top-p", "value"),
    State("repetition-penalty", "value"),
    State("max-new-tokens", "value"),
    State("advanced-settings", "is_open"),
    prevent_initial_call=True
)
def combined_callback(generate_script_clicks, generate_audio_clicks, advanced_settings_clicks, clear_clicks,
                      host1_name, host2_name, podcast_name, podcast_topic, prompt, uploaded_file, duration, num_hosts,
                      script_output, voice1, voice2, temperature, top_p, repetition_penalty, max_new_tokens, is_advanced_open):
    ctx = dash.callback_context
    if not ctx.triggered:
        return dash.no_update, dash.no_update, dash.no_update, dash.no_update, "", ""

    trigger_id = ctx.triggered[0]['prop_id'].split('.')[0]

    if trigger_id == "generate-script-btn":
        try:
            api_key = os.environ.get("GEMINI_API_KEY")
            if not api_key:
                raise ValueError("Gemini API key not found in environment variables")
            
            genai.configure(api_key=api_key)
            model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
            
            combined_content = prompt or ""
            
            if uploaded_file:
                content_type, content_string = uploaded_file.split(',')
                decoded = base64.b64decode(content_string)
                file_bytes = io.BytesIO(decoded)
                
                file_bytes.seek(0)
                if file_bytes.read(4) == b'%PDF':
                    file_bytes.seek(0)
                    pdf_reader = PyPDF2.PdfReader(file_bytes)
                    file_content = "\n".join([page.extract_text() for page in pdf_reader.pages])
                else:
                    file_bytes.seek(0)
                    try:
                        file_content = file_bytes.read().decode('utf-8')
                    except UnicodeDecodeError:
                        file_bytes.seek(0)
                        try:
                            doc = Document(file_bytes)
                            file_content = "\n".join([para.text for para in doc.paragraphs])
                        except:
                            raise ValueError("Unsupported file type or corrupted file")

                combined_content += "\n" + file_content if combined_content else file_content
            
            num_hosts = int(num_hosts) if num_hosts else 1
            
            prompt_template = f"""
            Create a podcast script for {num_hosts} {'person' if num_hosts == 1 else 'people'} discussing:
            {combined_content}
            
            Duration: {duration} minutes. Include natural speech, humor, and occasional off-topic thoughts.
            Use speech fillers like um, ah. Vary emotional tone.
            
            Format: {'Monologue' if num_hosts == 1 else 'Alternating dialogue'} without speaker labels.
            Separate {'paragraphs' if num_hosts == 1 else 'lines'} with blank lines.
            If the number of {num_hosts} is 1 then each paragraph will be no more than 3 sentences each
            Only provide the dialog for text to speech.
            Only use these emotion tags in angle brackets: {', '.join(EMOTIVE_TAGS)}.
            -Example: "I can't believe I stayed up all night <yawn> only to find out the meeting was canceled <groan>."
            Ensure content flows naturally and stays on topic. Match the script length to {duration} minutes.
            Do not include speaker labels like "jane:" or "john:" before dialogue.
            The intro always includes the ({host1_name} and/or {host2_name}) if it exists and should be in the same paragraph. 
            The outro always includes the ({host1_name} and/or {host2_name}) if it exists and should be in the same paragraph
            Do not include these types of transitions in the intro, outro or between paragraphs for example: "Intro Music fades in...".  Its just dialog.
            Keep each speaker's entire monologue in a single paragraph, regardless of length if the number of hosts is not 1.
            Start a new paragraph only when switching to a different speaker if the number of hosts is not 1.
            Maintain natural conversation flow and speech patterns within each monologue.
            Use context clues or subtle references to indicate who is speaking without explicit labels if the number of hosts is not 1.
            Use speaker names ({host1_name} and/or {host2_name}) sparingly, only when necessary for clarity or emphasis. Avoid starting every line with the other person's name.
            Rely more on context and speech patterns to indicate who is speaking, rather than always stating names.
            Use names primarily for transitions sparingly, definitely with agreements, or to draw attention to a specific point, not as a constant form of address.
            {'Make sure the script is a monologue for one person.' if num_hosts == 1 else f'Ensure the dialogue alternates between two distinct voices, with {host1_name} speaking on odd-numbered lines and {host2_name} on even-numbered lines.'}
            Always include intro with the speaker name and its the podcast name "{podcast_name}" in intoduce the topic of the podcast with "{podcast_topic}".
            Incorporate the podcast name and topic naturally into the intro and outro, and ensure the content stays relevant to the specified topic throughout the script.
            """
            
            response = model.generate_content(prompt_template)
            return re.sub(r'[^a-zA-Z0-9\s.,?!<>]', '', response.text), dash.no_update, dash.no_update, dash.no_update, "", ""
        except Exception as e:
            logger.error(f"Error generating podcast script: {str(e)}")
            return f"Error: {str(e)}", dash.no_update, dash.no_update, dash.no_update, "", ""

    elif trigger_id == "generate-audio-btn":
        if not script_output.strip():
            return dash.no_update, html.Div("No audio generated yet."), dash.no_update, dash.no_update, "", ""
        
        audio_buffer = generate_audio(script_output, voice1, voice2, num_hosts, temperature, top_p, repetition_penalty, max_new_tokens)
        
        if audio_buffer is not None:
            # Convert to base64 for audio playback
            audio_base64 = base64.b64encode(audio_buffer.getvalue()).decode('utf-8')
            src = f"data:audio/wav;base64,{audio_base64}"
            
            # Create a download link for the audio
            download_link = html.A("Download Audio", href=src, download="generated_audio.wav")
            
            return dash.no_update, html.Div([
                html.Audio(src=src, controls=True),
                html.Br(),
                download_link
            ]), dash.no_update, dash.no_update, "", ""
        else:
            return dash.no_update, html.Div("Error generating audio"), dash.no_update, dash.no_update, "", ""

    elif trigger_id == "advanced-settings-toggle":
        return dash.no_update, dash.no_update, not is_advanced_open, dash.no_update, "", ""

    elif trigger_id == "clear-btn":
        return "", html.Div("No audio generated yet."), dash.no_update, "", "", ""

    return dash.no_update, dash.no_update, dash.no_update, dash.no_update, "", ""

# Run the app
if __name__ == '__main__':
    print("Starting the Dash application...")
    app.run(debug=True, host='0.0.0.0', port=7860)
    print("Dash application has finished running.")