Blame for app/engine.py

import math
import re
import threading
import time
from collections import OrderedDict
from dataclasses import dataclass
from typing import Any, Optional

import numpy as np

from app.config import Settings

FRAME_DIMENSION = 133
FRAMES_PER_TOKEN = 4

MOTION_LAYOUT = {
    "upper_body": [0, 30],
    "left_hand": [30, 75],
    "right_hand": [75, 120],
    "jaw": [120, 123],
    "expression": [123, 133],
}

SENTENCE_BOUNDARY_PATTERN = re.compile(r"(?<=[.!?])\s+")
WORD_PATTERN = re.compile(r"\w+", re.UNICODE)


def count_words(text: str) -> int:
    return len(WORD_PATTERN.findall(text))


def split_sentences(text: str, maximum_words: int) -> list:
    cleaned_text = " ".join(text.split())
    if cleaned_text == "":
        return []

    rough_sentences = SENTENCE_BOUNDARY_PATTERN.split(cleaned_text)
    sentences = []

    for rough_sentence in rough_sentences:
        if WORD_PATTERN.search(rough_sentence) is None:
            continue

        words = rough_sentence.split(" ")

        if len(words) <= maximum_words:
            sentences.append(rough_sentence)
            continue

        for start_index in range(0, len(words), maximum_words):
            window = words[start_index:start_index + maximum_words]
            sentences.append(" ".join(window))

    return sentences


@dataclass(frozen=True)
class GenerationRequest:
    text: str
    token_count: int
    use_retrieval: bool
    do_sample: bool
    temperature: float
    top_k: int
    seed: int


@dataclass
class GenerationResult:
    motion: np.ndarray
    fps: float
    token_count: int
    elapsed_seconds: float
    from_cache: bool
    keywords: list


