inference.py

import os
import sys
import pdb
import json
import yaml
import random
import argparse

import torch
import numpy as np
import torch.nn.functional as F
from tqdm import tqdm
from scipy.signal import savgol_filter
from torchvision.utils import save_image
import torchvision.transforms.functional as TF

from xtalker import flow_model, myutils
from xtalker.custom_dataloader_liveportrait import *


def segmental_generate_images(live_model, source_images, pred_motion_emb, max_seg=16):
    num_frames = source_images.shape[0]
    if num_frames <= max_seg:
        return live_model.gen_image(source_images, pred_motion_emb[0])

    seg_outputs = [
        live_model.gen_image(source_images[start:end], pred_motion_emb[0, start:end])
        for start in range(0, num_frames, max_seg)
        for end in [min(start + max_seg, num_frames)]
    ]
    return torch.cat(seg_outputs, dim=0)

def build_curve_pose(curve_pose, n_motion_frames=200, curve_key=None):
    curve_keys = [key for key in curve_pose.keys() if key != "dist_info"]
    random_curve = curve_key if curve_key else random.choice(curve_keys)
    curve_data = curve_pose[random_curve]
    coords = torch.tensor(curve_data["coords"], dtype=torch.float32)
    # poses = torch.tensor(curve_data[random_curve.replace("curve", "pose")], dtype=torch.float32)
    poses = torch.tensor(curve_data["pose"], dtype=torch.float32)
    coords = F.interpolate(coords.T.unsqueeze(0), size=n_motion_frames, mode='linear', align_corners=False).transpose(1, 2)
    poses = F.interpolate(poses.T.unsqueeze(0), size=n_motion_frames, mode='linear', align_corners=False).transpose(1, 2)
    return coords, poses, random_curve

def assign_overlap_motion_emb(pred_motion_emb, pred_motion_emb_pre, overlap_win_size):
    if pred_motion_emb_pre is not None:
        c = torch.linspace(0, 1, overlap_win_size, device=pred_motion_emb.device, dtype=pred_motion_emb.dtype).view(1, -1, 1)
        pred_motion_emb[:, 0:overlap_win_size, :] = (
            c * pred_motion_emb[:, 0:overlap_win_size, :]
            + (1 - c) * pred_motion_emb_pre[:, -overlap_win_size:, :]
        )
    return pred_motion_emb


def motion_dim_report(name, motion_emb, dim, expr_offset=0, num_values=20):
    dim_idx = expr_offset + dim
    if motion_emb is None or dim_idx >= motion_emb.shape[-1]:
        return f"{name}=NA"

    curve = motion_emb[0, :, dim_idx].float()
    curve = torch.nan_to_num(curve, nan=0.0, posinf=0.0, neginf=0.0)
    values_n = min(num_values, curve.numel())
    first_values = ",".join(f"{v:.4f}" for v in curve[:values_n].detach().cpu().tolist())
    return (
        f"{name}: first20=[{first_values}], mean={float(curve.mean()):.4f}, "
        f"min={float(curve.min()):.4f}, max={float(curve.max()):.4f}"
    )


def postprocess_mouth_ema(pred_motion_emb, envelope, source_ref_emb, mouth_kps,
                          silence_threshold=0.03, silence_softness=0.08,
                          ema_alpha=0.65, silence_lerp=True,
                          use_source_ref=True, target_max_open=0.0,
                          target_mean_open=0.0,
                          lower_lip=58, upper_lip=61):
    """Mouth-only EMA: hold in quiet regions and softly release with envelope."""
    if pred_motion_emb is None or not mouth_kps:
        return pred_motion_emb

    valid_kps = [int(k) for k in mouth_kps if 0 <= int(k) < pred_motion_emb.shape[-1]]
    if not valid_kps:
        return pred_motion_emb

    alpha = max(0.0, min(1.0, float(ema_alpha)))
    if alpha >= 1.0:
        return pred_motion_emb

    out = pred_motion_emb.clone()
    mouth_idx = torch.tensor(valid_kps, device=out.device, dtype=torch.long)
    x = out[..., mouth_idx].float()
    y = x.clone()

    if envelope is None:
        speech = torch.ones(x.shape[:2] + (1,), device=out.device, dtype=torch.float32)
    else:
        env = envelope.to(device=out.device, dtype=torch.float32)
        if env.dim() == 2:
            env = env.unsqueeze(-1)
        if env.shape[1] != x.shape[1]:
            env = F.interpolate(
                env.transpose(1, 2),
                size=x.shape[1],
                mode="linear",
                align_corners=False,
            ).transpose(1, 2)
        softness = max(float(silence_softness), 1e-6)
        speech = ((env - float(silence_threshold)) / softness).clamp(0.0, 1.0)

    if use_source_ref and source_ref_emb is not None:
        source_mouth = source_ref_emb[..., mouth_idx + 7].to(
            device=out.device, dtype=torch.float32
        )
        if source_mouth.dim() == 2:
            source_mouth = source_mouth.unsqueeze(1)
    else:
        source_mouth = x[:, 0:1, :]
    y[:, 0, :] = speech[:, 0, :] * x[:, 0, :] + (1.0 - speech[:, 0, :]) * source_mouth[:, 0, :]

    for frame_idx in range(1, x.shape[1]):
        alpha_t = alpha * speech[:, frame_idx, :]
        y[:, frame_idx, :] = alpha_t * x[:, frame_idx, :] + (1.0 - alpha_t) * y[:, frame_idx - 1, :]

    if silence_lerp and envelope is not None:
        quiet = speech.squeeze(-1) <= 1e-6
        for batch_idx in range(quiet.shape[0]):
            frame_idx = 0
            while frame_idx < quiet.shape[1]:
                if not bool(quiet[batch_idx, frame_idx]):
                    frame_idx += 1
                    continue

                start = frame_idx
                while frame_idx + 1 < quiet.shape[1] and bool(quiet[batch_idx, frame_idx + 1]):
                    frame_idx += 1
                end = frame_idx

                left = (
                    source_mouth[batch_idx, 0]
                    if start == 0
                    else y[batch_idx, start - 1].clone()
                )
                if end + 1 < quiet.shape[1]:
                    right = y[batch_idx, end + 1].clone()
                    n_frames = end - start + 1
                    weights = torch.linspace(
                        1.0 / (n_frames + 1),
                        n_frames / (n_frames + 1),
                        n_frames,
                        device=out.device,
                        dtype=torch.float32,
                    ).view(-1, 1)
                    y[batch_idx, start:end + 1] = (
                        (1.0 - weights) * left.unsqueeze(0)
                        + weights * right.unsqueeze(0)
                    )
                else:
                    y[batch_idx, start:end + 1] = left

                frame_idx += 1

    target_max_open = float(target_max_open)
    if target_max_open > 0.0 and lower_lip in valid_kps and upper_lip in valid_kps:
        lower_pos = valid_kps.index(lower_lip)
        upper_pos = valid_kps.index(upper_lip)
        open_diff = y[..., lower_pos] - y[..., upper_pos]
        current_max = open_diff.max()

        if torch.isfinite(current_max) and float(current_max) < target_max_open:
            if float(current_max) > 1e-6:
                scale = target_max_open / (current_max + 1e-6)
                target_open = torch.where(open_diff > 0, open_diff * scale, open_diff)
            else:
                speech_weight = speech.squeeze(-1)
                weight_max = speech_weight.max()
                if torch.isfinite(weight_max) and float(weight_max) > 1e-6:
                    target_open = open_diff + target_max_open * (speech_weight / (weight_max + 1e-6))
                else:
                    target_open = open_diff

            delta_open = target_open - open_diff
            y[..., lower_pos] = y[..., lower_pos] + delta_open * (2.0 / 3.0)
            y[..., upper_pos] = y[..., upper_pos] - delta_open * (1.0 / 3.0)

        target_mean_open = float(target_mean_open)
        if target_mean_open > -0.1:
            open_diff = y[..., lower_pos] - y[..., upper_pos]
            current_mean = open_diff.mean()
            if torch.isfinite(current_mean) and float(current_mean) < target_mean_open:
                delta_mean = target_mean_open - current_mean
                y[..., lower_pos] = y[..., lower_pos] + delta_mean * (2.0 / 3.0)
                y[..., upper_pos] = y[..., upper_pos] - delta_mean * (1.0 / 3.0)

    out[..., mouth_idx] = y.to(dtype=out.dtype)
    return out


