app.py

"""
SATA Demo

Retarget:
  Task 1: Human -> Human
  Task 2: Animo -> Animo
  Task 3: Human -> Animo
  Task 4: Animo -> Human

Generation:
  Task 5: Text -> Motion
"""
import os
import sys
import threading
import time
import uuid

# Path setup
RELEASE_DIR = os.path.dirname(os.path.abspath(__file__))
GRADIO_DIR  = os.path.join(RELEASE_DIR, 'gradio')
sys.path.insert(0, GRADIO_DIR)
sys.path.insert(0, os.path.join(RELEASE_DIR, 'src'))
sys.path.insert(0, os.path.join(RELEASE_DIR, 'src', 'fairmotion'))
sys.path.insert(0, os.path.join(RELEASE_DIR, 'src', 'Visualization'))

# Gradio 6 reads DEFAULT_TEMP_DIR at import time, so set it first.
_gradio_tmp = os.path.join(GRADIO_DIR, ".gradio_tmp")
os.makedirs(_gradio_tmp, exist_ok=True)
os.environ["GRADIO_TEMP_DIR"] = _gradio_tmp

import gradio as gr

from config import (
    OUTPUT_DIR, MODEL_TYPE, MODEL_EPOCH, DEVICE,
    TASK1_Z_OPTIONS, TASK1_SKEL_OPTIONS,
    TASK2_Z_OPTIONS, TASK2_SKEL_OPTIONS,
    TASK3_PAIRS, TASK4_PAIRS,
    GEN_DEVICE, GEN_DEFAULT_MOTION_FRAMES,
    GEN_MIN_MOTION_FRAMES, GEN_MAX_MOTION_FRAMES,
    GEN_MODEL_OPTIONS, HUMAN_GEN_SKEL_OPTIONS,
)
from bvh_stick_viewer import get_bvh_viewer_html, iframe_html_for_motion


# ════════════════════════════════════════════════════════════
# Model loading (lazy and thread-safe)
# ════════════════════════════════════════════════════════════
_model_state = None
_model_lock  = threading.Lock()

def get_model_state():
    global _model_state
    if _model_state is None:
        with _model_lock:
            if _model_state is None:
                print("[model] Loading model...")
                from sata.utils.model_loading import load_model_by_type
                model, cfg, ms_dict = load_model_by_type(MODEL_TYPE, MODEL_EPOCH, DEVICE)
                model = model.to(DEVICE)
                model.eval()
                _model_state = (model, cfg, ms_dict)
                print("[model] Model loaded.")
    return _model_state


# ════════════════════════════════════════════════════════════
# Decode and render helpers
# ════════════════════════════════════════════════════════════
VIEWER_H = 360


def empty_viewer(title):
    return iframe_html_for_motion(None, title=title, height=VIEWER_H)


def option_label(options, key):
    return options[key].get('label', key)


def dropdown_choices(options):
    return [option_label(options, key) for key in options.keys()]


def default_choice(options):
    return option_label(options, list(options.keys())[0])


def unique_output_name(base_name):
    stamp = time.strftime('%y%m%d%H%M%S')
    suffix = uuid.uuid4().hex[:8]
    return f"{base_name}__{stamp}_{suffix}"


def resolve_option_key(options, selected):
    if selected in options:
        return selected
    for key in options.keys():
        if option_label(options, key) == selected:
            return key
    raise KeyError(selected)


def run_retarget(z_path, skel_path, output_name, task_name):
    """
    Decode to BVH and return the Three.js stick viewer.
    """
    from decode_z_sliding import decode_z_skel_to_bvh

    task_out_dir = os.path.join(OUTPUT_DIR, task_name)
    os.makedirs(task_out_dir, exist_ok=True)

    try:
        bvh_path = decode_z_skel_to_bvh(
            z_path=z_path,
            skel_path=skel_path,
            model_epoch=MODEL_EPOCH,
            output_dir=task_out_dir,
            output_name=output_name,
            device=DEVICE,
            model_type=MODEL_TYPE,
            model_state=get_model_state(),
        )
        return get_bvh_viewer_html(
            bvh_path, title=output_name, height=VIEWER_H
        )
    except Exception as e:
        print(f"[error] retarget failed: {e}")
        import traceback; traceback.print_exc()
        return empty_viewer("Retarget failed")