class SignMotionEngine:
    def __init__(self, settings: Settings) -> None:
        self.settings = settings
        self.ready = False
        self.device_name = "not loaded"
        self.load_seconds = 0.0
        self.torch: Any = None
        self.device: Any = None
        self.generator: Any = None
        self.vocab: Any = None
        self.tokenizer: Any = None
        self.mean: Optional[np.ndarray] = None
        self.std: Optional[np.ndarray] = None
        self.dictionary: dict = {}
        self.extract_keywords: Any = None
        self.denormalize_motion: Any = None
        self.maximum_motion_tokens = 0
        self.generation_lock = threading.Lock()
        self.cache_lock = threading.Lock()
        self.cache: OrderedDict = OrderedDict()
        self.cache_hits = 0
        self.cache_misses = 0

    def load(self) -> None:
        load_started_at = time.perf_counter()

        import torch
        from robostar.checkpoints import read_json, resolve_model_root
        from robostar.generator import RoboSTARGenerator
        from robostar.representation import denormalize_motion
        from robostar.retrieval import extract_keywords
        from robostar.tokenizer import RoboSTARTokenizer

        self.torch = torch
        self.extract_keywords = extract_keywords
        self.denormalize_motion = denormalize_motion

        requested_device = self.settings.device
        if requested_device.startswith("cuda") and not torch.cuda.is_available():
            selected_device = "cpu"
        else:
            selected_device = requested_device

        self.device = torch.device(selected_device)
        self.device_name = str(self.device)

        model_root = resolve_model_root(
            self.settings.model_id,
            allow_patterns=[
                "generator/*",
                "tokenizer/*",
                "stats/*",
                "retrieval/*",
                "README.md",
                "manifest.json",
            ],
        )

        generator, vocab = RoboSTARGenerator.from_pretrained(
            model_root,
            base_model=self.settings.base_model_id,
            device=self.device,
        )
        generator.eval()

        tokenizer = RoboSTARTokenizer.from_pretrained(model_root, device=self.device).eval()

        self.mean = np.load(model_root / "stats" / "mean133.npy").astype(np.float32)
        self.std = np.load(model_root / "stats" / "std133.npy").astype(np.float32)

        dictionary_path = model_root / "retrieval" / "word2code.json"
        if dictionary_path.is_file():
            self.dictionary = read_json(dictionary_path)
        else:
            self.dictionary = {}

        self.generator = generator
        self.vocab = vocab
        self.tokenizer = tokenizer
        self.maximum_motion_tokens = int(generator.cfg.max_motion_tokens)

        warm_up_request = GenerationRequest(
            text="Warm up sentence.",
            token_count=10,
            use_retrieval=False,
            do_sample=False,
            temperature=self.settings.temperature,
            top_k=self.settings.top_k,
            seed=self.settings.seed,
        )
        self.generate_uncached(warm_up_request)

        self.load_seconds = time.perf_counter() - load_started_at
        self.ready = True

    def plan_token_count(self, sentence: str, pace: float) -> int:
        word_count = count_words(sentence)
        planned_seconds = word_count * self.settings.seconds_per_word * pace
        planned_seconds = max(planned_seconds, self.settings.minimum_seconds)
        planned_seconds = min(planned_seconds, self.settings.maximum_seconds)

        token_count = int(math.ceil(planned_seconds * self.settings.fps / FRAMES_PER_TOKEN))
        token_count = max(1, token_count)

        if self.maximum_motion_tokens > 0:
            token_count = min(token_count, self.maximum_motion_tokens)

        return token_count

    def generate(self, request: GenerationRequest) -> GenerationResult:
        if not self.ready:
            raise RuntimeError("The model is still loading")

        with self.cache_lock:
            cached_result = self.cache.get(request)
            if cached_result is not None:
                self.cache.move_to_end(request)
                self.cache_hits += 1
                return GenerationResult(
                    motion=cached_result.motion,
                    fps=cached_result.fps,
                    token_count=cached_result.token_count,
                    elapsed_seconds=0.0,
                    from_cache=True,
                    keywords=cached_result.keywords,
                )
            self.cache_misses += 1

        with self.generation_lock:
            result = self.generate_uncached(request)

        with self.cache_lock:
            self.cache[request] = result
            self.cache.move_to_end(request)
            while len(self.cache) > self.settings.cache_size:
                self.cache.popitem(last=False)

        return result

    def generate_uncached(self, request: GenerationRequest) -> GenerationResult:
        torch = self.torch
        started_at = time.perf_counter()

        torch.manual_seed(request.seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(request.seed)

        keywords: list = []
        prompt = request.text

        if request.use_retrieval and self.dictionary:
            keywords = self.extract_keywords(request.text, self.dictionary, self.settings.max_keywords)
            prompt = self.vocab.retrieval_prompt(request.text, keywords, self.dictionary)

        encoded = self.vocab.encode_text([prompt], self.generator.cfg.max_source_length)
        input_ids = encoded["input_ids"].to(self.device)
        attention = encoded["attention_mask"].to(self.device)

        explicit_length = torch.tensor([request.token_count], dtype=torch.long, device=self.device)

        with torch.inference_mode():
            generated = self.generator.generate(
                input_ids,
                attention,
                explicit_length,
                temperature=request.temperature,
                top_k=request.top_k,
                do_sample=request.do_sample,
            )

            token_length = int(generated["full_lengths"][0].item())

            codes = {}
            for part_name in ("body", "left", "right"):
                codes[part_name] = generated["final_codes"][part_name][0:1, :token_length]

            normalized_motion = self.tokenizer.decode(codes)[0].float().cpu().numpy().astype(np.float32)

        motion = self.denormalize_motion(normalized_motion, self.mean, self.std)
        motion = np.ascontiguousarray(motion, dtype=np.float32)
        motion.setflags(write=False)

        elapsed_seconds = time.perf_counter() - started_at

        return GenerationResult(
            motion=motion,
            fps=self.settings.fps,
            token_count=token_length,
            elapsed_seconds=elapsed_seconds,
            from_cache=False,
            keywords=[str(keyword) for keyword in keywords],
        )

    def info(self) -> dict:
        return {
            "ready": self.ready,
            "device": self.device_name,
            "model_id": self.settings.model_id,
            "fps": self.settings.fps,
            "frame_dimension": FRAME_DIMENSION,
            "maximum_motion_tokens": self.maximum_motion_tokens,
            "maximum_seconds_per_sentence": (
                self.maximum_motion_tokens * FRAMES_PER_TOKEN / self.settings.fps
                if self.maximum_motion_tokens > 0
                else None
            ),
            "load_seconds": round(self.load_seconds, 2),
            "cache_entries": len(self.cache),
            "cache_hits": self.cache_hits,
            "cache_misses": self.cache_misses,
            "use_retrieval_default": self.settings.use_retrieval,
            "do_sample_default": self.settings.do_sample,
        }