Spaces:

Open-LiveTranslate
/

Demo-InfiniSST

Running on CPU Upgrade

Fix 'Backend: Unreachable' after the rebuild

63d7279 verified 6 days ago

11 kB

import re
import time

import gradio as gr
import numpy as np

import torch
import torchaudio.functional as F

def prepare_speech(new_chunk):
    sr, y = new_chunk
    # Convert to mono if stereo
    if y.ndim > 1:
        y = y.mean(axis=1)
        
    y = y.astype(np.float32)
    y /= 32768.0

    resampled_y = F.resample(torch.from_numpy(y), sr, 16000)

    return resampled_y.numpy()

def wav_array_to_base64(wav_array, sample_rate):
    """Convert a numpy audio array to base64 encoded WAV."""
    import base64
    import io
    import soundfile as sf
    
    buffer = io.BytesIO()
    sf.write(buffer, wav_array, sample_rate, format='WAV')
    buffer.seek(0)
    return base64.b64encode(buffer.read()).decode('utf-8')

LANGUAGE_PAIRS = {
    "English → Chinese": ("English", "Chinese"),
    "English → Spanish": ("English", "Spanish"),
    "English → German": ("English", "German"),
    "English → French": ("English", "French"),
    "English → Japanese": ("English", "Japanese"),
}

def prepare_inputs(messages, audio_base64, language_pair):
    src_lang, tgt_lang = LANGUAGE_PAIRS.get(language_pair, ("English", "Chinese"))
    if not messages:  # Check for None or empty list
        messages = [
            {
                "role": "system",
                "content": f"You are a professional simultaneous interpreter. You will be given chunks of {src_lang} audio and you need to translate the audio into {tgt_lang} text."
            },
        ]
    messages.append(
        {
            "role": "user",
            "content": [{"type": "audio_url", "audio_url": {"url": f"data:audio/wav;base64,{audio_base64}"}}]
        }
    )
    return messages

def check_backend_status():
    """Check if the vLLM backend is reachable."""
    from openai import OpenAI
    try:
        client = OpenAI(
            base_url="https://jaida-avian-irmgard.ngrok-free.dev/v1",
            api_key="EMPTY",
        )
        client.models.list()
        return "🟢 Backend: Connected"
    except Exception:
        return "🔴 Backend: Unreachable — try refreshing later"


def translate(messages, new_chunk, chunk_buffer, chunk_size_seconds, last_chunk_time, language_pair):
    """
    Translate audio chunks with buffering. No typewriter logic here — that runs on a separate timer.
    Returns:
        messages, chunk_buffer, current_time
    """
    from openai import OpenAI

    current_time = time.time()

    if new_chunk is None:
        return messages, chunk_buffer, last_chunk_time

    # Initialize messages if None
    if messages is None:
        messages = []

    # Initialize chunk_buffer if None
    if chunk_buffer is None:
        chunk_buffer = []

    # Check if there was a significant gap (> 2 seconds) - indicates pause/resume
    # Clear partial buffer to avoid concatenating audio from different time periods
    if last_chunk_time is not None and (current_time - last_chunk_time) > 2.0:
        if chunk_buffer:
            print(f"⚠️ Detected pause (gap: {current_time - last_chunk_time:.1f}s). Clearing {len(chunk_buffer)} partial chunks.")
        chunk_buffer = []

    # Prepare and buffer the new chunk
    y = prepare_speech(new_chunk)
    chunk_buffer.append(y)

    # Calculate how many 0.96s chunks we need to reach target size
    chunks_needed = int(chunk_size_seconds / 0.96)

    # If we haven't accumulated enough chunks yet, return without processing
    if len(chunk_buffer) < chunks_needed:
        return messages, chunk_buffer, current_time

    # We have enough chunks - concatenate and process
    concatenated_audio = np.concatenate(chunk_buffer[:chunks_needed])
    chunk_buffer = chunk_buffer[chunks_needed:]  # Keep any extra chunks for next iteration

    # Convert to base64
    audio_base64 = wav_array_to_base64(concatenated_audio, 16000)

    # Prepare messages
    messages = prepare_inputs(messages, audio_base64, language_pair)

    # Calculate context window size based on chunk size
    # Larger chunks = longer audio = can keep fewer messages in context
    # Base: 30 messages for 1.92s chunks, scale proportionally
    context_window = max(10, int(30 * (1.92 / chunk_size_seconds)))

    # Call OpenAI API
    # model owaski/Open-LiveTranslate-v0-En-Zh served locally with vllm
    client = OpenAI(
        base_url="https://jaida-avian-irmgard.ngrok-free.dev/v1",
        api_key="EMPTY",
    )

    model_path = "/data/user_data/siqiouya/ckpts/infinisst-omni/gigaspeech-zh-s_origin-bsz4/v1-20260122-055820-hf"

    try:
        completion = client.chat.completions.create(
            model=model_path,
            messages=[messages[0]] + messages[-context_window:],
            top_p=0.95,
            temperature=0.6,
            extra_body={"top_k": 20}
        )
        print(f"completion: {completion}")
        translation = completion.choices[0].message.content
    except Exception as e:
        print(f"Backend error: {e}")
        translation = ""

    messages.append(
        {
            "role": "assistant",
            "content": translation
        }
    )

    return messages, chunk_buffer, current_time


