| 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, |
| } |
|
|
