backend.py

import os

# Must be set before torch/OpenMP libraries load
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

# Hugging Face ZeroGPU: `spaces` has to be imported before torch. Outside a ZeroGPU Space it isn't
# installed, so @spaces.GPU becomes a no-op and the model simply runs on the device chosen below.
try:
    import spaces
except ImportError:
    class spaces:
        @staticmethod
        def GPU(fn=None, duration=None):
            return fn if fn else (lambda f: f)

import sys
import json
import time
import torch
import numpy as np
from dotenv import load_dotenv

if hasattr(sys.stdout, 'reconfigure'):
    try:
        sys.stdout.reconfigure(encoding='utf-8', errors='replace')
    except Exception:
        pass

# Import the translation model and helper functions from Uni_Sign
from Uni_Sign.models import Uni_Sign
from Uni_Sign.datasets import load_part_kp_YTASL, YTASL_GROUP_SIZES, _fill_missing_landmarks, select_frame_indices
from predict_pose import create_mediapipe_models, predict_pose, load_video_cv
from visualize_pose import render_keypoints_video
from utils import make_emitter

load_dotenv()


def select_device():
    """
    Device and precision for the translation model, from environment variables:
      SLT_DEVICE = auto (default: GPU if available) | cuda | cpu
      SLT_DTYPE  = auto (default: BF16 on GPU as in training, FP32 on CPU) | bf16 | float32
    """
    device_choice = os.environ.get("SLT_DEVICE", "auto").lower()
    dtype_choice = os.environ.get("SLT_DTYPE", "auto").lower()
    if device_choice not in ("auto", "cuda", "cpu"):
        raise ValueError(f"SLT_DEVICE must be auto, cuda or cpu, not {device_choice!r}")
    if dtype_choice not in ("auto", "bf16", "float32"):
        raise ValueError(f"SLT_DTYPE must be auto, bf16 or float32, not {dtype_choice!r}")
    if device_choice == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("SLT_DEVICE=cuda, but no GPU is available. Use SLT_DEVICE=cpu or auto.")

    use_cuda = device_choice == "cuda" or (device_choice == "auto" and torch.cuda.is_available())
    device = torch.device("cuda" if use_cuda else "cpu")
    use_bf16 = dtype_choice == "bf16" or (dtype_choice == "auto" and device.type == "cuda")
    return device, torch.bfloat16 if use_bf16 else torch.float32


DEVICE, DTYPE = select_device()
# YOLO runs outside the @spaces.GPU call, and on ZeroGPU CUDA is only usable inside it
POSE_DEVICE = "cpu" if os.environ.get("SPACES_ZERO_GPU") else DEVICE.type


# ---------- Diagnostics (set SLT_DEBUG=0 to silence) ----------
DEBUG = os.environ.get("SLT_DEBUG", "1") != "0"


def diag(message):
    if DEBUG:
        print(f"[diag] {message}", flush=True)


def tensor_summary(t):
    """Shape, dtype, device and value range of a tensor (NaN/Inf flagged)."""
    if not isinstance(t, torch.Tensor):
        return repr(t)[:60]
    f = t.detach().float()
    finite = f[torch.isfinite(f)]
    rng = f"min={finite.min().item():.4g} max={finite.max().item():.4g} mean={finite.mean().item():.4g}" if finite.numel() else "no finite values"
    bad = f" NaN={torch.isnan(f).sum().item()} Inf={torch.isinf(f).sum().item()}" if finite.numel() != f.numel() else ""
    return f"{tuple(t.shape)} {t.dtype} {t.device} {rng}{bad}"


def weights_fingerprint(m):
    """A few numbers that change if weights are lost, zeroed or garbled when moved between devices."""
    params = dict(m.named_parameters())
    names = [next(iter(params)), "mt5_model.shared.weight", "mt5_model.lm_head.weight", list(params)[-1]]
    return {n: round(params[n].detach().float().abs().sum().item(), 3) for n in names if n in params}


def weights_location(m):
    """How many parameters sit on each device/dtype."""
    counts = {}
    for p in m.parameters():
        key = f"{p.device}/{p.dtype}"
        counts[key] = counts.get(key, 0) + p.numel()
    return {k: f"{v / 1e6:.1f}M" for k, v in counts.items()}