# ════════════════════════════════════════════════════════════
# Tab callbacks
# ════════════════════════════════════════════════════════════

# ── Task 1 ──────────────────────────────────────────────────
def t1_src_change(z_key):
    z_key = resolve_option_key(TASK1_Z_OPTIONS, z_key)
    info = TASK1_Z_OPTIONS[z_key]
    return get_bvh_viewer_html(
        info['src_bvh'], title=f"Source: {option_label(TASK1_Z_OPTIONS, z_key)}", height=VIEWER_H,
    )

def t1_tgt_change(skel_key):
    skel_key = resolve_option_key(TASK1_SKEL_OPTIONS, skel_key)
    info = TASK1_SKEL_OPTIONS[skel_key]
    return get_bvh_viewer_html(
        info['tgt_bvh'], title=f"Target skeleton: {option_label(TASK1_SKEL_OPTIONS, skel_key)}",
        height=VIEWER_H, max_frames=1,
    )

def t1_run(z_key, skel_key):
    z_key = resolve_option_key(TASK1_Z_OPTIONS, z_key)
    skel_key = resolve_option_key(TASK1_SKEL_OPTIONS, skel_key)
    z    = TASK1_Z_OPTIONS[z_key]
    skel = TASK1_SKEL_OPTIONS[skel_key]
    name = unique_output_name(f"{z['name']}__{skel['name']}")
    return run_retarget(z['z_path'], skel['skel_path'], name, 'task1')

# ── Task 2 ──────────────────────────────────────────────────
def t2_src_change(z_key):
    z_key = resolve_option_key(TASK2_Z_OPTIONS, z_key)
    info = TASK2_Z_OPTIONS[z_key]
    return get_bvh_viewer_html(
        info['src_bvh'], title=f"Source: {option_label(TASK2_Z_OPTIONS, z_key)}", height=VIEWER_H,
    )

def t2_tgt_change(skel_key):
    skel_key = resolve_option_key(TASK2_SKEL_OPTIONS, skel_key)
    info = TASK2_SKEL_OPTIONS[skel_key]
    return get_bvh_viewer_html(
        info['tgt_bvh'], title=f"Target skeleton: {option_label(TASK2_SKEL_OPTIONS, skel_key)}",
        height=VIEWER_H, max_frames=1,
    )

def t2_run(z_key, skel_key):
    z_key = resolve_option_key(TASK2_Z_OPTIONS, z_key)
    skel_key = resolve_option_key(TASK2_SKEL_OPTIONS, skel_key)
    z    = TASK2_Z_OPTIONS[z_key]
    skel = TASK2_SKEL_OPTIONS[skel_key]
    name = unique_output_name(f"{z['name']}__{skel['name']}")
    return run_retarget(z['z_path'], skel['skel_path'], name, 'task2')

# ── Task 3 ──────────────────────────────────────────────────
def t3_pair_change(pair_key):
    pair_key = resolve_option_key(TASK3_PAIRS, pair_key)
    pair = TASK3_PAIRS[pair_key]
    label = option_label(TASK3_PAIRS, pair_key)
    src = get_bvh_viewer_html(
        pair['src_bvh'], title=f"Source: {label}", height=VIEWER_H,
    )
    tgt = get_bvh_viewer_html(
        pair['tgt_bvh'], title=f"Target skeleton: {label}",
        height=VIEWER_H, max_frames=1,
    )
    return src, tgt

def t3_run(pair_key):
    pair_key = resolve_option_key(TASK3_PAIRS, pair_key)
    pair = TASK3_PAIRS[pair_key]
    name = unique_output_name(pair['output_name'])
    return run_retarget(pair['z_path'], pair['skel_path'], name, 'task3')

# ── Task 4 ──────────────────────────────────────────────────
def t4_pair_change(pair_key):
    pair_key = resolve_option_key(TASK4_PAIRS, pair_key)
    pair = TASK4_PAIRS[pair_key]
    label = option_label(TASK4_PAIRS, pair_key)
    src = get_bvh_viewer_html(
        pair['src_bvh'], title=f"Source: {label}", height=VIEWER_H,
    )
    tgt = get_bvh_viewer_html(
        pair['tgt_bvh'], title=f"Target skeleton: {label}",
        height=VIEWER_H, max_frames=1,
    )
    return src, tgt