def correct_reference_mouth_open(source_ref_emb, lower_lip=58, upper_lip=61,
                                 expr_offset=7, target_open=0.0):
    """
    Only correct the inference reference embedding when the first-frame mouth is over-closed.

    open = lower_lip_y - upper_lip_y = dim58 - dim61.
    If open < 0, move only the reference embedding to open=0.
    The correction is split as lower:upper = 2:1:
      - dim58 (lower lip y) moves down by 2/3 delta, so dim58 increases.
      - dim61 (upper lip y) moves up by 1/3 delta, so dim61 decreases.

    This does not modify source image, source video motion, predicted a2v motion, or v2v.
    """
    lower_idx = expr_offset + lower_lip
    upper_idx = expr_offset + upper_lip
    if lower_idx >= source_ref_emb.shape[-1] or upper_idx >= source_ref_emb.shape[-1]:
        return source_ref_emb, None

    corrected = source_ref_emb.clone()
    open_before = corrected[..., lower_idx] - corrected[..., upper_idx]
    delta = torch.clamp(float(target_open) - open_before, min=0.0)
    if float(delta.max()) <= 0.0:
        return corrected, None

    lower_delta = delta * (2.0 / 3.0)
    upper_delta = delta * (1.0 / 3.0)
    corrected[..., lower_idx] = corrected[..., lower_idx] + lower_delta
    corrected[..., upper_idx] = corrected[..., upper_idx] - upper_delta
    return corrected, None


def mouth_open_sequence_report(name, motion_emb, lower_lip=58, upper_lip=61, expr_offset=7):
    lower_idx = expr_offset + lower_lip
    upper_idx = expr_offset + upper_lip
    if motion_emb is None or lower_idx >= motion_emb.shape[-1] or upper_idx >= motion_emb.shape[-1]:
        return None

    open_diff = (motion_emb[..., lower_idx] - motion_emb[..., upper_idx]).float()
    return (
        f"[MouthOpenSeq] file={name}, "
        f"max={float(open_diff.max()):.4f}, "
        f"min={float(open_diff.min()):.4f}, "
        f"mean={float(open_diff.mean()):.4f}"
    )


def init_seg_interp(T=200, D=66, seg_len=25, first_frame_latent=None, batch_size=1, device=None, dtype=None):
    # every 25 frames (1s) interpolate once to get smooth noise latent
    S = (T + seg_len - 1) // seg_len  # 段数
    keys = torch.randn(batch_size, S+1, 1, D, device=device, dtype=dtype)  # 关键点
    if first_frame_latent is not None:
        keys[:, 0:1] = first_frame_latent

    parts = []
    for s in range(S):
        n = min(seg_len, T - s*seg_len)
        t = torch.linspace(0, 1, n, device=device, dtype=dtype).view(-1, 1)
        seg = (1 - t) * keys[:, s] + t * keys[:, s+1]
        parts.append(seg)
    return torch.cat(parts, 1)  # [1,T,D]

