Spaces:

Interhuman
/

inter2-audio-playground

Running on Zero

Inter-2 Audio playground

f821a4a about 10 hours ago

9.66 kB

"""Inter-2 Audio playground: a Gradio UI and API around the Inter-2 Audio model.
Task presets (tasks.py) fill in the prompts for emotion, or for engagement, social
signals, emotion and vocal delivery together; each can optionally be constrained to
its JSON schema. Prompts and decoding are editable.
"""

from __future__ import annotations

import html
import json
import logging
import os
import sys
import time
from collections.abc import Callable
from pathlib import Path
from typing import Any, TypeVar

import gradio as gr
import spaces
import torch
from huggingface_hub import snapshot_download
from inference import (
    DEFAULT_MAX_NEW_TOKENS,
    DEFAULT_REPETITION_PENALTY,
    ENCODER_WINDOW_S,
    TARGET_SR,
    AudioInputError,
    build_messages,
    format_summary,
    load_audio,
    parse_json,
)
from tasks import DEFAULT_TASK, TASKS
from theme import CSS, FOOTER, HERO, THEME

MODEL_ID = os.environ["MODEL_ID"]  # a Space variable: the model repo (or a local copy of it)
# cuda on a GPU or ZeroGPU Space (ZeroGPU has no GPU at import time), else cpu.
_HAS_GPU = torch.cuda.is_available() or bool(os.getenv("SPACES_ZERO_GPU"))
MODEL_DEVICE = os.getenv("MODEL_DEVICE") or ("cuda" if _HAS_GPU else "cpu")
SPACES_GPU_DURATION = int(os.getenv("SPACES_GPU_DURATION", "90"))  # rationale-heavy tasks write more
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


def _load() -> tuple[Any, Any]:
    """(model, processor), via the loader that ships with the model repo."""
    model_dir = MODEL_ID if Path(MODEL_ID).is_dir() else snapshot_download(MODEL_ID)
    sys.path.insert(0, model_dir)
    from inter2_loader import load

    loaded = load(model_dir, MODEL_DEVICE)
    logger.info("Model loaded")
    return loaded


def _compile_grammars() -> dict[str, Any]:
    """Each task's answer schema for xgrammar, in compact JSON layout."""
    try:
        import xgrammar as xgr
    except ImportError:
        logger.warning("xgrammar is not installed; the schema toggle will be ignored.")
        return {}
    vocab_size = model.get_output_embeddings().weight.shape[0]
    compiler = xgr.GrammarCompiler(xgr.TokenizerInfo.from_huggingface(processor.tokenizer, vocab_size=vocab_size))
    return {
        name: compiler.compile_json_schema(
            json.dumps(task.schema), any_whitespace=False, indent=None, separators=(",", ":")
        )
        for name, task in TASKS.items()
    }


model, processor = _load()
grammars = _compile_grammars()

F = TypeVar("F", bound=Callable[..., Any])


def _spaces_gpu(func: F) -> F:
    setting = os.getenv("ENABLE_SPACES_GPU", "auto").lower()
    enabled = setting in {"1", "true", "yes"} or (setting == "auto" and bool(os.getenv("SPACE_ID")))
    return spaces.GPU(duration=SPACES_GPU_DURATION)(func) if enabled else func