def t4_run(pair_key):
    pair_key = resolve_option_key(TASK4_PAIRS, pair_key)
    pair = TASK4_PAIRS[pair_key]
    name = unique_output_name(pair['output_name'])
    return run_retarget(pair['z_path'], pair['skel_path'], name, 'task4')


# ════════════════════════════════════════════════════════════
# Task 5: generation callbacks
# ════════════════════════════════════════════════════════════
def run_generation(text: str, skel_key: str, generation_label: str,
                   motion_frames: int, skel_options: dict) -> str:
    """
    Generation pipeline: text -> latent/code -> BVH -> viewer.
    """
    from gen_runner import generate_motion_from_text

    text = text.strip()
    if not text:
        return empty_viewer("Enter a text prompt")

    motion_frames = int(round(float(motion_frames)))
    motion_frames = min(max(motion_frames, GEN_MIN_MOTION_FRAMES), GEN_MAX_MOTION_FRAMES)

    if not skel_key:
        skel_key = list(skel_options.keys())[0]
    else:
        skel_key = resolve_option_key(skel_options, skel_key)

    skel_info   = skel_options[skel_key]
    skel_path   = skel_info['skel_path']
    display_name = unique_output_name(skel_info['name'])
    generation_type = GEN_MODEL_OPTIONS[generation_label]
    gen_out_dir = os.path.join(OUTPUT_DIR, 'generation', generation_type)
    os.makedirs(gen_out_dir, exist_ok=True)

    try:
        bvh_path = generate_motion_from_text(
            text=text,
            skel_npz_path=skel_path,
            generation_type=generation_type,
            output_dir=gen_out_dir,
            device_str=GEN_DEVICE,
            skel_name=skel_info['name'],
            motion_frames=motion_frames,
        )
        if not bvh_path:
            return empty_viewer("Generation failed")
        return get_bvh_viewer_html(
            bvh_path,
            title=f"{display_name} - {generation_label} - {motion_frames} frames - {text}",
            height=VIEWER_H,
        )
    except Exception as e:
        print(f"[error] generation failed: {e}")
        import traceback; traceback.print_exc()
        return empty_viewer("Generation failed")


def t5_run(text, skel_key, generation_label, motion_frames):
    return run_generation(text, skel_key, generation_label, motion_frames, HUMAN_GEN_SKEL_OPTIONS)


def t5_tgt_change(skel_key):
    info = HUMAN_GEN_SKEL_OPTIONS[skel_key]
    return get_bvh_viewer_html(
        info['tgt_bvh'], title=f"Target skeleton: {skel_key}",
        height=VIEWER_H, max_frames=1,
    )


# ════════════════════════════════════════════════════════════
# Page-load previews
# ════════════════════════════════════════════════════════════
def on_load():
    try:
        return (
            t1_src_change(list(TASK1_Z_OPTIONS.keys())[0]),
            t1_tgt_change(list(TASK1_SKEL_OPTIONS.keys())[0]),
            empty_viewer("Source motion"),
            empty_viewer("Target skeleton"),
            empty_viewer("Source human motion"),
            empty_viewer("Target animal skeleton"),
            empty_viewer("Source animal motion"),
            empty_viewer("Target human skeleton"),
            t5_tgt_change('mousey'),
        )
    except Exception as e:
        print(f"[on_load] Preview preload failed, possibly because a BVH file is missing: {e}")
        return (None,) * 9


# ════════════════════════════════════════════════════════════
# Gradio UI
# ════════════════════════════════════════════════════════════
DEFAULT_TEXT_PROMPT = "A person walks forward, bends down to pick something up off the ground."