def environment_report():
    import transformers
    try:
        import spaces as _spaces_module  # noqa: F401
        has_spaces = not isinstance(spaces, type)
    except ImportError:
        has_spaces = False
    cuda = torch.cuda.is_available()
    diag(f"python {sys.version.split()[0]} | torch {torch.__version__} (cuda build {torch.version.cuda}) | transformers {transformers.__version__}")
    diag(f"spaces package: {has_spaces} | SPACES_ZERO_GPU={os.environ.get('SPACES_ZERO_GPU')!r} | "
         f"SLT_DEVICE={os.environ.get('SLT_DEVICE', 'auto')!r} SLT_DTYPE={os.environ.get('SLT_DTYPE', 'auto')!r}")
    diag(f"torch.cuda.is_available()={cuda}" + (f" | {torch.cuda.device_count()} GPU(s): {torch.cuda.get_device_name(0)} | bf16 supported: {torch.cuda.is_bf16_supported()}" if cuda else ""))
    diag(f"selected DEVICE={DEVICE} DTYPE={DTYPE} POSE_DEVICE={POSE_DEVICE}")

# Global variables to cache models and config so they are not reloaded on every video translation
model = None
pose_models = None
args = None

class InferenceConfig:
    """Configuration class for Uni_Sign model inference."""
    def __init__(self):
        # Path to the pre-trained weights checkpoint
        self.finetune = os.environ.get("UNISIGN_WEIGHTS", r"./Uni_Sign/unisign_model/best_checkpoint-wlasl.pth")
        self.dataset = "YTASL"
        self.task = "ISLR"
        self.max_length = 256
        self.normalization = "none"
        self.layout = "pruned"
        self.n_registers = 0
        self.hidden_dim = 256
        self.rgb_support = False
        self.no_adaptive_gcn = False
        self.register_position = "before_all"
        self.label_smoothing = 0.2
        self.batch_size = 1

def initialize_model():
    """Initializes the Uni_Sign translator and the Mediapipe/YOLO pose estimation models."""
    global model, pose_models, args

    if model is not None and pose_models is not None:
        return

    print("Initializing Uni_Sign model...")
    environment_report()
    args = InferenceConfig()
    model = Uni_Sign(args=args)

    # Load model weights
    if args.finetune and os.path.exists(args.finetune):
        print(f"Loading checkpoint weights from: {args.finetune}")
        diag(f"checkpoint size: {os.path.getsize(args.finetune) / 1e6:.1f} MB")
        state_dict = torch.load(args.finetune, map_location='cpu')['model']
        result = model.load_state_dict(state_dict, strict=False)
        diag(f"checkpoint tensors: {len(state_dict)} | model tensors: {len(model.state_dict())} | "
             f"missing: {len(result.missing_keys)} {result.missing_keys[:5]} | unexpected: {len(result.unexpected_keys)} {result.unexpected_keys[:5]}")
    else:
        print(f"WARNING: Checkpoint weights not found at '{args.finetune}'!")

    diag(f"weights fingerprint after loading (cpu, float32): {weights_fingerprint(model)}")
    model.to(device=DEVICE, dtype=DTYPE).eval()
    print(f"Uni_Sign model on {DEVICE} ({DTYPE}); pose detection on {POSE_DEVICE}")
    diag(f"weights location after .to(): {weights_location(model)}")
    diag(f"weights fingerprint after .to(): {weights_fingerprint(model)}")
    diag(f"mT5 generation config: decoder_start={model.mt5_model.config.decoder_start_token_id} "
         f"eos={model.mt5_model.config.eos_token_id} pad={model.mt5_model.config.pad_token_id}")

    print("Initializing Mediapipe and YOLO models...")
    pose_checkpoint_folder = 'checkpoints/pose/'
    pose_models = create_mediapipe_models(pose_checkpoint_folder)


EXAMPLES_JSON = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples.json")

def load_examples():
    """
    Loads examples from examples.json ({name: {title, video_path, ...}}) as a list in file order,
    resolving their paths relative to the repo root.
    """
    if not os.path.exists(EXAMPLES_JSON):
        return []
    base_dir = os.path.dirname(EXAMPLES_JSON)
    with open(EXAMPLES_JSON, "r", encoding="utf-8") as f:
        data = json.load(f)
    examples = []
    for name, ex in data.items():
        ex = {"id": name, "title": ex.get("title", name), **ex}
        examples.append(ex)
        for key in ("video_path", "keypoints_path", "keypoints_video_path"):
            if ex.get(key):
                ex[key] = os.path.normpath(os.path.join(base_dir, ex[key]))
    return examples

