app.py

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces
import sys
import shutil
import tempfile
import json
import math
import time
import gc
from pathlib import Path

import torch
import torch.distributed as dist
import numpy as np
import imageio
import gradio as gr
from PIL import Image
from huggingface_hub import snapshot_download

# ── Paths ──────────────────────────────────────────────────────────────────
SPACE_ROOT = Path(__file__).parent
CKPT_DIR = Path("/tmp/checkpoints")
CKPT_DIR.mkdir(parents=True, exist_ok=True)

# ── Download all model weights at module scope ──────────────────────────────
print("[HOMA] Downloading model weights...", flush=True)

# HOMA checkpoint - only the model states file
snapshot_download(
    "ProAudience/homa_checkpoint",
    local_dir=str(CKPT_DIR / "homa_checkpoint"),
    repo_type="model",
    allow_patterns=["tp_rank_00_pp_rank_00_model_states.pt", "*.md"],
)
print("[HOMA] HOMA checkpoint downloaded.", flush=True)

# Base models
MODEL_BASE = CKPT_DIR / "pretrained_models"
MODEL_BASE.mkdir(parents=True, exist_ok=True)

# DINOv2-giant (required by model architecture - 1536 dim)
DINO_DIR = MODEL_BASE / "dinov2-giant"
snapshot_download(
    "facebook/dinov2-giant",
    local_dir=str(DINO_DIR),
    repo_type="model",
    allow_patterns=["*.safetensors", "config.json", "preprocessor_config.json"],
)
print("[HOMA] DINOv2-giant downloaded.", flush=True)

# CLIP-L
snapshot_download(
    "openai/clip-vit-large-patch14",
    local_dir=str(MODEL_BASE / "openai_clip-vit-large-patch14"),
    repo_type="model",
    allow_patterns=["pytorch_model.bin", "config.json", "tokenizer*", "vocab.json", "merges.txt", "special_tokens_map.json", "preprocessor_config.json"],
)
print("[HOMA] CLIP-L downloaded.", flush=True)