with gr.Blocks(title="SATA Demo") as app:

    gr.Markdown(
        """
        # SATA: Semantic-Aware Motion Encoding for Topology-Agnostic Character Animation
        Interactive visualization for motion retargeting and text-to-motion generation.
        """
    )

    gr.Markdown("## Motion Retargeting")

    with gr.Tabs() as retarget_tabs:

        # ╔══════════════════════════════════════════════════╗
        # ║  Task 1: Human -> Human                          ║
        # ╚══════════════════════════════════════════════════╝
        with gr.TabItem("Task 1 - Human -> Human") as t1_tab:
            gr.Markdown("**Free selection**: choose any source motion and target skeleton.")
            with gr.Row():
                t1_z_dd   = gr.Dropdown(dropdown_choices(TASK1_Z_OPTIONS),
                                        value=default_choice(TASK1_Z_OPTIONS),
                                        label="Source motion", scale=2)
                t1_sk_dd  = gr.Dropdown(dropdown_choices(TASK1_SKEL_OPTIONS),
                                        value=default_choice(TASK1_SKEL_OPTIONS),
                                        label="Target skeleton", scale=2)
                t1_btn    = gr.Button("Run Model", variant="primary", scale=1)
            with gr.Row():
                t1_v_src  = gr.HTML(value=empty_viewer("Source motion"))
                t1_v_tgt  = gr.HTML(value=empty_viewer("Target skeleton"))
                t1_v_out  = gr.HTML(value=empty_viewer("Retarget result"))

            t1_z_dd.change(t1_src_change,  t1_z_dd,           t1_v_src)
            t1_sk_dd.change(t1_tgt_change, t1_sk_dd,          t1_v_tgt)
            t1_btn.click(t1_run,           [t1_z_dd, t1_sk_dd], t1_v_out)

        # ╔══════════════════════════════════════════════════╗
        # ║  Task 2: Animo -> Animo                          ║
        # ╚══════════════════════════════════════════════════╝
        with gr.TabItem("Task 2 - Animo -> Animo") as t2_tab:
            gr.Markdown("**Zero-shot retargeting**: choose any animal source motion and target skeleton.")
            with gr.Row():
                t2_z_dd   = gr.Dropdown(dropdown_choices(TASK2_Z_OPTIONS),
                                        value=default_choice(TASK2_Z_OPTIONS),
                                        label="Source motion", scale=2)
                t2_sk_dd  = gr.Dropdown(dropdown_choices(TASK2_SKEL_OPTIONS),
                                        value=default_choice(TASK2_SKEL_OPTIONS),
                                        label="Target skeleton", scale=2)
                t2_btn    = gr.Button("Run Model", variant="primary", scale=1)
            with gr.Row():
                t2_v_src  = gr.HTML(value=empty_viewer("Source motion"))
                t2_v_tgt  = gr.HTML(value=empty_viewer("Target skeleton"))
                t2_v_out  = gr.HTML(value=empty_viewer("Retarget result"))

            t2_z_dd.change(t2_src_change,  t2_z_dd,           t2_v_src)
            t2_sk_dd.change(t2_tgt_change, t2_sk_dd,          t2_v_tgt)
            t2_btn.click(t2_run,           [t2_z_dd, t2_sk_dd], t2_v_out)

        # ╔══════════════════════════════════════════════════╗
        # ║  Task 3: Human -> Animo                          ║
        # ╚══════════════════════════════════════════════════╝
        with gr.TabItem("Task 3 - Human -> Animo") as t3_tab:
            gr.Markdown("**Zero-shot retargeting**: selected human-to-animal samples.")
            with gr.Row():
                t3_dd     = gr.Dropdown(dropdown_choices(TASK3_PAIRS),
                                        value=default_choice(TASK3_PAIRS),
                                        label="Sample pair", scale=3)
                t3_btn    = gr.Button("Run Model", variant="primary", scale=1)
            with gr.Row():
                t3_v_src  = gr.HTML(value=empty_viewer("Source human motion"))
                t3_v_tgt  = gr.HTML(value=empty_viewer("Target animal skeleton"))
                t3_v_out  = gr.HTML(value=empty_viewer("Retarget result"))

            t3_dd.change(t3_pair_change, t3_dd, [t3_v_src, t3_v_tgt])
            t3_btn.click(t3_run,         t3_dd,  t3_v_out)

        # ╔══════════════════════════════════════════════════╗
        # ║  Task 4: Animo -> Human                          ║
        # ╚══════════════════════════════════════════════════╝
        with gr.TabItem("Task 4 - Animo -> Human") as t4_tab:
            gr.Markdown("**Zero-shot retargeting**: selected animal-to-human samples.")
            with gr.Row():
                t4_dd     = gr.Dropdown(dropdown_choices(TASK4_PAIRS),
                                        value=default_choice(TASK4_PAIRS),
                                        label="Sample pair", scale=3)
                t4_btn    = gr.Button("Run Model", variant="primary", scale=1)
            with gr.Row():
                t4_v_src  = gr.HTML(value=empty_viewer("Source animal motion"))
                t4_v_tgt  = gr.HTML(value=empty_viewer("Target human skeleton"))
                t4_v_out  = gr.HTML(value=empty_viewer("Retarget result"))

            t4_dd.change(t4_pair_change, t4_dd, [t4_v_src, t4_v_tgt])
            t4_btn.click(t4_run,         t4_dd,  t4_v_out)

        # Refresh source/target previews whenever a tab is selected.
        t1_tab.select(
            lambda z_key, skel_key: (t1_src_change(z_key), t1_tgt_change(skel_key)),
            [t1_z_dd, t1_sk_dd],
            [t1_v_src, t1_v_tgt],
        )
        t2_tab.select(
            lambda z_key, skel_key: (t2_src_change(z_key), t2_tgt_change(skel_key)),
            [t2_z_dd, t2_sk_dd],
            [t2_v_src, t2_v_tgt],
       )
        t3_tab.select(t3_pair_change, t3_dd, [t3_v_src, t3_v_tgt])
        t4_tab.select(t4_pair_change, t4_dd, [t4_v_src, t4_v_tgt])

    # ════════════════════════════════════════════════════════
    # Generation
    # ════════════════════════════════════════════════════════
    gr.Markdown("---")
    gr.Markdown("## Text-to-motion Generation")
    gr.Markdown(
        "Enter a text prompt, choose a target skeleton and generation method, then click **Generate**."
    )

    with gr.Tabs():

        # ╔══════════════════════════════════════════════════╗
        # ║  Task 5: Text-to-motion generation               ║
        # ╚══════════════════════════════════════════════════╝
        with gr.TabItem("Task 5 - Text-to-motion Generation"):
            gr.Markdown("Generate a motion from text and decode it onto the selected character skeleton.")
            gr.Markdown("Generation models MDM and MoMask are large; each run can take about 2 minutes.")
            with gr.Row():
                t5_text = gr.Textbox(
                    label="Text prompt (English)",
                    value=DEFAULT_TEXT_PROMPT,
                    lines=2,
                )
            with gr.Row():
                t5_sk_dd = gr.Dropdown(
                    dropdown_choices(HUMAN_GEN_SKEL_OPTIONS),
                    value="mousey",
                    label="Target skeleton",
                    scale=2,
                )
                t5_model_dd = gr.Dropdown(
                    list(GEN_MODEL_OPTIONS.keys()),
                    value="RVQ",
                    label="Method",
                    scale=1,
                )
                t5_motion_frames = gr.Slider(
                    minimum=GEN_MIN_MOTION_FRAMES,
                    maximum=GEN_MAX_MOTION_FRAMES,
                    value=GEN_DEFAULT_MOTION_FRAMES,
                    step=1,
                    label="Motion length",
                    scale=2,
                )
                t5_btn = gr.Button("Generate", variant="primary", scale=1)
            with gr.Row():
                t5_v_tgt = gr.HTML(value=empty_viewer("Selected target skeleton"))
                t5_v_out = gr.HTML(value=empty_viewer("Generation result"))

            t5_sk_dd.change(t5_tgt_change, t5_sk_dd, t5_v_tgt)
            t5_btn.click(t5_run, [t5_text, t5_sk_dd, t5_model_dd, t5_motion_frames], t5_v_out)

    # Preload default interactive previews.
    app.load(
        fn=on_load,
        outputs=[t1_v_src, t1_v_tgt,
                 t2_v_src, t2_v_tgt,
                 t3_v_src, t3_v_tgt,
                 t4_v_src, t4_v_tgt,
                 t5_v_tgt],
    )


def _env_flag(name: str, default: bool = False) -> bool:
    value = os.environ.get(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


if __name__ == "__main__":
    app.launch(
        server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
        server_port=int(os.environ.get("PORT", "7860")),
        share=_env_flag("GRADIO_SHARE", False),
        debug=_env_flag("GRADIO_DEBUG", False),
        show_error=True,
        theme=gr.themes.Soft(primary_hue="blue"),
    )