def process_pose_data_in_memory(pose_results, args):
    """
    Processes pose keypoints in-memory from dictionary results.
    Converts keypoint data to Tensors formatted for Uni_Sign input.
    """
    raw_pose = pose_results.get('cropped_keypoints', [])
    if not raw_pose:
        raise ValueError("Video does not contain any detected pose (cropped_keypoints).")

    # Format raw coordinates into standard [X, Y] coordinates for each body part
    pose = []
    for frame_data in raw_pose:
        formatted_frame = {}
        for part, expected_size in YTASL_GROUP_SIZES.items():
            kps = frame_data.get(part, [])
            
            # Use empty list if data is missing; _fill_missing_landmarks will fill it
            if kps is None or len(kps) == 0:
                 formatted_frame[part] = [] 
            else:
                 # Standardize to only X and Y coordinates (ignoring confidence/Z columns)
                 formatted_frame[part] = np.array(kps)[:, :2].tolist()
        pose.append(formatted_frame)

    detected = {part: sum(1 for f in pose if len(f[part]) > 0) for part in YTASL_GROUP_SIZES}
    diag(f"keypoints: {len(pose)} frames; frames with detections per part: "
         + ", ".join(f"{part}={n}/{len(pose)}" for part, n in detected.items()))

    # Subsample frames to match target length (max_length)
    duration = len(pose)
    tmp = select_frame_indices(duration, args.max_length, phase='test')
    skeletons = [pose[i] for i in tmp]

    # Fill in missing landmarks to maintain consistent shape dimensions
    confs = []
    for i, skeleton in enumerate(skeletons):
        conf = {}
        for group_name, expected_size in YTASL_GROUP_SIZES.items():
            _fill_missing_landmarks(
                skeleton=skeleton,
                conf=conf,
                group_name=group_name,
                expected_size=expected_size,
                clip_name="gradio_video",
                frame_idx=i,
                error_group_label=f"group '{group_name}'",
                include_size_details=False,
                strict_key_access=True,
            )
        confs.append(conf)

    # Normalize keypoints and structure them for YTASL dataset
    kps_with_scores = load_part_kp_YTASL(skeletons, confs, args.normalization, args.layout)

    # Batch features by unsqueezing to add batch dimension (batch_size = 1)
    src_input = {}
    for key, val in kps_with_scores.items():
        src_input[key] = val.unsqueeze(0)

    # Generate attention mask and sequence lengths
    seq_len = src_input['body'].shape[1]
    mask_gen = torch.ones([seq_len]) + 7
    src_input['attention_mask'] = (mask_gen != 0).long().unsqueeze(0)
    src_input['src_length_batch'] = torch.LongTensor([seq_len])
    src_input['name_batch'] = ["gradio_video"]

    for key, value in src_input.items():
        if isinstance(value, torch.Tensor):
            diag(f"model input {key}: {tensor_summary(value)}")
    return src_input


@spaces.GPU(duration=60)  # on ZeroGPU, only this call holds a GPU (time limit in seconds)
def run_translation(src_input):
    """Runs Uni_Sign on prepared pose tensors and returns the decoded translation."""
    # Inside the GPU call: is CUDA really there, and did the weights arrive intact?
    cuda = torch.cuda.is_available()
    diag(f"in run_translation: torch.cuda.is_available()={cuda}"
         + (f", current device {torch.cuda.current_device()} ({torch.cuda.get_device_name(0)})" if cuda else ""))
    diag(f"weights location in run_translation: {weights_location(model)}")
    diag(f"weights fingerprint in run_translation: {weights_fingerprint(model)}")

    for key, value in src_input.items():
        if isinstance(value, torch.Tensor):
            src_input[key] = value.to(device=DEVICE, dtype=DTYPE if value.is_floating_point() else value.dtype)
    tgt_input = {'gt_sentence': [""], 'gt_gloss': [""]}
    tokenizer = model.mt5_tokenizer

    with torch.inference_mode():
        stack_out = model(src_input, tgt_input)
        diag(f"encoder input embeds: {tensor_summary(stack_out['inputs_embeds'])}")
        diag(f"encoder attention mask: {tensor_summary(stack_out['attention_mask'])}")

        # What the model wants to say first: top-5 candidates for the first output token
        start = torch.full((1, 1), model.mt5_model.config.decoder_start_token_id, device=DEVICE, dtype=torch.long)
        logits = model.mt5_model(inputs_embeds=stack_out['inputs_embeds'], attention_mask=stack_out['attention_mask'],
                                 decoder_input_ids=start).logits[0, -1].float()
        diag(f"first-token logits: {tensor_summary(logits)}")
        probs, ids = logits.softmax(-1).topk(5)
        diag("first-token top-5: " + ", ".join(f"{tok!r}={p:.3f}" for tok, p in
                                               zip(tokenizer.convert_ids_to_tokens(ids.tolist()), probs.tolist())))

        greedy = model.generate(stack_out, max_new_tokens=100, num_beams=1)
        diag(f"greedy: ids {greedy[0].tolist()[:20]} -> {tokenizer.batch_decode(greedy, skip_special_tokens=True)[0]!r}")
        output = model.generate(stack_out, max_new_tokens=100, num_beams=4)
        diag(f"beam search (4): ids {output[0].tolist()[:20]} -> {tokenizer.batch_decode(output, skip_special_tokens=True)[0]!r}")

    return tokenizer.batch_decode(output, skip_special_tokens=True)[0].strip()