# LLaVA-Llama-3-8B (used as text encoder AND for tokenizer extraction)
LLAVA_DIR = MODEL_BASE / "llava-llama-3-8b-v1_1-transformers"
snapshot_download(
    "xtuner/llava-llama-3-8b-v1_1-transformers",
    local_dir=str(LLAVA_DIR),
    repo_type="model",
    allow_patterns=["model-*.safetensors", "model.safetensors.index.json", "config.json", "generation_config.json", "preprocessor_config.json", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json"],
)
print("[HOMA] LLaVA-Llama-3-8B downloaded.", flush=True)

# Extract pure LLaMA tokenizer from LLaVA (just copy tokenizer files, no model loading)
pure_llama_path = MODEL_BASE / "llava-llama-3-8b-v1_1-pure-llama"
if not pure_llama_path.exists():
    print("[HOMA] Copying LLaMA tokenizer from LLaVA...", flush=True)
    pure_llama_path.mkdir(parents=True, exist_ok=True)
    for fname in ["tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
                   "added_tokens.json", "config.json", "generation_config.json"]:
        src = LLAVA_DIR / fname
        if src.exists():
            shutil.copy2(src, pure_llama_path / fname)
    print("[HOMA] LLaMA tokenizer ready.", flush=True)

# HunyuanVideo 3D VAE
vae_tmp = CKPT_DIR / "_hyvae_tmp"
vae_dir = MODEL_BASE / "vae_3d" / "hyvae"
vae_dir.mkdir(parents=True, exist_ok=True)
snapshot_download(
    "tencent/HunyuanVideo",
    local_dir=str(vae_tmp),
    repo_type="model",
    allow_patterns=["hunyuan-video-t2v-720p/vae/*"],
)
hyvae_src = vae_tmp / "hunyuan-video-t2v-720p" / "vae"
for f in hyvae_src.iterdir():
    shutil.copy2(f, vae_dir / f.name)
# Cleanup temp
shutil.rmtree(vae_tmp, ignore_errors=True)
print("[HOMA] VAE downloaded.", flush=True)

# Aux models
MODEL_AUX = CKPT_DIR / "aux"
MODEL_AUX.mkdir(parents=True, exist_ok=True)

snapshot_download(
    "openai/whisper-tiny",
    local_dir=str(MODEL_AUX / "ckpts" / "whisper-tiny"),
    repo_type="model",
    allow_patterns=["model.safetensors", "config.json", "generation_config.json", "normalizer.json", "preprocessor_config.json", "tokenizer.json", "tokenizer_config.json", "added_tokens.json", "special_tokens_map.json"],
)
print("[HOMA] Whisper-tiny downloaded.", flush=True)

# Face detector
detface_tmp = CKPT_DIR / "_detface_tmp"
detface_dir = MODEL_AUX / "ckpts" / "det_align"
detface_dir.mkdir(parents=True, exist_ok=True)
snapshot_download(
    "tencent/HunyuanVideo-Avatar",
    local_dir=str(detface_tmp),
    repo_type="model",
    allow_patterns=["ckpts/det_align/detface.pt"],
)
detface_src = detface_tmp / "ckpts" / "det_align" / "detface.pt"
if detface_src.exists():
    shutil.copy2(detface_src, detface_dir / "detface.pt")
shutil.rmtree(detface_tmp, ignore_errors=True)
print("[HOMA] Face detector downloaded.", flush=True)

# ── Set environment variables ─────────────────────────────────────────────
os.environ["MODEL_BASE"] = str(MODEL_BASE)
os.environ["MODEL_AUX_PATH"] = str(MODEL_AUX)
os.environ["RELEASE_ROOT"] = str(SPACE_ROOT)
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
os.environ.setdefault("MASTER_PORT", "29500")
os.environ.setdefault("RANK", "0")
os.environ.setdefault("WORLD_SIZE", "1")
os.environ.setdefault("LOCAL_RANK", "0")
os.environ["SP_SIZE"] = "1"

sys.path.insert(0, str(SPACE_ROOT))

# ── Initialize the inference pipeline ──────────────────────────────────────
from hyavatar.infer_args import parse_inference_args
from hyavatar.runtime import set_global_args, get_args
from hyavatar.runtime.ulysses import init_sequence_parallel
from hyavatar.utils.torch_utils import set_manual_seed
from hyavatar.sample_all_in_one_infer import Evaluator, resolve_sparse_pose_parts
from loguru import logger

sys.argv = [
    "app.py",
    "--load", str(CKPT_DIR / "homa_checkpoint"),
    "--inference-meta-file", str(SPACE_ROOT / "examples" / "demo_cases.csv"),
    "--seed", "128",
    "--infer-steps", "50",
    "--cfg-scale", "3.0",
    "--pose-control-type", "full",
    "--pose-parts", "both_arms",
    "--attn-mode", "torch",
]
parse_inference_args()
args = get_args()

# Initialize distributed (single-process for ZeroGPU)
if not dist.is_initialized():
    dist.init_process_group(backend="gloo", init_method="env://")
init_sequence_parallel(1)
set_manual_seed(args.seed)

print("[HOMA] Loading Evaluator (model, VAE, text encoders, DINO)...", flush=True)
evaluator = Evaluator.from_pretrained(
    mode="eval",
    ckpt=str(CKPT_DIR / "homa_checkpoint" / "tp_rank_00_pp_rank_00_model_states.pt"),
    extra_model_base=None,
    args_path=None,
    world_size=1,
    rank=0,
    device="cuda",
    logger=logger,
    flow_shift=1.0,
)
print("[HOMA] Evaluator loaded.", flush=True)


# ── Inference helper ───────────────────────────────────────────────────────
def run_inference(
    human_image_path: str,
    object_image_path: str,
    audio_path: str,
    prompt: str,
    infer_steps: int = 50,
    cfg_scale: float = 3.0,
    seed: int = 128,
    pose_control_type: str = "full",
):
    """Run HOMA inference for a single case."""
    import pandas as pd
    from hyavatar.sample_all_in_one_infer import encode_audio, get_facemask
    from hyavatar.data_kits.face_align import AlignImage
    from transformers import WhisperModel, AutoFeatureExtractor
    from einops import rearrange
    from hyavatar.runtime import print_rank_0
    from torch.utils.data import DataLoader

    # Update args
    args.infer_steps = infer_steps
    args.cfg_scale = cfg_scale
    args.seed = seed
    args.pose_control_type = pose_control_type
    args.pose_parts = "both_arms"
    set_manual_seed(seed)

    # Build CSV for this single case
    tmp_dir = Path(tempfile.mkdtemp())
    case_dir = tmp_dir / "case"
    case_dir.mkdir(parents=True, exist_ok=True)

    human_img = Image.open(human_image_path).convert("RGB")
    human_img.save(case_dir / "human.png")

    object_img = Image.open(object_image_path).convert("RGBA")
    object_img.save(case_dir / "object.png")
    object_img.convert("RGB").save(case_dir / "object_llava.png")

    shutil.copy2(audio_path, case_dir / "audio.wav")

    # Use pre-computed dwpose from example cases
    default_dwpose = SPACE_ROOT / "examples" / "assets" / "case_a" / "dwpose.pkl"
    default_bbox = SPACE_ROOT / "examples" / "assets" / "case_a" / "object_bbox.json"
    shutil.copy2(default_dwpose, case_dir / "dwpose.pkl")
    shutil.copy2(default_bbox, case_dir / "object_bbox.json")

    csv_path = tmp_dir / "cases.csv"
    df = pd.DataFrame([{
        "videoid": "user_case",
        "image": str(case_dir / "human.png"),
        "audio": str(case_dir / "audio.wav"),
        "prompt": prompt,
        "fps": 25.0,
        "dwpose": str(case_dir / "dwpose.pkl"),
        "object": str(case_dir / "object.png"),
        "object_llava": str(case_dir / "object_llava.png"),
        "object_bbox": str(case_dir / "object_bbox.json"),
        "draw_type": "full",
        "expand_ratio": 0.2,
    }])
    df.to_csv(csv_path, index=False)

    # Build data loader
    MODEL_AUX_PATH = os.environ.get("MODEL_AUX_PATH")
    wav2vec = WhisperModel.from_pretrained(f"{MODEL_AUX_PATH}/ckpts/whisper-tiny/").to(
        device=torch.device("cuda"), dtype=torch.float32
    )
    wav2vec.requires_grad_(False)

    det_path = os.path.join(MODEL_AUX_PATH, "ckpts", "det_align", "detface.pt")
    align_instance = AlignImage("cuda", det_path=det_path)
    feature_extractor = AutoFeatureExtractor.from_pretrained(f"{MODEL_AUX_PATH}/ckpts/whisper-tiny/")

    draw_type_override = None if args.pose_control_type == "csv" else args.pose_control_type
    sparse_pose_parts = resolve_sparse_pose_parts(args.pose_parts)

    kwargs = {
        "load_pixel_values": False,
        "ratios": (8, 8, 4),
        "text_encoder": evaluator.text_encoder,
        "text_encoder_2": evaluator.text_encoder_2,
        "wav2vec": wav2vec,
        "feature_extractor": feature_extractor,
        "test_len": args.text_len,
        "vae_ratios": (8, 8, 4),
        "uncond_ref_p": 0.1,
        "uncond_text_p": 0.1,
        "uncond_audio_p": 0.1,
        "bbox_only": False,
        "first_n_frame_only": False,
        "first_n_frame": 129,
        "draw_type": draw_type_override,
        "draw_parts": sparse_pose_parts,
        "draw_hand": True,
        "gen_long_video": False,
        "obj_align": True,
        "apply_offset": True,
        "offset_scale": 0.15,
        "load_lama_info": False,
        "load_sam_bbox_info": True,
        "is_user_edited": False,
        "left_or_right": "right",
        "use_min_size": False,
        "largest_resolution": (64, 112),
    }

    caption_sample_ratio = '{"long caption": 0.50, "short caption": 0.45,"background": 0.8,"shot type":0.8,"style": 0.8,"light":0.8,"atmosphere":0.8,"camera movement":0.8}'
    from hyavatar.data_kits.video_loader_all_in_one_infer import VideoAudioTextLoaderVal
    video_dataset = VideoAudioTextLoaderVal(
        meta_file=[str(csv_path)],
        sample_n_frames=129,
        resolution=(512, 512),
        logger=None,
        dtype_encode="video",
        resolution_type="540p",
        caption_sample_ratio=caption_sample_ratio,
        **kwargs,
    )
    video_loader = DataLoader(
        video_dataset,
        batch_size=1,
        shuffle=False,
        num_workers=0,
        pin_memory=True,
        drop_last=True,
    )
    args.video_sampler = None
    args.patch_size = evaluator.model.patch_size
    args.hidden_size = evaluator.model.hidden_size
    args.num_heads = evaluator.model.num_heads
    args.rope_dim_list = evaluator.model.rope_dim_list

    output_dir = tmp_dir / "output"
    output_dir.mkdir(parents=True, exist_ok=True)

    for batch_id, batch in enumerate(video_loader):
        evaluator.vae.enable_tiling()

        prompt_str = batch["text_prompt"][0]
        image_path = str(batch["image_path"][0])
        audio_path = str(batch["audio_path"][0])
        videoid = batch["videoid"][0]
        fps = batch["fps"].to(evaluator.device)
        audio_prompts_raw = batch["audio_prompts"].to(evaluator.device)
        sample_n_frames = batch["sample_n_frames"].item()
        weight_dtype = torch.float16

        pixel_value_ref = batch["pixel_value_ref"].to(evaluator.device)
        uncond_pixel_value_ref = torch.zeros_like(pixel_value_ref)
        pixel_value_ref = pixel_value_ref / 127.5 - 1.
        uncond_pixel_value_ref = uncond_pixel_value_ref * 2 - 1

        pixel_value_ref_object = batch["pixel_value_ref_object"].to(evaluator.device)
        uncond_pixel_value_ref_object = torch.zeros_like(pixel_value_ref_object)
        pixel_value_ref_object = pixel_value_ref_object / 127.5 - 1.
        uncond_pixel_value_ref_object = uncond_pixel_value_ref_object * 2 - 1

        pixel_value_ref_object_seq = batch["pixel_value_ref_object_seq"].to(evaluator.device)
        uncond_pixel_value_ref_object_seq = torch.zeros_like(pixel_value_ref_object_seq)
        pixel_value_ref_object_seq = pixel_value_ref_object_seq / 127.5 - 1.
        uncond_pixel_value_ref_object_seq = uncond_pixel_value_ref_object_seq * 2 - 1
        pixel_value_ref_object_seq = rearrange(pixel_value_ref_object_seq, "b f c h w -> b c f h w")
        uncond_pixel_value_ref_object_seq = rearrange(uncond_pixel_value_ref_object_seq, "b f c h w -> b c f h w")

        face_masks = get_facemask((pixel_value_ref.clone() + 1.) * 127.5, align_instance, area=1.5)

        pixel_value_dwpose = batch["pixel_value_sparse_dwpose"]
        uncond_pixel_value_dwpose = torch.zeros_like(pixel_value_dwpose)
        pixel_value_dwpose = pixel_value_dwpose / 127.5 - 1.
        uncond_pixel_value_dwpose = uncond_pixel_value_dwpose * 2 - 1
        pixel_value_dwpose = pixel_value_dwpose.to(evaluator.device)
        uncond_pixel_value_dwpose = uncond_pixel_value_dwpose.to(evaluator.device)

        num_frames = pixel_value_dwpose.size(1)
        audio_prompts = [encode_audio(wav2vec, audio_feat.to(dtype=wav2vec.dtype), fps.item(), num_frames=num_frames)
                         for audio_feat in audio_prompts_raw]
        audio_prompts = torch.cat(audio_prompts, dim=0).to(device=evaluator.device, dtype=weight_dtype)
        uncond_audio_prompts = torch.zeros_like(audio_prompts)

        pixel_value_bboxes_mask = batch["pixel_value_bboxes_mask"]
        pixel_value_bboxes_mask = pixel_value_bboxes_mask / 255.
        pixel_value_bboxes_mask = pixel_value_bboxes_mask.to(evaluator.device)

        pixel_value_object_dot = batch["pixel_value_object_dot"]
        uncond_pixel_value_object_dot = torch.zeros_like(pixel_value_object_dot)
        pixel_value_object_dot = pixel_value_object_dot / 127.5 - 1.
        uncond_pixel_value_object_dot = uncond_pixel_value_object_dot * 2 - 1
        pixel_value_object_dot = pixel_value_object_dot.to(evaluator.device)
        uncond_pixel_value_object_dot = uncond_pixel_value_object_dot.to(evaluator.device)
        if not args.apply_obj_dot:
            pixel_value_object_dot = torch.zeros_like(pixel_value_object_dot)
            pixel_value_object_dot = pixel_value_object_dot * 2 - 1

        pixel_value_ref_for_vae = rearrange(pixel_value_ref, "b f c h w -> b c f h w")
        uncond_uncond_pixel_value_ref = rearrange(uncond_pixel_value_ref, "b f c h w -> b c f h w")
        pixel_value_ref_object_for_vae = rearrange(pixel_value_ref_object, "b f c h w -> b c f h w")
        uncond_pixel_value_ref_object_for_vae = rearrange(uncond_pixel_value_ref_object, "b f c h w -> b c f h w")
        pixel_value_dwpose = rearrange(pixel_value_dwpose, "b f c h w -> b c f h w")
        uncond_pixel_value_dwpose = rearrange(uncond_pixel_value_dwpose, "b f c h w -> b c f h w")
        pixel_value_object_dot = rearrange(pixel_value_object_dot, "b f c h w -> b c f h w")
        uncond_pixel_value_object_dot = rearrange(uncond_pixel_value_object_dot, "b f c h w -> b c f h w")
        pixel_value_bboxes_mask = rearrange(pixel_value_bboxes_mask, "b f c h w -> b c f h w")

        pixel_value_llava = batch["pixel_value_ref_llava"].to(evaluator.device)
        pixel_value_llava = rearrange(pixel_value_llava, "b f c h w -> (b f) c h w")
        uncond_pixel_value_llava = pixel_value_llava.clone()
        if args.zero_uncond_llava:
            uncond_pixel_value_llava = batch["uncond_pixel_value_ref_llava"].to(evaluator.device)
            uncond_pixel_value_llava = rearrange(uncond_pixel_value_llava, "b f c h w -> (b f) c h w")

        pixel_value_object_llava = batch["pixel_value_ref_object_llava"].to(evaluator.device)
        pixel_value_object_llava = rearrange(pixel_value_object_llava, "b f c h w -> (b f) c h w")
        uncond_pixel_value_object_llava = pixel_value_object_llava.clone()
        if args.zero_uncond_llava:
            uncond_pixel_value_object_llava = batch["uncond_pixel_value_ref_object_llava"].to(evaluator.device)
            uncond_pixel_value_object_llava = rearrange(uncond_pixel_value_object_llava, "b f c h w -> (b f) c h w")

        pixel_value_ref_object_ip = batch["pixel_value_ref_object_ip"].to(evaluator.device)
        pixel_value_ref_object_ip = rearrange(pixel_value_ref_object_ip, "b f c h w -> b c f h w")
        uncond_pixel_value_ref_object_ip = torch.zeros_like(pixel_value_ref_object_ip)

        pixel_value_ref_object_ip_2 = batch["pixel_value_ref_object_ip_2"].to(evaluator.device)
        uncond_pixel_value_ref_object_pil_2 = Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))
        uncond_pixel_value_ref_object_ip_2 = evaluator.dino_image_processor(
            uncond_pixel_value_ref_object_pil_2, return_tensors="pt")["pixel_values"].unsqueeze(1).cuda()

        # Encode latents
        ref_latents, uncond_ref_latents = evaluator.encode_image_w_vae3d(
            pixel_value_ref_for_vae, uncond_uncond_pixel_value_ref, args)
        object_latents, uncond_object_latents = evaluator.encode_image_w_vae3d(
            pixel_value_ref_object_for_vae, uncond_pixel_value_ref_object_for_vae, args)
        object_latents_seq, uncond_object_latents_seq = evaluator.encode_image_w_vae3d(
            pixel_value_ref_object_seq, uncond_pixel_value_ref_object_seq, args)
        ip_latents, uncond_ip_latents = evaluator.encode_image_w_vae3d(
            pixel_value_ref_object_ip, uncond_pixel_value_ref_object_ip, args)
        ip_vec, _, uncond_ip_vec, _ = evaluator.encode_image_w_dino(
            pixel_value_ref_object_ip_2, uncond_pixel_value_ref_object_ip_2, args)
        dwpose_latents, uncond_dwpose_latents = evaluator.encode_image_w_vae3d(
            pixel_value_dwpose, uncond_pixel_value_dwpose, args)
        object_motion_latents, uncond_object_motion_latents = evaluator.encode_image_w_vae3d(
            pixel_value_object_dot, uncond_pixel_value_object_dot, args)

        face_masks = torch.nn.functional.interpolate(
            face_masks.float().squeeze(2),
            (ref_latents.shape[-2], ref_latents.shape[-1]),
            mode="nearest").unsqueeze(2).to(dtype=ref_latents.dtype)

        size = (batch["pixel_value_ref"].shape[-2], batch["pixel_value_ref"].shape[-1])
        target_length = sample_n_frames

        def align_to(value, alignment):
            return int(math.ceil(value / alignment) * alignment)

        target_height = align_to(size[0], 16)
        target_width = align_to(size[1], 16)

        concat_dict = {"mode": "timecat-w", "bias": -1}
        freqs_cos, freqs_sin = evaluator.get_rotary_pos_embed(
            129, target_height, target_width, concat_dict)
        n_tokens = freqs_cos.shape[0]

        concat_dict_ip = {"mode": "channelcat", "bias": -2}
        ip_freqs_cos, ip_freqs_sin = evaluator.get_rotary_pos_embed(
            1,
            align_to(pixel_value_ref_object_ip.shape[-2], 16),
            align_to(pixel_value_ref_object_ip.shape[-1], 16),
            concat_dict_ip)

        generator = torch.Generator(device=evaluator.device).manual_seed(seed)

        neg_prompt = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion, blurring"

        pipeline_kwargs = {
            "text_encoder_type": "llava-llama-3-8b-vision",
            "rank": 0,
            "additional_cfg": args.additional_cfg,
            "prompt_cfg": args.prompt_cfg,
            "ref_cfg": args.ref_cfg,
            "use_image_weight_prompt": args.use_image_weight_prompt,
            "iwp_weight": args.iwp_weight,
            "ip_freqs_cis": (ip_freqs_cos, ip_freqs_sin),
            "ip_scale": args.ip_scale,
            "dynamic_ip_scale": args.dynamic_ip_scale,
            "shift_offset": args.shift_offset,
        }

        samples = evaluator.pipeline(
            prompt=prompt_str,
            height=target_height,
            width=target_width,
            frame=target_length,
            num_inference_steps=infer_steps,
            guidance_scale=cfg_scale,
            negative_prompt=neg_prompt,
            num_images_per_prompt=args.num_images,
            generator=generator,
            prompt_embeds=None,
            ref_latents=ref_latents,
            uncond_ref_latents=uncond_ref_latents,
            ref_object_latents=object_latents,
            uncond_ref_object_latents=uncond_object_latents,
            ref_object_latents_seq=object_latents_seq,
            uncond_ref_object_latents_seq=uncond_object_latents_seq,
            ref_object_latents_as_prompt=object_latents_seq,
            ip_latents=ip_latents,
            uncond_ip_latents=uncond_ip_latents,
            ip_vec=ip_vec,
            uncond_ip_vec=uncond_ip_vec,
            pixel_value_llava=pixel_value_llava,
            uncond_pixel_value_llava=uncond_pixel_value_llava,
            pixel_value_object_llava=pixel_value_object_llava,
            uncond_pixel_value_object_llava=uncond_pixel_value_object_llava,
            pixel_value_dwpose=dwpose_latents,
            uncond_pixel_value_dwpose=uncond_dwpose_latents,
            pixel_value_object_dot=object_motion_latents,
            uncond_pixel_value_object_dot=uncond_object_motion_latents,
            object_mask=pixel_value_bboxes_mask,
            face_mask=face_masks,
            audio_prompts=audio_prompts,
            uncond_audio_prompts=uncond_audio_prompts,
            fps=fps,
            ip_cfg_scale=args.ip_cfg_scale,
            attention_mask=None,
            negative_prompt_embeds=None,
            negative_attention_mask=None,
            output_type="pil",
            freqs_cis=(freqs_cos, freqs_sin),
            n_tokens=n_tokens,
            flux_cfg_scale=args.flux_cfg_scale,
            data_type="video" if target_length > 1 else "image",
            is_progress_bar=True,
            vae_ver=args.vae,
            enable_tiling=args.vae_tiling,
            **pipeline_kwargs,
        )[0]

        sample = samples[0].unsqueeze(0)
        video = rearrange(sample[0], "c f h w -> f h w c")
        video = (video * 255.).data.cpu().numpy().astype(np.uint8)

        out_video_path = str(output_dir / f"{videoid}.mp4")
        out_audio_path = str(output_dir / f"{videoid}_audio.mp4")
        imageio.mimsave(out_video_path, video, fps=fps.item())

        os.system(f"ffmpeg -i '{out_video_path}' -i '{audio_path}' -shortest '{out_audio_path}' -y -loglevel quiet")

        del wav2vec
        gc.collect()
        torch.cuda.empty_cache()

        return out_audio_path

    return None