def init_smoothed_noise(n_motion_frames=200, latent_dim=66, ignore_indices=None, device=None, dtype=None, batch_size=1,
                        mouth_kps=None, mouth_kps_latent=None, envelope=None, first_frame_latent=None,
                        source_init_latent=None, source_init_kps_latent=None,
                        noise_type="seg_interp"):
    # Use caller-supplied latent indices when available (remapped valid-dim space);
    # fall back to legacy hardcoded indices for backward compatibility.
    if mouth_kps_latent is not None:
        mouth_kps = torch.tensor(mouth_kps_latent, device=device, dtype=torch.long)
    elif mouth_kps is not None:
        mouth_kps = torch.tensor(mouth_kps, device=device, dtype=torch.long) + 2
    else:
        mouth_kps = torch.tensor([58, 61, 43, 53], device=device, dtype=torch.long) + 2  # legacy: original 65-dim latent
    # mouth_kps = torch.tensor([58, 61], device=device) + 2 # add 2 pose index, mouth_kps 53 O
    # mouth_std = torch.tensor([0.01248989279474321, 0.006130238289615321], device=device, dtype=dtype)
    
    # pdb.set_trace()
    if noise_type=="random":
        noise_latents = torch.randn([batch_size, n_motion_frames, latent_dim], device=device, dtype=dtype)
    elif noise_type=="seg_interp":
        noise_latents = init_seg_interp(n_motion_frames, latent_dim, 5, first_frame_latent, batch_size=batch_size, device=device, dtype=dtype)

    # Mouth channels should not inherit the fixed-interval linear interpolation
    # prior; use per-frame random noise so audio/viseme cues drive the rhythm.
    valid_mouth_kps = mouth_kps[(mouth_kps >= 0) & (mouth_kps < latent_dim)]
    if valid_mouth_kps.numel() > 0:
        noise_latents[:, :, valid_mouth_kps] = torch.randn(
            batch_size,
            n_motion_frames,
            valid_mouth_kps.numel(),
            device=device,
            dtype=dtype,
        )
    # Keep mouth latents unbiased; the model should close/open from audio features,
    # not from an envelope-driven open-mouth noise prior.
    weights_jitter = torch.full((latent_dim,), 0.1, device=device, dtype=dtype) # 0.3 will cause jitter
    no_jitter_idxs = torch.cat([
        torch.tensor([0, 1], device=device, dtype=torch.long),
        valid_mouth_kps,
    ])
    weights_jitter[no_jitter_idxs] = 0 # no jitter on head pose and mouth kps
    jitter = torch.randn_like(noise_latents) * weights_jitter
    noise_latents = noise_latents + jitter  # [1, T, D]
    if ignore_indices is not None:
        noise_latents[:, :, ignore_indices+2] = 0.0 # add 2 pose index

    if source_init_latent is not None and source_init_kps_latent is not None:
        source_init_kps = torch.tensor(source_init_kps_latent, device=device, dtype=torch.long)
        source_init_kps = source_init_kps[
            (source_init_kps >= 0) & (source_init_kps < latent_dim)
        ]
        if source_init_kps.numel() > 0:
            source_values = source_init_latent.to(device=device, dtype=dtype)
            if source_values.dim() == 2:
                source_values = source_values.unsqueeze(1)
            if source_values.shape[-1] == latent_dim:
                init_values = source_values[..., source_init_kps]
            elif source_values.shape[-1] == latent_dim - 2:
                expr_kps = source_init_kps - 2
                valid_expr = (expr_kps >= 0) & (expr_kps < source_values.shape[-1])
                source_init_kps = source_init_kps[valid_expr]
                expr_kps = expr_kps[valid_expr]
                init_values = source_values[..., expr_kps]
            else:
                raise ValueError(
                    "source_init_latent last dim must match latent_dim or latent_dim - 2"
                )
            if source_init_kps.numel() > 0:
                noise_latents[:, :, source_init_kps] = init_values.expand(
                    batch_size, n_motion_frames, -1
                )

    end_latent = noise_latents[:, -1:, :].clone()
    # noise_latents = (noise_latents - noise_latents.mean(dim=-2, keepdim=True)) / noise_latents.std(dim=-2, keepdim=True)
    return noise_latents, end_latent
    