def typewriter_tick(messages, displayed_text, language_pair):
    """Advance the typewriter by one unit. Runs on a fast gr.Timer, independent of the audio stream."""
    if displayed_text is None:
        displayed_text = ""
    full_translation = ''.join([m["content"] for m in messages if m["role"] == "assistant"]) if messages else ""
    displayed_text = _typewriter_step(displayed_text, full_translation, language_pair)
    return _trim_display(displayed_text), displayed_text


CJK_LANGUAGES = {"Chinese", "Japanese"}

def _typewriter_step(displayed_text, full_text, language_pair="English → Chinese"):
    """Advance displayed_text toward full_text by one unit at a time.
    For CJK target languages (Chinese, Japanese): one character per step.
    For other languages (English, Spanish, etc.): one word per step.
    """
    if displayed_text is None:
        displayed_text = ""
    if not full_text:
        return displayed_text
    if len(displayed_text) >= len(full_text):
        return full_text

    _, tgt_lang = LANGUAGE_PAIRS.get(language_pair, ("English", "Chinese"))
    remaining = full_text[len(displayed_text):]

    if tgt_lang in CJK_LANGUAGES:
        # One character at a time
        return full_text[:len(displayed_text) + 1]
    else:
        # One word at a time: find the end of the next word (skip leading spaces, then find next space)
        i = 0
        # Include leading whitespace with the word
        while i < len(remaining) and remaining[i] in (' ', '\n', '\t'):
            i += 1
        # Include the word itself
        while i < len(remaining) and remaining[i] not in (' ', '\n', '\t'):
            i += 1
        return full_text[:len(displayed_text) + max(i, 1)]


def _trim_display(text, max_lines=5):
    """Keep only the last N non-empty lines for display."""
    if not text:
        return ""
    translation_lines = text.split('\n')
    non_empty_lines = [line for line in translation_lines if line.strip()]

    if len(non_empty_lines) > max_lines:
        count = 0
        for i in range(len(translation_lines) - 1, -1, -1):
            if translation_lines[i].strip():
                count += 1
                if count == max_lines:
                    return '\n'.join(translation_lines[i:])
    return text