@_spaces_gpu
def _generate(inputs: dict[str, Any], gen_kwargs: dict[str, Any], schema_task: str | None) -> tuple[str, int]:
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    if "input_features" in inputs:
        inputs["input_features"] = inputs["input_features"].to(model.dtype)
    if schema_task is not None:
        from xgrammar.contrib.hf import LogitsProcessor  # fresh per call: it keeps per-sequence state

        gen_kwargs = {**gen_kwargs, "logits_processor": [LogitsProcessor(grammars[schema_task])]}
    prompt_len = inputs["input_ids"].shape[-1]
    with torch.inference_mode():
        output_ids = model.generate(**inputs, **gen_kwargs)
    new_ids = output_ids[:, prompt_len:]
    text = processor.batch_decode(new_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    return text.strip(), int(new_ids.shape[-1])


def predict(
    audio_path: str | None,
    system_prompt: str,
    user_prompt: str,
    audio_first: bool,
    constrain_json: bool,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    repetition_penalty: float,
    task: str = DEFAULT_TASK,
) -> tuple[str, str, str]:
    """Run one clip. Returns (summary HTML, parsed JSON, raw text). `task` picks the schema used
    when constrain_json is on; the prompts are whatever was passed in."""
    if not audio_path:
        return _message("Please upload or record an audio clip."), "", ""
    if not user_prompt.strip():
        return _message("The user prompt is empty."), "", ""
    try:
        wav, duration = load_audio(audio_path)
        chat = processor.apply_chat_template(
            build_messages(user_prompt, system_prompt.strip(), audio_first), tokenize=False, add_generation_prompt=True
        )
        inputs = dict(processor(text=[chat], audio=[wav], sampling_rate=TARGET_SR, return_tensors="pt", padding=True))
        gen_kwargs: dict[str, Any] = {
            "max_new_tokens": int(max_new_tokens),
            "repetition_penalty": float(repetition_penalty),
            "do_sample": temperature > 0,
        }
        if temperature > 0:
            gen_kwargs.update(temperature=float(temperature), top_p=float(top_p))
        start = time.time()
        schema_task = task if constrain_json and task in grammars else None
        raw, n_tokens = _generate(inputs, gen_kwargs, schema_task)
        logger.info("%s: %d tokens in %.1fs (constrained=%s)", task, n_tokens, time.time() - start, bool(schema_task))
    except Exception as exc:
        logger.exception("predict failed")
        if isinstance(exc, AudioInputError):  # our own input checks; other errors may carry internals
            return _message(str(exc)), "", ""
        return _message("Something went wrong while analysing this clip. Please try again or use another file."), "", ""
    payload = parse_json(raw)
    pretty = json.dumps(payload, indent=2, ensure_ascii=False) if payload is not None else ""
    summary = format_summary(payload, duration if duration > ENCODER_WINDOW_S else None, n_tokens >= max_new_tokens)
    return summary, pretty, raw


def _message(text: str) -> str:
    return f'<div class="ih-card"><p class="ih-note">{html.escape(text)}</p></div>'


def _preset(task: str) -> tuple[str, str, bool]:
    preset = TASKS[task]
    return preset.system_prompt, preset.user_prompt, preset.constrain


with gr.Blocks(title="Inter-2-Audio Playground") as demo:
    gr.HTML(HERO)
    with gr.Row(equal_height=False):
        with gr.Column(scale=5):
            audio_in = gr.Audio(label="Audio clip", sources=["upload", "microphone"], type="filepath")
            task_in = gr.Dropdown(list(TASKS), value=DEFAULT_TASK, label="Task")
            user_prompt_in = gr.Textbox(label="Prompt", value=TASKS[DEFAULT_TASK].user_prompt, lines=8, max_lines=12)
            system_prompt_in = gr.Textbox(
                label="System prompt (empty = the model's default)",
                value=TASKS[DEFAULT_TASK].system_prompt,
                lines=3,
                max_lines=8,
            )
            constrain_in = gr.Checkbox(
                value=TASKS[DEFAULT_TASK].constrain, label="Constrain the answer to the task's JSON schema"
            )
            with gr.Accordion("Advanced", open=False):
                audio_first_in = gr.Checkbox(value=True, label="Audio before text")
                max_tokens_in = gr.Slider(16, 1024, value=DEFAULT_MAX_NEW_TOKENS, step=16, label="Max new tokens")
                temperature_in = gr.Slider(0.0, 1.5, value=0.0, step=0.05, label="Temperature (0 = greedy)")
                top_p_in = gr.Slider(0.05, 1.0, value=1.0, step=0.05, label="Top-p (sampling only)")
                rep_in = gr.Slider(1.0, 1.5, value=DEFAULT_REPETITION_PENALTY, step=0.05, label="Repetition penalty")
            with gr.Row():
                submit_btn = gr.Button("Analyze", variant="primary", scale=2)
                reset_btn = gr.Button("Reset", variant="secondary", scale=1)
        with gr.Column(scale=4):
            summary_out = gr.HTML(_message("Upload a clip and press Analyze."))
            with gr.Accordion("JSON", open=False):
                json_out = gr.Code(show_label=False, language="json", wrap_lines=True)
            with gr.Accordion("Raw model output", open=False):
                raw_out = gr.Textbox(show_label=False, lines=6)
    gr.HTML(FOOTER)
    inputs = [
        audio_in,
        system_prompt_in,
        user_prompt_in,
        audio_first_in,
        constrain_in,
        max_tokens_in,
        temperature_in,
        top_p_in,
        rep_in,
        task_in,
    ]
    submit_btn.click(predict, inputs=inputs, outputs=[summary_out, json_out, raw_out], api_name="generate")
    task_in.change(_preset, inputs=task_in, outputs=[system_prompt_in, user_prompt_in, constrain_in], api_name=False)
    reset_btn.click(
        lambda: (
            *_preset(DEFAULT_TASK),
            True,
            DEFAULT_MAX_NEW_TOKENS,
            0.0,
            1.0,
            DEFAULT_REPETITION_PENALTY,
            DEFAULT_TASK,
        ),
        outputs=[system_prompt_in, user_prompt_in, constrain_in, audio_first_in, max_tokens_in]
        + [temperature_in, top_p_in, rep_in, task_in],
        api_name=False,
    )


def _gradio_auth() -> tuple[str, str] | None:
    raw_auth = os.getenv("GRADIO_AUTH", "").strip()
    if not raw_auth:
        return None
    if ":" not in raw_auth:
        raise ValueError("GRADIO_AUTH must be formatted as username:password")
    user, password = raw_auth.split(":", 1)
    return user, password


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        server_name=os.getenv("GRADIO_SERVER_NAME") or None,
        server_port=int(os.getenv("GRADIO_SERVER_PORT", "7860")),
        auth=_gradio_auth(),
        theme=THEME,
        css=CSS,
    )

	"""Inter-2 Audio playground: a Gradio UI and API around the Inter-2 Audio model.

	Task presets (tasks.py) fill in the prompts for emotion, or for engagement, social
	signals, emotion and vocal delivery together; each can optionally be constrained to
	its JSON schema. Prompts and decoding are editable.
	"""

	from __future__ import annotations

	import html
	import json
	import logging
	import os
	import sys
	import time
	from collections.abc import Callable
	from pathlib import Path
	from typing import Any, TypeVar

	import gradio as gr
	import spaces
	import torch
	from huggingface_hub import snapshot_download
	from inference import (
	DEFAULT_MAX_NEW_TOKENS,
	DEFAULT_REPETITION_PENALTY,
	ENCODER_WINDOW_S,
	TARGET_SR,
	AudioInputError,
	build_messages,
	format_summary,
	load_audio,
	parse_json,
	)
	from tasks import DEFAULT_TASK, TASKS
	from theme import CSS, FOOTER, HERO, THEME

	MODEL_ID = os.environ["MODEL_ID"] # a Space variable: the model repo (or a local copy of it)
	# cuda on a GPU or ZeroGPU Space (ZeroGPU has no GPU at import time), else cpu.
	_HAS_GPU = torch.cuda.is_available() or bool(os.getenv("SPACES_ZERO_GPU"))
	MODEL_DEVICE = os.getenv("MODEL_DEVICE") or ("cuda" if _HAS_GPU else "cpu")
	SPACES_GPU_DURATION = int(os.getenv("SPACES_GPU_DURATION", "90")) # rationale-heavy tasks write more
	logger = logging.getLogger(__name__)
	logging.basicConfig(level=logging.INFO)


	def _load() -> tuple[Any, Any]:
	"""(model, processor), via the loader that ships with the model repo."""
	model_dir = MODEL_ID if Path(MODEL_ID).is_dir() else snapshot_download(MODEL_ID)
	sys.path.insert(0, model_dir)
	from inter2_loader import load

	loaded = load(model_dir, MODEL_DEVICE)
	logger.info("Model loaded")
	return loaded


	def _compile_grammars() -> dict[str, Any]:
	"""Each task's answer schema for xgrammar, in compact JSON layout."""
	try:
	import xgrammar as xgr
	except ImportError:
	logger.warning("xgrammar is not installed; the schema toggle will be ignored.")
	return {}
	vocab_size = model.get_output_embeddings().weight.shape[0]
	compiler = xgr.GrammarCompiler(xgr.TokenizerInfo.from_huggingface(processor.tokenizer, vocab_size=vocab_size))
	return {
	name: compiler.compile_json_schema(
	json.dumps(task.schema), any_whitespace=False, indent=None, separators=(",", ":")
	)
	for name, task in TASKS.items()
	}


	model, processor = _load()
	grammars = _compile_grammars()

	F = TypeVar("F", bound=Callable[..., Any])


	def _spaces_gpu(func: F) -> F:
	setting = os.getenv("ENABLE_SPACES_GPU", "auto").lower()
	enabled = setting in {"1", "true", "yes"} or (setting == "auto" and bool(os.getenv("SPACE_ID")))
	return spaces.GPU(duration=SPACES_GPU_DURATION)(func) if enabled else func


	@_spaces_gpu
	def _generate(inputs: dict[str, Any], gen_kwargs: dict[str, Any], schema_task: str \| None) -> tuple[str, int]:
	inputs = {k: v.to(model.device) for k, v in inputs.items()}
	if "input_features" in inputs:
	inputs["input_features"] = inputs["input_features"].to(model.dtype)
	if schema_task is not None:
	from xgrammar.contrib.hf import LogitsProcessor # fresh per call: it keeps per-sequence state

	gen_kwargs = {**gen_kwargs, "logits_processor": [LogitsProcessor(grammars[schema_task])]}
	prompt_len = inputs["input_ids"].shape[-1]
	with torch.inference_mode():
	output_ids = model.generate(inputs, gen_kwargs)
	new_ids = output_ids[:, prompt_len:]
	text = processor.batch_decode(new_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
	return text.strip(), int(new_ids.shape[-1])


	def predict(
	audio_path: str \| None,
	system_prompt: str,
	user_prompt: str,
	audio_first: bool,
	constrain_json: bool,
	max_new_tokens: int,
	temperature: float,
	top_p: float,
	repetition_penalty: float,
	task: str = DEFAULT_TASK,
	) -> tuple[str, str, str]:
	"""Run one clip. Returns (summary HTML, parsed JSON, raw text). `task` picks the schema used
	when constrain_json is on; the prompts are whatever was passed in."""
	if not audio_path:
	return _message("Please upload or record an audio clip."), "", ""
	if not user_prompt.strip():
	return _message("The user prompt is empty."), "", ""
	try:
	wav, duration = load_audio(audio_path)
	chat = processor.apply_chat_template(
	build_messages(user_prompt, system_prompt.strip(), audio_first), tokenize=False, add_generation_prompt=True
	)
	inputs = dict(processor(text=[chat], audio=[wav], sampling_rate=TARGET_SR, return_tensors="pt", padding=True))
	gen_kwargs: dict[str, Any] = {
	"max_new_tokens": int(max_new_tokens),
	"repetition_penalty": float(repetition_penalty),
	"do_sample": temperature > 0,
	}
	if temperature > 0:
	gen_kwargs.update(temperature=float(temperature), top_p=float(top_p))
	start = time.time()
	schema_task = task if constrain_json and task in grammars else None
	raw, n_tokens = _generate(inputs, gen_kwargs, schema_task)
	logger.info("%s: %d tokens in %.1fs (constrained=%s)", task, n_tokens, time.time() - start, bool(schema_task))
	except Exception as exc:
	logger.exception("predict failed")
	if isinstance(exc, AudioInputError): # our own input checks; other errors may carry internals
	return _message(str(exc)), "", ""
	return _message("Something went wrong while analysing this clip. Please try again or use another file."), "", ""
	payload = parse_json(raw)
	pretty = json.dumps(payload, indent=2, ensure_ascii=False) if payload is not None else ""
	summary = format_summary(payload, duration if duration > ENCODER_WINDOW_S else None, n_tokens >= max_new_tokens)
	return summary, pretty, raw


	def _message(text: str) -> str:
	return f'<div class="ih-card"><p class="ih-note">{html.escape(text)}</p></div>'


	def _preset(task: str) -> tuple[str, str, bool]:
	preset = TASKS[task]
	return preset.system_prompt, preset.user_prompt, preset.constrain


	with gr.Blocks(title="Inter-2-Audio Playground") as demo:
	gr.HTML(HERO)
	with gr.Row(equal_height=False):
	with gr.Column(scale=5):
	audio_in = gr.Audio(label="Audio clip", sources=["upload", "microphone"], type="filepath")
	task_in = gr.Dropdown(list(TASKS), value=DEFAULT_TASK, label="Task")
	user_prompt_in = gr.Textbox(label="Prompt", value=TASKS[DEFAULT_TASK].user_prompt, lines=8, max_lines=12)
	system_prompt_in = gr.Textbox(
	label="System prompt (empty = the model's default)",
	value=TASKS[DEFAULT_TASK].system_prompt,
	lines=3,
	max_lines=8,
	)
	constrain_in = gr.Checkbox(
	value=TASKS[DEFAULT_TASK].constrain, label="Constrain the answer to the task's JSON schema"
	)
	with gr.Accordion("Advanced", open=False):
	audio_first_in = gr.Checkbox(value=True, label="Audio before text")
	max_tokens_in = gr.Slider(16, 1024, value=DEFAULT_MAX_NEW_TOKENS, step=16, label="Max new tokens")
	temperature_in = gr.Slider(0.0, 1.5, value=0.0, step=0.05, label="Temperature (0 = greedy)")
	top_p_in = gr.Slider(0.05, 1.0, value=1.0, step=0.05, label="Top-p (sampling only)")
	rep_in = gr.Slider(1.0, 1.5, value=DEFAULT_REPETITION_PENALTY, step=0.05, label="Repetition penalty")
	with gr.Row():
	submit_btn = gr.Button("Analyze", variant="primary", scale=2)
	reset_btn = gr.Button("Reset", variant="secondary", scale=1)
	with gr.Column(scale=4):
	summary_out = gr.HTML(_message("Upload a clip and press Analyze."))
	with gr.Accordion("JSON", open=False):
	json_out = gr.Code(show_label=False, language="json", wrap_lines=True)
	with gr.Accordion("Raw model output", open=False):
	raw_out = gr.Textbox(show_label=False, lines=6)
	gr.HTML(FOOTER)
	inputs = [
	audio_in,
	system_prompt_in,
	user_prompt_in,
	audio_first_in,
	constrain_in,
	max_tokens_in,
	temperature_in,
	top_p_in,
	rep_in,
	task_in,
	]
	submit_btn.click(predict, inputs=inputs, outputs=[summary_out, json_out, raw_out], api_name="generate")
	task_in.change(_preset, inputs=task_in, outputs=[system_prompt_in, user_prompt_in, constrain_in], api_name=False)
	reset_btn.click(
	lambda: (
	*_preset(DEFAULT_TASK),
	True,
	DEFAULT_MAX_NEW_TOKENS,
	0.0,
	1.0,
	DEFAULT_REPETITION_PENALTY,
	DEFAULT_TASK,
	),
	outputs=[system_prompt_in, user_prompt_in, constrain_in, audio_first_in, max_tokens_in]
	+ [temperature_in, top_p_in, rep_in, task_in],
	api_name=False,
	)


	def _gradio_auth() -> tuple[str, str] \| None:
	raw_auth = os.getenv("GRADIO_AUTH", "").strip()
	if not raw_auth:
	return None
	if ":" not in raw_auth:
	raise ValueError("GRADIO_AUTH must be formatted as username:password")
	user, password = raw_auth.split(":", 1)
	return user, password


	if __name__ == "__main__":
	demo.queue(default_concurrency_limit=1).launch(
	server_name=os.getenv("GRADIO_SERVER_NAME") or None,
	server_port=int(os.getenv("GRADIO_SERVER_PORT", "7860")),
	auth=_gradio_auth(),
	theme=THEME,
	css=CSS,
	)