""" 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"), )