with gr.Blocks(css="""
    .large-font textarea {
        font-size: 20px !important;
        font-weight: 500;
        overflow-y: auto !important;
    }
    .large-font label {
        font-size: 20px !important;
        font-weight: bold;
    }
    .markdown-text p {
        font-size: 20px !important;
    }
    .note-text p {
        font-size: 16px !important;
        color: #666;
    }
""") as demo:
    gr.Markdown("# Streaming Speech-to-Text Translation with Open-LiveTranslate")
    gr.Markdown("""**Instructions:**
- Select a language pair and chunk size
- Click the microphone to start recording
- Refresh the page to reset the translation history
- If you experience no response or slow response, please refresh — the serving engine (vLLM) may be experiencing a cold start""", elem_classes=["markdown-text"])
    
    # State components
    messages_state = gr.State(value=[])
    chunk_buffer_state = gr.State(value=[])
    last_chunk_time_state = gr.State(value=None)
    displayed_text_state = gr.State(value="")

    with gr.Row():
        with gr.Column():
            with gr.Row():
                language_pair_selector = gr.Dropdown(
                    choices=list(LANGUAGE_PAIRS.keys()),
                    value="English → Chinese",
                    label="Language Pair",
                    info="Select source → target language."
                )
                chunk_size_selector = gr.Dropdown(
                    choices=[0.96, 1.92, 2.88, 3.84, 4.80, 5.76, 6.72, 7.68, 8.64, 9.60],
                    value=1.92,
                    label="Chunk Size (seconds)",
                    info="Larger chunks = more context but slower response. Must be multiple of 0.96s."
                )
            audio_input = gr.Audio(sources=["microphone"], streaming=True, label="Audio Input")

    with gr.Row():
        with gr.Column():
            status_indicator = gr.Textbox(
                label="Backend Status",
                value="Checking...",
                interactive=False,
                max_lines=1,
            )
            translation_output = gr.Textbox(
                label="Translation",
                lines=3,
                max_lines=5,
                interactive=False,
                elem_classes=["large-font"],
                autoscroll=True,
                show_copy_button=True
            )

    # Check backend status on load
    demo.load(check_backend_status, outputs=[status_indicator])

    # Audio stream — handles buffering and API calls at 0.96s intervals
    audio_input.stream(
        translate,
        inputs=[messages_state, audio_input, chunk_buffer_state, chunk_size_selector, last_chunk_time_state, language_pair_selector],
        outputs=[messages_state, chunk_buffer_state, last_chunk_time_state],
        show_progress=False,
        stream_every=0.96
    )

    # Typewriter timer — ticks every 0.02s to reveal text one unit at a time
    typewriter_timer = gr.Timer(0.02)
    typewriter_timer.tick(
        typewriter_tick,
        inputs=[messages_state, displayed_text_state, language_pair_selector],
        outputs=[translation_output, displayed_text_state],
    )

    gr.Markdown("**Note:** Currently supports English → Chinese translation only. This model is trained with the method described in the ACL 2025 paper [InfiniSST: Simultaneous Translation of Unbounded Speech with Large Language Model](https://aclanthology.org/2025.findings-acl.157/) by Siqi Ouyang, Xi Xu, and Lei Li.", elem_classes=["note-text"])