class XTalkerPerformer:
    """
    XTalkerPerformer handles the inference process for XTalker models.
    * It performs inference on validation cases, processes audio and video data,
    * and generates videos based on the model's predictions.
    TODO: It also supports single image inference with audio input.
    """

    def __init__(self, cfg, wav2vec, live_portrait_model, model=None):
        if (
            "valid_expr_dims" not in cfg
            or "motion_std" not in cfg
            or "noise_mouth_kps_latent" not in cfg
        ):
            cfg = myutils.configure_motion_filter(cfg)
        self.cfg = cfg
        self.output_dir = cfg["output_dir"]
        self.device = cfg["device"]
        self.dtype = cfg["weight_dtype"]
        
        self.model = model
        self.wav2vec = wav2vec
        self.live_model = live_portrait_model
        self.relative_motion = myutils.is_relative_motion(self.cfg)
        self.live_warp_stitch_renderer = None
        if self.cfg.get("live_warp_stitch", False):
            from liveportrait.joyvasa_decoder import LiveWarpStitchRenderer
            self.live_warp_stitch_renderer = LiveWarpStitchRenderer(
                device=str(self.device),
                checkpoint_dir=self.cfg.get("liveportrait_model_path"),
                retargeting_dir=self.cfg.get("liveportrait_retargeting_model_path"),
                models_config=self.cfg.get("liveportrait_models_config", "config/liveportrait_models.yaml"),
            )

        self.n_motion_frames = cfg["n_motion_frames"]
        self.latent_dim = cfg["latent_dim"]
        self.overlap_win_size = cfg["overlap_win_size"]
        self.train_steps = cfg["train_steps"]
        self.inference_steps = cfg["inference_steps"]
        
        _std_full             = torch.tensor(self.cfg["motion_std"], dtype=self.dtype, device=self.device)
        self.valid_expr_dims  = self.cfg["valid_expr_dims"]   # list[int] in original 63-dim space
        _v                    = torch.tensor(self.valid_expr_dims, dtype=torch.long, device=self.device)
        self.motion_std_valid = _std_full[_v]    # (n_valid,) for normalizing valid dims
        self.motion_std_full  = _std_full        # (63,)  for scattering back at reconstruction
        
        self.emo_delta = load_emotion_delta(self.cfg["emo_delta_path"])
        self.expr2idx_onehot = {key: torch.nn.functional.one_hot(torch.tensor(idx), \
                num_classes=len(self.cfg["expr2idx"])).float() \
                for key, idx in self.cfg["expr2idx"].items()}
        self._random_inference_emo_idx = 0  # round-robin cursor for random_inference_emo
        self.curve_pose = load_curve_pose(self.cfg["curve_pose_path"])
        self.curve_mean = torch.tensor(self.curve_pose["dist_info"]["mean"], dtype=self.dtype, device=self.device)
        self.curve_std = torch.tensor(self.curve_pose["dist_info"]["std"], dtype=self.dtype, device=self.device)
        self.use_viseme_prior = self.cfg.get("use_viseme_prior", True)
        self.viseme_label_cache = {}
        if self.use_viseme_prior:
            self.viseme15_template_matrix = load_viseme15_template_matrix(
                self.cfg.get("viseme15_template_path", "data_process/priors/viseme15_mouth_templates.json"),
                valid_expr_dims=self.valid_expr_dims,
                motion_std=self.cfg.get("motion_std"),
                normalize=self.cfg.get("normalize_viseme_prior", True),
                scale=self.cfg.get("viseme_prior_scale", 1.0),
                clip_value=self.cfg.get("viseme_prior_clip", 10.0),
            )
        else:
            self.viseme15_template_matrix = np.zeros((15, len(self.valid_expr_dims)), dtype=np.float32)
        self.viseme15_eval_template_matrix = load_viseme15_template_matrix(
            self.cfg.get("viseme15_template_path", "data_process/priors/viseme15_mouth_templates.json"),
            valid_expr_dims=None,
            motion_std=None,
            normalize=False,
            scale=1.0,
            clip_value=None,
        )
        

    def perform_cases(self, model, output_dir=None):
        """
        Use validation cases to perform inference for training or batch inferences.
        """
        self.model = model
        self.output_dir = output_dir
        
        # load demo video yaml config
        with open(self.cfg["demos_videos_path"], "r") as fid:
            demo_videos_info = yaml.load(fid, Loader=yaml.Loader)
        video_paths = demo_videos_info["demo_mp4_paths"]
        audio_paths = demo_videos_info["demo_wav_paths"]
        debug_num_cases = self.cfg.get("debug_num_cases")
        if debug_num_cases is not None:
            debug_num_cases = int(debug_num_cases)
            video_paths = video_paths[:debug_num_cases]
            audio_paths = audio_paths[:debug_num_cases]
        if len(video_paths) == 0:
            raise (f'empty video_path_paths in rank={self.cfg["machine_rank"]}')

        for video_path, audio_path in tqdm(zip(video_paths, audio_paths)):
            self.perform_once(video_path, None, audio_path)
            # if "final" in self.output_dir: break

    def perform_once(self, video_path=None, img_path=None, audio_path=None, ):
        """
        Perform inference on a single demo video and audio file.
        This function prepares the video and audio data, runs the inference,
        and saves the output video along with denoise metrics.
        
        """
        if (video_path is None and img_path is None) or audio_path is None:
            raise ValueError("(video_path or img_path) and audio_path must be provided.")
        
        (video_fps, source_images, source_motion_emb, \
            audio_embed, envelope_embed) = self.preprocess_video_wav(
            video_path, audio_path
        )
        self.video_length = len(source_images)
        self.video_fps = video_fps
        
        # v2v to test the generation ability of the liveportrait model
        v2v_path_export = os.path.join(self.output_dir, \
            os.path.basename(video_path).replace(".", "_v2v."))
        v2v_frames = self.perform_live_v2v(source_images, source_motion_emb, audio_path, v2v_path_export)

        # a2v to test the generation ability of the DiT model
        a2v_path_export = os.path.join(self.output_dir, \
            os.path.basename(video_path).replace(".", "_a2v."))
        metrics_path_export = os.path.join(self.output_dir, \
            os.path.basename(video_path).replace(".mp4", "_metric.json"))
        a2v_frames = self.perform_xtalker(
            source_images, source_motion_emb,
            audio_embed, envelope_embed,
            video_path=video_path,
            audio_path=audio_path,
            a2v_path_export=a2v_path_export,
            metrics_path_export=metrics_path_export
        )

        # Save side-by-side comparison: v2v (left) | a2v (right)
        side_by_side_path = os.path.join(self.output_dir,
            os.path.basename(video_path).replace(".", "_compare."))
        min_len = min(len(v2v_frames), len(a2v_frames))
        combined = torch.cat(
            [v2v_frames[:min_len].cpu().float(), a2v_frames[:min_len].cpu().float()],
            dim=2  # width dimension: (T,H,W,3) → (T,H,2W,3)
        )
        myutils.tensor_to_video_inbatch(
            combined, side_by_side_path,
            audio_source=audio_path,
            fps=self.video_fps, save_fps=self.video_fps,
            draw_frame_numbers_enabled=self.cfg.get("draw_frame_numbers", False),
        )

    def preprocess_video_wav(self, video_path, audio_path):
        video_frame_tensor, video_fps = myutils.extract_video_frame_to_tensor(
            video_path, resolution=256)  # (f,3,512,512), value [0,1]
        
        source_images = video_frame_tensor.to(device=self.device, dtype=self.dtype)
        source_motion_emb = self.live_model.get_motion_emb(source_images)  # (f,d)
        source_motion_emb = source_motion_emb.unsqueeze(0)  # (1,f,d)
        
        audio_embed, envelope_embed = compute_audio_features(self.wav2vec, \
            audio_path, is_envelope=True, fps=video_fps)

        max_idx = min(len(video_frame_tensor), len(audio_embed))
        return (video_fps, source_images[:max_idx], source_motion_emb[:, :max_idx], audio_embed[:max_idx], envelope_embed[:max_idx])

    def load_viseme_data_file(self, label_path):
        if not label_path or not os.path.exists(label_path):
            return {}
        if label_path not in self.viseme_label_cache:
            self.viseme_label_cache[label_path] = load_viseme_label_data(label_path)
        return self.viseme_label_cache[label_path]

    def find_viseme15_ids(self, video_path, audio_path, video_length, warn=True):
        stems = []
        for path in (audio_path, video_path):
            if path:
                stems.append(os.path.splitext(os.path.basename(path))[0])

        candidates = []
        for path in (audio_path, video_path):
            if path:
                path_dir = os.path.dirname(path)
                candidates.append(os.path.join(path_dir, "viseme_labels.json"))
                parent_dir = os.path.dirname(path_dir)
                dir_name = os.path.basename(path_dir)
                if dir_name.endswith("_data"):
                    candidates.append(os.path.join(
                        parent_dir,
                        dir_name[:-len("_data")] + "_phoneme",
                        "viseme_labels.json",
                    ))
        candidates.append(self.cfg.get("viseme_label_path"))

        info = None
        found_label_path = None
        found_stem = None
        for label_path in dict.fromkeys(p for p in candidates if p):
            label_data = self.load_viseme_data_file(label_path)
            for stem in stems:
                if stem in label_data:
                    info = label_data[stem]
                    found_label_path = label_path
                    found_stem = stem
                    break
            if info is not None:
                break

        if info is None:
            if warn:
                print(
                    f"[Warn] no viseme prior found for stems={stems}; "
                    f"audio_path={audio_path}; video_path={video_path}; "
                    f"checked={list(dict.fromkeys(p for p in candidates if p))}"
                )
            return np.zeros((video_length,), dtype=np.int64)

        ids = parse_int_list_from_str(info.get("viseme15_ids", ""))
        return fit_1d_array_length(ids, video_length, pad_value=0)

    def get_viseme15_prior(self, video_path, audio_path, video_length):
        n_dim = self.viseme15_template_matrix.shape[1]
        if not self.use_viseme_prior:
            return torch.zeros((1, video_length, n_dim), dtype=self.dtype, device=self.device)

        ids = self.find_viseme15_ids(video_path, audio_path, video_length)
        prior_curve = build_viseme_prior_curve(ids, self.viseme15_template_matrix)

        return torch.tensor(prior_curve, dtype=self.dtype, device=self.device).unsqueeze(0)

    def perform_live_v2v(
        self, source_images, source_motion_emb,
        audio_path, v2v_path_export,
        ):
        pred_video_frame_list = []
        infer_batch_size = 1  # cfg["batch_size"]
        for frame_id in range(0, self.video_length, infer_batch_size):
            i = frame_id
            j = min(self.video_length, i + infer_batch_size)

            source_images_repeat = source_images[0:1].repeat([j - i, 1, 1, 1])
            drive_motion_emb = source_motion_emb[:, i:j, :][0]  # (b,f,d)
            # drive_motion_emb[:,0:7] = drive_motion_emb[0:1, 0:7] # try
            model_output_decoded = self.live_model.gen_image(
                source_images_repeat, drive_motion_emb)
            pred_video_frame_list.append(model_output_decoded.cpu().detach().float())

        pred_video_frame_tensor = torch.concat(pred_video_frame_list, dim=0)
        pred_video_frame_tensor = pred_video_frame_tensor.clamp(0, 1).permute(0, 2, 3, 1).contiguous()  # (B,H,W,3)

        myutils.tensor_to_video_inbatch(
            pred_video_frame_tensor,
            v2v_path_export,
            audio_source=audio_path,
            fps=self.video_fps,
            save_fps=self.video_fps,
            draw_frame_numbers_enabled=self.cfg.get("draw_frame_numbers", False),
        )
        return pred_video_frame_tensor  # (T,H,W,3) float, [0,1]

    def prepare_missing_frames(self, source_motion_emb, audio_embed, envelope_embed, viseme_prior=None):
        # pdb.set_trace()
        
        mis_frames = 0
        total_frames = self.video_length + (self.video_length - 1) // (self.n_motion_frames - self.overlap_win_size) * self.overlap_win_size
        if total_frames % self.n_motion_frames != 0:
            mis_frames = self.n_motion_frames - (total_frames % self.n_motion_frames)

        # Extend source_motion_emb
        last_motion_emb = source_motion_emb[:, -1:, :]
        source_motion_emb = torch.cat(
            [source_motion_emb, last_motion_emb.repeat(1, mis_frames, 1)], dim=1
        )

        # Extend audio_embed
        last_audio_embed = audio_embed[-1:, :]
        audio_embed = torch.cat(
            [audio_embed, last_audio_embed.repeat(mis_frames, 1)], dim=0
        )

        # Extend envelope_embed
        last_envelope_embed = envelope_embed[-1:]
        envelope_embed = torch.cat(
            [envelope_embed, last_envelope_embed.repeat(mis_frames, 1)], dim=0
        )

        if viseme_prior is not None:
            last_viseme_prior = viseme_prior[:, -1:, :]
            viseme_prior = torch.cat(
                [viseme_prior, last_viseme_prior.repeat(1, mis_frames, 1)], dim=1
            )

        return source_motion_emb, audio_embed.unsqueeze(0), envelope_embed.unsqueeze(0), viseme_prior, mis_frames
        
    def perform_xtalker(
        self,
        source_images, source_motion_emb,
        audio_embed, envelope_embed,
        emo_label=None, 
        video_path=None, audio_path=None, a2v_path_export=None, metrics_path_export=None
        ):
        """
        Perform inference using the XTalker model on the provided video frames and audio.
        """
        if "final" in self.output_dir and "0001_a2v" in a2v_path_export: 
            np.save(os.path.join(self.output_dir, "source.npy"), source_motion_emb.float().cpu().numpy())
        source_ref_img = source_images[0:1]
        source_ref_emb = source_motion_emb[:, 0:1, :]  # (b, 1, d)
        # print(
        #     f"\n[MouthO] file={os.path.basename(a2v_path_export)}, "
        #     f"{motion_dim_report('v2v_o53', source_motion_emb, dim=53, expr_offset=7)}"
        # )
        source_ref_emb, _ = correct_reference_mouth_open(
            source_ref_emb,
            target_open=self.cfg.get("ref_mouth_target_open", 0.005),
        )
        pose_ref_emb = source_ref_emb[:, :, 0:2].repeat(1, self.n_motion_frames, 1)  # (b, n, 4)
        deform_ref_emb = source_ref_emb[:, :, 2:7].repeat(1, self.n_motion_frames, 1)  # (b, n, 4)

        if emo_label is not None:
            emo_label_resolved = emo_label
            drive_emo_onehot = self.expr2idx_onehot[emo_label]
        elif self.cfg.get("random_inference_emo", False):
            emo_labels = list(self.cfg["expr2idx"])
            emo_label_resolved = emo_labels[self._random_inference_emo_idx % len(emo_labels)]
            self._random_inference_emo_idx += 1
            drive_emo_onehot = self.expr2idx_onehot[emo_label_resolved].unsqueeze(0).unsqueeze(0)
        else:
            emo_label_resolved = self.cfg["inference_expr_label"]
            drive_emo_onehot = self.expr2idx_onehot[emo_label_resolved].unsqueeze(0).unsqueeze(0)
        drive_curve, drive_pose, random_curve = build_curve_pose(curve_pose=self.curve_pose, n_motion_frames=self.n_motion_frames)
        drive_pose = drive_pose[..., :2] # only use pitch yaw
        
        viseme15_ids = self.find_viseme15_ids(video_path, audio_path, self.video_length, warn=False)
        viseme_prior = self.get_viseme15_prior(video_path, audio_path, self.video_length)
        source_motion_emb_, audio_embed_, envelope_embed_, viseme_prior_, mis_frames = \
            self.prepare_missing_frames(source_motion_emb, audio_embed, envelope_embed, viseme_prior)
        if mis_frames > 0:
            pad_value = int(viseme15_ids[-1]) if len(viseme15_ids) > 0 else 0
            viseme15_ids = np.concatenate([
                viseme15_ids,
                np.full((mis_frames,), pad_value, dtype=np.int64),
            ])
        mouth_o_viseme_ids = [int(v) for v in self.cfg.get("mouth_o_viseme_ids", [13])]
        mouth_o_soft_mask = None
        if mouth_o_viseme_ids:
            mouth_o_transition_frames = int(self.cfg.get("mouth_o_transition_frames", 3))
            mouth_o_hard = torch.tensor(
                np.isin(
                    viseme15_ids,
                    np.asarray(mouth_o_viseme_ids, dtype=np.int64),
                ),
                dtype=torch.bool,
            ).view(1, -1)
            mouth_o_soft_mask = myutils.soften_binary_mask(
                mouth_o_hard,
                transition_frames=mouth_o_transition_frames,
            ).squeeze(0).cpu().numpy()
        
        pred_motion_emb_combined_pre = None
        pred_source_emo_emb_cached = None
        pred_video_frame_list = []
        pred_motion_emb_list = []
        denoising_loss_metric_list = []
        
        if "final" in self.output_dir and "0001_a2v" in a2v_path_export: debug_pred_motion_emb = []
        for loop_count, frame_start in enumerate(
            range(0, self.video_length, self.n_motion_frames - self.overlap_win_size)
        ):
            frame_end = min(frame_start + self.n_motion_frames, self.video_length + mis_frames)
            # pdb.set_trace()
            
            drive_motion_emb = source_motion_emb_[:, frame_start:frame_end, :]  # (1,f,d)
            drive_audio_emb = audio_embed_[:, frame_start:frame_end, :]
            drive_env_emb = envelope_embed_[:, frame_start:frame_end, :]
            drive_viseme_prior = (
                viseme_prior_[:, frame_start:frame_end, :]
                if viseme_prior_ is not None
                else None
            )

            # denoise 步骤
            pred_motion_emb, pred_source_emo_emb, pred_head_pose, \
            denoising_loss_metric_chunk, last_frame_latent = self.flow_denoising_loop(
                source_ref_emb=source_ref_emb, drive_emo_onehot=drive_emo_onehot,
                drive_audio_emb=drive_audio_emb, drive_env_emb=drive_env_emb,
                drive_viseme_prior=drive_viseme_prior,
                drive_curve=drive_curve, frame_start=frame_start, 
                drive_motion_emb=drive_motion_emb[..., 7:], drive_pose=drive_pose,
                first_frame_latent=None if loop_count==0 else last_frame_latent
            )
            denoising_loss_metric_list.append(denoising_loss_metric_chunk)

            ################### post-process motion emb to anti jitter ###################
            # pred_motion_emb_combined = drive_motion_emb #try
            # pred_motion_emb = myutils.lowpass_temporal(pred_motion_emb, kernel_size=5, sigma=1.5)
            # pred_motion_emb = myutils.lowpass_temporal(pred_motion_emb, kernel_size=7, sigma=2)
            # invalid kps now come from keypoint.yaml and are source-filled.
            # invalid_idx = [6*3, 6*3+1, 8*3, 8*3+1, 9*3, 9*3+1, 12*3, 12*3+1, 18*3, 18*3+1]
            # pred_motion_emb[..., invalid_idx] = 0*pred_motion_emb[..., invalid_idx]
            # K2-XY and K5-XY scaling changes source-filled head/neck axes; keep raw values.
            # scale_idx = [2*3, 2*3+1, 5*3, 5*3+1]
            # pred_motion_emb[..., scale_idx] = pred_motion_emb[..., scale_idx]/3
            # Bias-align stable dims to source/previous chunk. For the ablation,
            # keep mouth dims as absolute predictions unless explicitly enabled.
            mouth_bias_exclude = []
            if not self.cfg.get("align_mouth_to_source", False):
                mouth_bias_exclude = self.cfg.get("mouth_kps_orig", self.cfg.get("mouth_kps", []))
            bias_exclude = set([34, 40, 46, 49])
            bias_exclude.update(int(i) for i in mouth_bias_exclude)
            indices_exc = torch.tensor(
                [i for i in range(pred_motion_emb.shape[-1]) if i not in bias_exclude],
                device=pred_motion_emb.device,
            )
            if pred_motion_emb_combined_pre is not None:
                pred_motion_emb[..., indices_exc] = pred_motion_emb[..., indices_exc]-(pred_motion_emb[:, 0:1, indices_exc]-\
                    pred_motion_emb_combined_pre[:, -self.overlap_win_size:-self.overlap_win_size+1, indices_exc+7])
            else:
                pred_motion_emb[..., indices_exc] = pred_motion_emb[..., indices_exc]-(pred_motion_emb[:, 0:1, indices_exc]-\
                    source_ref_emb[..., indices_exc+7])
                
                # 58 (lower lip), 61 (upper lip), 53 (O/stretch) mouth kps indices in motion emb
                # pred_motion_emb[..., 61] = pred_motion_emb[..., 61]-(pred_motion_emb[:, 0:1, 61]-0)

            mouth_o_bias = float(self.cfg.get("mouth_o_bias", 0.0))
            if mouth_o_bias != 0.0 and pred_motion_emb.shape[-1] > 53 and mouth_o_soft_mask is not None:
                o_mask_np = mouth_o_soft_mask[frame_start:frame_end][:pred_motion_emb.shape[1]]
                if float(o_mask_np.max()) > 0.0:
                    o_mask = torch.tensor(
                        o_mask_np,
                        device=pred_motion_emb.device,
                        dtype=pred_motion_emb.dtype,
                    ).view(1, -1)
                    pred_motion_emb[:, :o_mask.shape[1], 53] = (
                        pred_motion_emb[:, :o_mask.shape[1], 53] - mouth_o_bias * o_mask
                    )

            if self.cfg.get("use_mouth_postprocess", True):
                pred_motion_emb = postprocess_mouth_ema(
                    pred_motion_emb,
                    drive_env_emb,
                    source_ref_emb,
                    self.cfg.get("mouth_kps_orig", self.cfg.get("mouth_kps", [])),
                    silence_threshold=self.cfg.get("silence_env_threshold", 0.03),
                    silence_softness=self.cfg.get("mouth_post_silence_softness", 0.08),
                    ema_alpha=self.cfg.get("mouth_post_ema_alpha", 0.65),
                    silence_lerp=self.cfg.get("mouth_post_silence_lerp", True),
                    use_source_ref=self.cfg.get("align_mouth_to_source", False),
                    target_max_open=self.cfg.get("mouth_post_target_max_open", 0.04),
                    target_mean_open=self.cfg.get("mouth_post_target_mean_open", 0.005),
                )
            
            # post_process pred_head_pose
            # pred_head_pose = myutils.lowpass_temporal(pred_head_pose, kernel_size=7, sigma=2)
            # pred_head_pose = pred_head_pose/4  # reduce head pose amplitude
            # pdb.set_trace()
            pred_head_pose = pred_head_pose - (pred_head_pose[:, 0:1, :]-pose_ref_emb[:, 0:1, :])
            ################### post-process motion emb ###################
            
            pred_motion_emb_combined = torch.cat(
                [pred_head_pose, deform_ref_emb, pred_motion_emb], dim=-1)
                # [pose_ref_emb, deform_ref_emb, pred_motion_emb], dim=-1)
            # if "final" not in self.output_dir: 

            if pred_motion_emb_combined_pre is not None:
                pred_motion_emb_combined = assign_overlap_motion_emb(pred_motion_emb_combined, 
                                    pred_motion_emb_combined_pre, self.overlap_win_size)
            pred_motion_emb_combined_pre = pred_motion_emb_combined
            
            if "final" in self.output_dir and "0001_a2v" in a2v_path_export: debug_pred_motion_emb.append(pred_motion_emb_combined)
            # drive_motion_emb = drive_motion_emb * self.motion_std + self.motion_mean  # (1,f,d)

            if pred_source_emo_emb_cached is None:
                pred_source_emo_emb_cached = pred_source_emo_emb
                pred_emo_img_cached = self.live_model.gen_image(source_ref_img, pred_source_emo_emb_cached.squeeze(0))
                emo_img_to_save = pred_emo_img_cached
                if self.cfg.get("random_inference_emo", False):
                    emo_img_to_save = myutils.draw_text_label(pred_emo_img_cached, emo_label_resolved)
                save_image(emo_img_to_save.detach().cpu() , a2v_path_export.replace(".mp4", "_emo_ref.png"), normalize=True, value_range=(0, 1))
                # Key step to avoid size mismatch
                pred_emo_img_cached = TF.resize(pred_emo_img_cached, (256, 256)) 
                # if "final" not in self.output_dir: 
                pred_emo_img_cached = source_ref_img # for debug without emotion effect
                
            # generate reference images for the following warping
            # pdb.set_trace()
            source_emo_imgs = pred_emo_img_cached.repeat([self.n_motion_frames, 1, 1, 1])
            
            # generate target images with liveportrait warping model
            if self.live_warp_stitch_renderer is not None:
                pred_target_images = self.live_warp_stitch_renderer.render(
                    source_emo_imgs[0], pred_motion_emb_combined
                )
            else:
                pred_target_images = segmental_generate_images(
                    self.live_model, source_emo_imgs, pred_motion_emb_combined, max_seg=16
                )
            
            # process the overlap frames of previous chunk and missing frames of last chunk 
            if loop_count == len(range(0, self.video_length, self.n_motion_frames - self.overlap_win_size)) - 1:
                if mis_frames > 0:
                    pred_video_frame_list.append(pred_target_images[:-mis_frames])
                    pred_motion_emb_list.append(pred_motion_emb_combined[:, :-mis_frames, :])
                else:
                    pred_video_frame_list.append(pred_target_images)
                    pred_motion_emb_list.append(pred_motion_emb_combined)
            else:
                pred_video_frame_list.append(pred_target_images[:-self.overlap_win_size])
                pred_motion_emb_list.append(pred_motion_emb_combined[:, :-self.overlap_win_size, :])

        if "final" in self.output_dir and "0001_a2v" in a2v_path_export: 
            debug_pred_motion_emb = torch.cat(debug_pred_motion_emb, dim=1)
            np.save(os.path.join(self.output_dir, "pred.npy"), debug_pred_motion_emb.float().cpu().numpy())
        pred_motion_emb_tensor = torch.cat(pred_motion_emb_list, dim=1)
        mouth_open_report = mouth_open_sequence_report(
            os.path.basename(a2v_path_export) if a2v_path_export else "unknown",
            pred_motion_emb_tensor,
        )
        if mouth_open_report is not None:
            print(mouth_open_report)
        viseme_report = myutils.viseme_motion_eval_report(
            os.path.basename(a2v_path_export) if a2v_path_export else "unknown",
            pred_motion_emb_tensor,
            self.find_viseme15_ids(video_path, audio_path, pred_motion_emb_tensor.shape[1], warn=False),
            self.viseme15_eval_template_matrix,
            mouth_dims=self.cfg.get("mouth_kps_orig", self.cfg.get("mouth_kps", [])),
        )
        if viseme_report is not None:
            print(viseme_report)
        # print(
        #     f"[MouthO] file={os.path.basename(a2v_path_export)}, "
        #     f"{motion_dim_report('a2v_o53', pred_motion_emb_tensor, dim=53, expr_offset=7)}"
        # )
        pred_video_frame_tensor = torch.cat(pred_video_frame_list, dim=0).clamp_(0, 1).permute(0, 2, 3, 1)

        pred_video_frame_tensor = pred_video_frame_tensor.cpu().detach().float()
        myutils.tensor_to_video_inbatch(
            pred_video_frame_tensor,
            a2v_path_export,
            audio_source=audio_path,
            fps=self.video_fps,
            save_fps=self.video_fps,
            draw_frame_numbers_enabled=self.cfg.get("draw_frame_numbers", False),
        )

        with open(metrics_path_export, "w") as f:
            json.dump(denoising_loss_metric_list, f, indent=2)

        return pred_video_frame_tensor  # (T,H,W,3) float, [0,1]
            
                
    def flow_denoising_loop(self,
        source_ref_emb=None, drive_emo_onehot=None,
        drive_audio_emb=None, drive_env_emb=None,
        drive_viseme_prior=None,
        drive_curve=None, frame_start=None,
        drive_motion_emb=None, drive_pose=None,
        first_frame_latent=None,
        ):
        """
        Perform the denoising loop for the given model and inputs with flow matching.
        """
        
        delta_t = 1.0 / self.inference_steps
        # Slice source to valid expression dims and normalize.  In relative mode,
        # zero is the source expression and the model predicts source-relative deltas.
        source_expr_norm = source_ref_emb[..., 7:][:, :, self.valid_expr_dims]
        source_expr_norm = source_expr_norm / self.motion_std_valid
        if self.relative_motion:
            emotion_latents = torch.zeros_like(source_expr_norm)
            motion_init_latent = torch.zeros_like(source_expr_norm)
        else:
            emotion_latents = source_expr_norm
            motion_init_latent = source_expr_norm

        noise_latents, last_frame_latent = init_smoothed_noise(self.n_motion_frames, self.latent_dim,
                                device=self.device, dtype=self.dtype,
                                mouth_kps_latent=self.cfg.get("noise_mouth_kps_latent"),
                                source_init_latent=motion_init_latent,
                                source_init_kps_latent=self.cfg.get("weak_kps_latent"),
                                envelope=drive_env_emb,
                                first_frame_latent=first_frame_latent)
        initial_noise_latents = noise_latents.clone()
        # source_ref_emb_in = source_ref_emb[..., 7:] / self.motion_std_ig
        
        drive_pose = drive_pose.to(device=self.device)
        
        denoising_loss_metric_chunk = []
        for tau in range(self.inference_steps):
            # convert testing time step to training time steps
            timesteps = torch.full((noise_latents.shape[0],),
                int(tau / self.inference_steps * self.train_steps),
                device=self.device, dtype=torch.long)

            v_motion_emb, v_source_emo, v_head_pose = self.model(
                latent=noise_latents.to(device=self.device, dtype=self.dtype),
                emotion_latents=emotion_latents.to(device=self.device, dtype=self.dtype),
                emo_onehot=drive_emo_onehot.to(device=self.device, dtype=self.dtype),
                curve=drive_curve.to(device=self.device, dtype=self.dtype),
                audio=drive_audio_emb.to(device=self.device, dtype=self.dtype),
                envelope=drive_env_emb.to(device=self.device, dtype=self.dtype),
                viseme_prior=(
                    drive_viseme_prior.to(device=self.device, dtype=self.dtype)
                    if drive_viseme_prior is not None
                    else None
                ),
                timestep=timesteps,
                # control_emo=source_ref_emb_in
            )
            v_motion_comb = torch.cat([v_head_pose, v_motion_emb], dim=-1)  # (B,F,2+n_valid)

            noise_latents = noise_latents + v_motion_comb * delta_t
            emotion_latents = emotion_latents + v_source_emo * delta_t
 
            # drive_motion_emb is still full 63-dim; filter + normalize for metric only
            drive_motion_norm = drive_motion_emb[:, :, self.valid_expr_dims].to(torch.float32) / self.motion_std_valid
            if self.relative_motion:
                drive_motion_target = drive_motion_norm - source_expr_norm.to(torch.float32)
            else:
                drive_motion_target = drive_motion_norm
            target_latents = torch.cat([drive_pose, drive_motion_target], dim=-1).to(torch.float32)

            t_float = (tau + 1) / self.inference_steps
            initial_latents = initial_noise_latents.to(torch.float32)
            current_latents = noise_latents.to(torch.float32)
            velocity_target = target_latents - initial_latents
            expected_noise_latents = (1 - t_float) * initial_latents + t_float * target_latents
            
            denoising_loss_metric_chunk.append(
            {
                "frame_id": frame_start,
                "infer_step": tau,
                "main_loss": torch.nn.functional.mse_loss(v_motion_comb.float(), velocity_target).item(),
                "xt_loss": torch.nn.functional.mse_loss(current_latents, expected_noise_latents).item(),
                "x1_loss": torch.nn.functional.mse_loss(current_latents, target_latents).item(),
            }
        )
        
        pred_head_pose = (noise_latents[:, :, :2] * self.curve_std[:2]).float()  # (B,F,2)

        # Denormalize valid expression dims; cast to source dtype (LivePortrait uses float32).
        # Relative mode predicts deltas, which are added back to the source expression.
        pred_motion_valid = (noise_latents[:, :, 2:] * self.motion_std_valid).float()  # (B,F,n_valid)

        # Reconstruct full 63-dim: invalid dims filled with source reference values
        B, F_len = pred_motion_valid.shape[:2]
        source_expr = source_ref_emb[..., 7:].float().expand(B, F_len, -1).clone()    # (B,F,63)
        pred_motion_emb = source_expr.clone()
        if self.relative_motion:
            pred_motion_emb[:, :, self.valid_expr_dims] = (
                source_expr[:, :, self.valid_expr_dims] + pred_motion_valid
            )
        else:
            pred_motion_emb[:, :, self.valid_expr_dims] = pred_motion_valid

        # Emotion: denorm valid dims, scatter into full 63-dim with source filling invalid
        emo_valid  = (emotion_latents * self.motion_std_valid).float()                 # (B,1,n_valid)
        source_emo = source_ref_emb[..., 7:].float().expand(B, 1, -1).clone()         # (B,1,63)
        if self.relative_motion:
            source_emo[:, :, self.valid_expr_dims] = (
                source_emo[:, :, self.valid_expr_dims] + emo_valid
            )
        else:
            source_emo[:, :, self.valid_expr_dims] = emo_valid
        # Reconstruct full 70-dim: [pose_0:7 | expr_7:70]
        pred_source_emo_emb = torch.cat(
            (source_ref_emb[..., :7].float().expand(B, 1, -1), source_emo), dim=-1    # (B,1,70)
        )

        return pred_motion_emb, pred_source_emo_emb, pred_head_pose, denoising_loss_metric_chunk, last_frame_latent


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Simple example of training script.")
    parser.add_argument(
        "--config",
        type=str,
        default="./config/default_config.yaml",
    )
    args, _ = parser.parse_known_args()

    # load cfg and update it w.r.t argv
    with open(args.config) as fid:
        cfg = yaml.load(fid, Loader=yaml.Loader)
    cfg = myutils.update_cfg_from_argv(cfg=cfg, argv=sys.argv)