@spaces.GPU(duration=300)
def generate(
    human_image,
    object_image,
    audio,
    prompt,
    infer_steps: int = 50,
    cfg_scale: float = 3.0,
    seed: int = 128,
    pose_control_type: str = "full",
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a human-object interaction video from a reference person image,
    object image, speech audio, and a text prompt."""
    if human_image is None or object_image is None or audio is None:
        return None, "Please provide all inputs: human image, object image, and audio."

    try:
        result = run_inference(
            human_image_path=human_image,
            object_image_path=object_image,
            audio_path=audio,
            prompt=prompt,
            infer_steps=infer_steps,
            cfg_scale=cfg_scale,
            seed=seed,
            pose_control_type=pose_control_type,
        )
        if result and os.path.exists(result):
            return result, "Video generated successfully!"
        return None, "Generation failed - no output produced."
    except Exception as e:
        import traceback
        traceback.print_exc()
        return None, f"Error: {str(e)}"


# ── Gradio UI ──────────────────────────────────────────────────────────────
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    gr.Markdown("""
    # 🎬 HOMA: Human-Object Interaction Animation

    Generate realistic human-object interaction videos from a reference person image,
    an object image, speech audio, and a text prompt.

    Based on the paper [HOMA: Towards Generic Human-Object Interaction in Multimodal Driven Human Animation with Weak Conditions](https://arxiv.org/abs/2506.08797) (SIGGRAPH Asia 2025).
    """)

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            human_img = gr.Image(label="Reference Person Image", type="filepath", scale=1)
            object_img = gr.Image(label="Object Image (PNG with transparency preferred)", type="filepath", scale=1)

        with gr.Row():
            audio_input = gr.Audio(label="Speech Audio (wav)", type="filepath")
            prompt_input = gr.Textbox(
                label="Text Prompt",
                placeholder="Describe the interaction, e.g., 'A woman holds a makeup palette and gestures while speaking.'",
                lines=3,
            )

        run_btn = gr.Button("Generate Video", variant="primary")

        video_output = gr.Video(label="Generated Video")
        status_output = gr.Textbox(label="Status", interactive=False)

        with gr.Accordion("Advanced Settings", open=False):
            infer_steps_slider = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=50, step=1)
            cfg_scale_slider = gr.Slider(label="CFG Scale", minimum=1.0, maximum=10.0, value=3.0, step=0.1)
            seed_input = gr.Number(label="Seed", value=128, precision=0)
            pose_control = gr.Radio(
                label="Pose Control Type",
                choices=["full", "sparse"],
                value="full",
                info="Full uses complete body pose; sparse uses only arm keypoints.",
            )

    gr.Examples(
        examples=[
            [
                str(SPACE_ROOT / "examples" / "assets" / "case_a" / "human.png"),
                str(SPACE_ROOT / "examples" / "assets" / "case_a" / "object.png"),
                str(SPACE_ROOT / "examples" / "assets" / "case_a" / "audio.wav"),
                "A fair-skinned woman with long black hair wears a pink sweater. She holds an open, rectangular makeup palette with orange, red, yellow, green, blue, purple, and white shades. The background features a room with furniture, plants, and various items stacked on top of each other.",
                50, 3.0, 128, "full",
            ],
            [
                str(SPACE_ROOT / "examples" / "assets" / "case_b" / "human.png"),
                str(SPACE_ROOT / "examples" / "assets" / "case_b" / "object.png"),
                str(SPACE_ROOT / "examples" / "assets" / "case_b" / "audio.wav"),
                "An Asian male wearing a black baseball cap, glasses, and a purple sweatshirt sits on a couch holding a white computer keyboard. He is looking at the keyboard and gesturing with his left hand. Behind him are blue walls decorated with framed pictures, stuffed animals, and a green curtain.",
                50, 3.0, 128, "full",
            ],
        ],
        inputs=[human_img, object_img, audio_input, prompt_input, infer_steps_slider, cfg_scale_slider, seed_input, pose_control],
        outputs=[video_output, status_output],
        fn=generate,
        # generate() can return (None, "error…") — Gradio's cached-example machinery
        # crashes on cached None outputs, so run on click instead of caching.
        cache_examples=False,
        run_on_click=True,
    )

    run_btn.click(
        fn=generate,
        inputs=[human_img, object_img, audio_input, prompt_input, infer_steps_slider, cfg_scale_slider, seed_input, pose_control],
        outputs=[video_output, status_output],
    )

demo.launch(mcp_server=True, show_error=True)