def process_input(input_video_path, progress=None, log_callback=None):
    """
    Main entry point for Gradio. Runs pose extraction, pre-processing, 
    and model translation inference via Uni_Sign.
    """
    try:
        start_time = time.time()
        emit = make_emitter(start_time, progress, log_callback)

        emit("0. Initializing models...", 0.02, "0. Initializing models...")
        initialize_model()
        
        # Check if precomputed keypoints JSON & pre-rendered video exist (e.g. from examples.json or file pattern)
        video_dir = os.path.dirname(os.path.abspath(input_video_path))
        video_stem = os.path.splitext(os.path.basename(input_video_path))[0]
        
        # Normalize stem in case it ends with _keypoints
        clean_stem = video_stem.replace("_keypoints", "")
        possible_json = os.path.join(video_dir, f"{clean_stem}_keypoints.json")
        possible_video = os.path.join(video_dir, f"{clean_stem}_keypoints.mp4")

        # Also check examples.json registry
        input_key = os.path.normcase(os.path.abspath(input_video_path))
        for item in load_examples():
            if os.path.normcase(item.get("video_path", "")) == input_key:
                if item.get("keypoints_path") and os.path.exists(item["keypoints_path"]):
                    possible_json = item["keypoints_path"]
                if item.get("keypoints_video_path") and os.path.exists(item["keypoints_video_path"]):
                    possible_video = item["keypoints_video_path"]
                break

        if os.path.exists(possible_json) and os.path.getsize(possible_json) > 0:
            emit("1. Loading precomputed keypoints from JSON...", 0.50, "1. Loading precomputed keypoints...")
            with open(possible_json, "r", encoding="utf-8") as f:
                pose_results = json.load(f)
            
            keypoints_video_path = possible_video if os.path.exists(possible_video) else ""
            emit("  [1e] Precomputed keypoints loaded from cache", 0.90, "  [1e] Precomputed keypoints loaded")
        else:
            emit("1. Extracting keypoints from video...", 0.05, "1. Extracting keypoints from video...")
            video_frames, fps = load_video_cv(input_video_path)
            if fps is None or fps <= 0 or np.isnan(fps):
                fps = 25.0

            num_frames = len(video_frames)
            dur = num_frames / fps if fps > 0 else 0
            emit(f"  [1a] Video loaded: {num_frames} frames ({fps:.1f} FPS, duration {dur:.1f} s)", 0.06, f"  [1a] Video loaded ({num_frames} frames)")

            pose_results = predict_pose(video_frames, pose_models, progress=progress, start_time=start_time, log_callback=log_callback, yolo_device=POSE_DEVICE)

            # Attach bounding boxes to each frame's keypoints
            kps_list = pose_results.get("keypoints", [])
            bbox_face_list = pose_results.get("bbox_face", [])
            bbox_lh_list = pose_results.get("bbox_left_hand", [])
            bbox_rh_list = pose_results.get("bbox_right_hand", [])

            for i, kp in enumerate(kps_list):
                if i < len(bbox_face_list):
                    kp['bbox_face'] = bbox_face_list[i]
                if i < len(bbox_lh_list):
                    kp['bbox_left_hand'] = bbox_lh_list[i]
                if i < len(bbox_rh_list):
                    kp['bbox_right_hand'] = bbox_rh_list[i]

            # Render keypoint overlay video
            keypoints_video_path = ""
            try:
                keypoints_video_path = render_keypoints_video(video_frames, kps_list, fps=fps, progress=progress, start_time=start_time, log_callback=log_callback)
                emit("  [1e] Keypoint video saved successfully", 0.90, "  [1e] Keypoint video saved")
            except Exception as kp_err:
                emit(f"  [1e] Warning rendering keypoint video: {kp_err}", 0.90, "  [1e] Warning keypoint video")
        
        emit("2. Pre-processing visual features...", 0.92, "2. Pre-processing visual features...")
        emit("3. Converting keypoints to Tensors for Uni-Sign...", 0.94, "3. Converting keypoints to Tensors...")
        src_input = process_pose_data_in_memory(pose_results, args)

        emit("4. Generating translation with Uni-Sign...", 0.96, "4. Generating translation...")
        result = run_translation(src_input)
        if not result:
            result = "The model was unable to generate a translation. Please try a different video."
        
        emit("5. Translation completed!", 1.0, "5. Translation completed!")
        return result, keypoints_video_path
                
    except Exception as e:
        # Raised rather than returned as text, so a translation can never be mistaken for an error
        raise RuntimeError(f"Could not process the video: {e}") from e