demo.launch(share=True)

	import re
	import time

	import gradio as gr
	import numpy as np

	import torch
	import torchaudio.functional as F

	def prepare_speech(new_chunk):
	sr, y = new_chunk
	# Convert to mono if stereo
	if y.ndim > 1:
	y = y.mean(axis=1)

	y = y.astype(np.float32)
	y /= 32768.0

	resampled_y = F.resample(torch.from_numpy(y), sr, 16000)

	return resampled_y.numpy()

	def wav_array_to_base64(wav_array, sample_rate):
	"""Convert a numpy audio array to base64 encoded WAV."""
	import base64
	import io
	import soundfile as sf

	buffer = io.BytesIO()
	sf.write(buffer, wav_array, sample_rate, format='WAV')
	buffer.seek(0)
	return base64.b64encode(buffer.read()).decode('utf-8')

	LANGUAGE_PAIRS = {
	"English → Chinese": ("English", "Chinese"),
	"English → Spanish": ("English", "Spanish"),
	"English → German": ("English", "German"),
	"English → French": ("English", "French"),
	"English → Japanese": ("English", "Japanese"),
	}

	def prepare_inputs(messages, audio_base64, language_pair):
	src_lang, tgt_lang = LANGUAGE_PAIRS.get(language_pair, ("English", "Chinese"))
	if not messages: # Check for None or empty list
	messages = [
	{
	"role": "system",
	"content": f"You are a professional simultaneous interpreter. You will be given chunks of {src_lang} audio and you need to translate the audio into {tgt_lang} text."
	},
	]
	messages.append(
	{
	"role": "user",
	"content": [{"type": "audio_url", "audio_url": {"url": f"data:audio/wav;base64,{audio_base64}"}}]
	}
	)
	return messages

	def check_backend_status():
	"""Check if the vLLM backend is reachable."""
	from openai import OpenAI
	try:
	client = OpenAI(
	base_url="https://jaida-avian-irmgard.ngrok-free.dev/v1",
	api_key="EMPTY",
	)
	client.models.list()
	return "🟢 Backend: Connected"
	except Exception:
	return "🔴 Backend: Unreachable — try refreshing later"


	def translate(messages, new_chunk, chunk_buffer, chunk_size_seconds, last_chunk_time, language_pair):
	"""
	Translate audio chunks with buffering. No typewriter logic here — that runs on a separate timer.

	Returns:
	messages, chunk_buffer, current_time
	"""
	from openai import OpenAI

	current_time = time.time()

	if new_chunk is None:
	return messages, chunk_buffer, last_chunk_time

	# Initialize messages if None
	if messages is None:
	messages = []

	# Initialize chunk_buffer if None
	if chunk_buffer is None:
	chunk_buffer = []

	# Check if there was a significant gap (> 2 seconds) - indicates pause/resume
	# Clear partial buffer to avoid concatenating audio from different time periods
	if last_chunk_time is not None and (current_time - last_chunk_time) > 2.0:
	if chunk_buffer:
	print(f"⚠️ Detected pause (gap: {current_time - last_chunk_time:.1f}s). Clearing {len(chunk_buffer)} partial chunks.")
	chunk_buffer = []

	# Prepare and buffer the new chunk
	y = prepare_speech(new_chunk)
	chunk_buffer.append(y)

	# Calculate how many 0.96s chunks we need to reach target size
	chunks_needed = int(chunk_size_seconds / 0.96)

	# If we haven't accumulated enough chunks yet, return without processing
	if len(chunk_buffer) < chunks_needed:
	return messages, chunk_buffer, current_time

	# We have enough chunks - concatenate and process
	concatenated_audio = np.concatenate(chunk_buffer[:chunks_needed])
	chunk_buffer = chunk_buffer[chunks_needed:] # Keep any extra chunks for next iteration

	# Convert to base64
	audio_base64 = wav_array_to_base64(concatenated_audio, 16000)

	# Prepare messages
	messages = prepare_inputs(messages, audio_base64, language_pair)

	# Calculate context window size based on chunk size
	# Larger chunks = longer audio = can keep fewer messages in context
	# Base: 30 messages for 1.92s chunks, scale proportionally
	context_window = max(10, int(30 * (1.92 / chunk_size_seconds)))

	# Call OpenAI API
	# model owaski/Open-LiveTranslate-v0-En-Zh served locally with vllm
	client = OpenAI(
	base_url="https://jaida-avian-irmgard.ngrok-free.dev/v1",
	api_key="EMPTY",
	)

	model_path = "/data/user_data/siqiouya/ckpts/infinisst-omni/gigaspeech-zh-s_origin-bsz4/v1-20260122-055820-hf"

	try:
	completion = client.chat.completions.create(
	model=model_path,
	messages=[messages[0]] + messages[-context_window:],
	top_p=0.95,
	temperature=0.6,
	extra_body={"top_k": 20}
	)
	print(f"completion: {completion}")
	translation = completion.choices[0].message.content
	except Exception as e:
	print(f"Backend error: {e}")
	translation = ""

	messages.append(
	{
	"role": "assistant",
	"content": translation
	}
	)

	return messages, chunk_buffer, current_time


	def typewriter_tick(messages, displayed_text, language_pair):
	"""Advance the typewriter by one unit. Runs on a fast gr.Timer, independent of the audio stream."""
	if displayed_text is None:
	displayed_text = ""
	full_translation = ''.join([m["content"] for m in messages if m["role"] == "assistant"]) if messages else ""
	displayed_text = _typewriter_step(displayed_text, full_translation, language_pair)
	return _trim_display(displayed_text), displayed_text


	CJK_LANGUAGES = {"Chinese", "Japanese"}

	def _typewriter_step(displayed_text, full_text, language_pair="English → Chinese"):
	"""Advance displayed_text toward full_text by one unit at a time.

	For CJK target languages (Chinese, Japanese): one character per step.
	For other languages (English, Spanish, etc.): one word per step.
	"""
	if displayed_text is None:
	displayed_text = ""
	if not full_text:
	return displayed_text
	if len(displayed_text) >= len(full_text):
	return full_text

	_, tgt_lang = LANGUAGE_PAIRS.get(language_pair, ("English", "Chinese"))
	remaining = full_text[len(displayed_text):]

	if tgt_lang in CJK_LANGUAGES:
	# One character at a time
	return full_text[:len(displayed_text) + 1]
	else:
	# One word at a time: find the end of the next word (skip leading spaces, then find next space)
	i = 0
	# Include leading whitespace with the word
	while i < len(remaining) and remaining[i] in (' ', '\n', '\t'):
	i += 1
	# Include the word itself
	while i < len(remaining) and remaining[i] not in (' ', '\n', '\t'):
	i += 1
	return full_text[:len(displayed_text) + max(i, 1)]


	def _trim_display(text, max_lines=5):
	"""Keep only the last N non-empty lines for display."""
	if not text:
	return ""
	translation_lines = text.split('\n')
	non_empty_lines = [line for line in translation_lines if line.strip()]

	if len(non_empty_lines) > max_lines:
	count = 0
	for i in range(len(translation_lines) - 1, -1, -1):
	if translation_lines[i].strip():
	count += 1
	if count == max_lines:
	return '\n'.join(translation_lines[i:])
	return text


	with gr.Blocks(css="""
	.large-font textarea {
	font-size: 20px !important;
	font-weight: 500;
	overflow-y: auto !important;
	}
	.large-font label {
	font-size: 20px !important;
	font-weight: bold;
	}
	.markdown-text p {
	font-size: 20px !important;
	}
	.note-text p {
	font-size: 16px !important;
	color: #666;
	}
	""") as demo:
	gr.Markdown("# Streaming Speech-to-Text Translation with Open-LiveTranslate")
	gr.Markdown("""Instructions:
	- Select a language pair and chunk size
	- Click the microphone to start recording
	- Refresh the page to reset the translation history
	- If you experience no response or slow response, please refresh — the serving engine (vLLM) may be experiencing a cold start""", elem_classes=["markdown-text"])

	# State components
	messages_state = gr.State(value=[])
	chunk_buffer_state = gr.State(value=[])
	last_chunk_time_state = gr.State(value=None)
	displayed_text_state = gr.State(value="")

	with gr.Row():
	with gr.Column():
	with gr.Row():
	language_pair_selector = gr.Dropdown(
	choices=list(LANGUAGE_PAIRS.keys()),
	value="English → Chinese",
	label="Language Pair",
	info="Select source → target language."
	)
	chunk_size_selector = gr.Dropdown(
	choices=[0.96, 1.92, 2.88, 3.84, 4.80, 5.76, 6.72, 7.68, 8.64, 9.60],
	value=1.92,
	label="Chunk Size (seconds)",
	info="Larger chunks = more context but slower response. Must be multiple of 0.96s."
	)
	audio_input = gr.Audio(sources=["microphone"], streaming=True, label="Audio Input")

	with gr.Row():
	with gr.Column():
	status_indicator = gr.Textbox(
	label="Backend Status",
	value="Checking...",
	interactive=False,
	max_lines=1,
	)
	translation_output = gr.Textbox(
	label="Translation",
	lines=3,
	max_lines=5,
	interactive=False,
	elem_classes=["large-font"],
	autoscroll=True,
	show_copy_button=True
	)

	# Check backend status on load
	demo.load(check_backend_status, outputs=[status_indicator])

	# Audio stream — handles buffering and API calls at 0.96s intervals
	audio_input.stream(
	translate,
	inputs=[messages_state, audio_input, chunk_buffer_state, chunk_size_selector, last_chunk_time_state, language_pair_selector],
	outputs=[messages_state, chunk_buffer_state, last_chunk_time_state],
	show_progress=False,
	stream_every=0.96
	)

	# Typewriter timer — ticks every 0.02s to reveal text one unit at a time
	typewriter_timer = gr.Timer(0.02)
	typewriter_timer.tick(
	typewriter_tick,
	inputs=[messages_state, displayed_text_state, language_pair_selector],
	outputs=[translation_output, displayed_text_state],
	)

	gr.Markdown("Note: Currently supports English → Chinese translation only. This model is trained with the method described in the ACL 2025 paper [InfiniSST: Simultaneous Translation of Unbounded Speech with Large Language Model](https://aclanthology.org/2025.findings-acl.157/) by Siqi Ouyang, Xi Xu, and Lei Li.", elem_classes=["note-text"])

	demo.launch(share=True)