Spaces:
Running on Zero
Running on Zero
| """Strict checkpoint-backed Falcon inference for synchronized two-second windows. | |
| This is windowed inference, not a stateful streaming architecture. The checkpoint | |
| expects mono 16 kHz audio and 25 fps, aligned grayscale mouth crops of 96×96. | |
| The browser/server is responsible for lip tracking and capture synchronization. | |
| """ | |
| from __future__ import annotations | |
| import copy | |
| import os | |
| from pathlib import Path | |
| import threading | |
| import time | |
| import numpy as np | |
| import torch | |
| import yaml | |
| from falcon_model.models.Falcon_Contrastive import FalconSpeakerContrastive | |
| from falcon_model.models.Speaker_Contrastive_v2 import SpeakerContrastiveV2 | |
| from falcon_model.videomodels.frcnn_videomodel import FRCNNVideoModel | |
| SAMPLE_RATE = 16_000 | |
| VIDEO_FPS = 25 | |
| WINDOW_SAMPLES = 32_000 | |
| WINDOW_FRAMES = 50 | |
| ROOT = Path(__file__).resolve().parent | |
| # Weights come from these Hugging Face short links: downloaded on first use and cached in | |
| # ~/.cache/huggingface. A local file or hf://<org>/<repo>/<file> can be passed instead. | |
| VARIANTS = { | |
| "small": {"config": "Falcon.yml", "weights": "falcon.pt", "repo": "JusperLee/Falcon", | |
| "architecture": "FalconSpeakerContrastive"}, | |
| "big": {"config": "FalconBig.yml", "weights": "falcon-big.pt", "repo": "JusperLee/Falcon-Big", | |
| "architecture": "SpeakerContrastiveV2"}, | |
| } | |
| MODEL_CLASSES = {"FalconSpeakerContrastive": FalconSpeakerContrastive, "SpeakerContrastiveV2": SpeakerContrastiveV2} | |
| # The Hub counts a model download when config.yaml is requested; requests for the .pt weights are not counted. | |
| HUB_COUNTED_FILE = "config.yaml" | |
| def variant_for(name: str) -> str: | |
| """Variant id for a short link ("JusperLee/Falcon", "JusperLee/Falcon-Big") or "small"/"big".""" | |
| for variant, info in VARIANTS.items(): | |
| if name in (variant, info["repo"]): | |
| return variant | |
| raise ValueError(f"Unknown Falcon model {name!r}; use {' or '.join(v['repo'] for v in VARIANTS.values())}") | |
| def weights_source(variant: str) -> str | None: | |
| """'cached' (already downloaded), 'huggingface' (downloads on first load) or None.""" | |
| try: | |
| from huggingface_hub import try_to_load_from_cache | |
| except ImportError: | |
| return None | |
| cached = try_to_load_from_cache(VARIANTS[variant]["repo"], VARIANTS[variant]["weights"]) | |
| return "cached" if isinstance(cached, str) else "huggingface" | |
| def resolve_checkpoint(variant: str, checkpoint: str | Path | None = None) -> Path: | |
| """The variant's Hugging Face weights, or a local file / hf://<org>/<repo>/<file> override.""" | |
| spec = str(checkpoint) if checkpoint else None | |
| if spec and not spec.startswith("hf://"): | |
| return Path(spec) | |
| if spec: | |
| repo_id, _, filename = spec[len("hf://"):].rpartition("/") | |
| else: | |
| repo_id, filename = VARIANTS[variant]["repo"], VARIANTS[variant]["weights"] | |
| if not repo_id or not filename: | |
| raise ValueError(f"Expected hf://<org>/<repo>/<file>, got {spec}") | |
| try: | |
| from huggingface_hub import hf_hub_download | |
| except ImportError as exc: | |
| raise FileNotFoundError("huggingface_hub is not installed (pip install huggingface_hub); " | |
| "or pass a local weights file.") from exc | |
| try: | |
| path = Path(hf_hub_download(repo_id=repo_id, filename=filename)) | |
| except Exception as exc: | |
| raise FileNotFoundError(f"Could not download {filename} from Hugging Face repo {repo_id}: {exc}. " | |
| "Set HF_ENDPOINT to a mirror, or pass a local weights file.") from exc | |
| count_download(repo_id) | |
| return path | |
| def count_download(repo_id: str) -> threading.Thread: | |
| """HEAD the repo's config.yaml in the background so this load shows in the Hub's download stats. | |
| Best effort and never delays loading: offline, a mirror or a repo without the file only skips | |
| the count. HF_HUB_DISABLE_TELEMETRY=1 turns it off. | |
| """ | |
| def head(): | |
| try: | |
| from huggingface_hub import constants, get_hf_file_metadata, hf_hub_url | |
| if not constants.HF_HUB_DISABLE_TELEMETRY: | |
| get_hf_file_metadata(hf_hub_url(repo_id, HUB_COUNTED_FILE)) | |
| except Exception: | |
| pass | |
| thread = threading.Thread(target=head, name="hub-download-count", daemon=True) | |
| thread.start() | |
| return thread | |
| def pick_device(requested: str = "auto") -> str: | |
| """'auto' prefers CUDA, then Apple Metal (MPS), then the CPU.""" | |
| if requested not in {"auto", "cpu", "cuda", "mps"}: | |
| raise ValueError("device must be 'auto', 'cpu', 'cuda' or 'mps'") | |
| cuda, mps = torch.cuda.is_available(), torch.backends.mps.is_available() | |
| if requested == "cuda" and not cuda: | |
| raise RuntimeError("CUDA is unavailable in this Python process; install a CUDA build of PyTorch or choose CPU.") | |
| if requested == "mps" and not mps: | |
| raise RuntimeError("Apple Metal/MPS is unavailable in this Python process; choose CPU.") | |
| if requested == "auto": | |
| return "cuda" if cuda else "mps" if mps else "cpu" | |
| return requested | |
| def preprocess_mouth(frames: np.ndarray) -> torch.Tensor: | |
| """Match get_preprocessing_pipelines()['val'] from the training repository.""" | |
| if not isinstance(frames, np.ndarray) or frames.dtype != np.uint8: | |
| raise ValueError("Mouth frames must be a uint8 NumPy array.") | |
| if frames.shape != (WINDOW_FRAMES, 96, 96): | |
| raise ValueError(f"Expected {WINDOW_FRAMES} grayscale 96×96 mouth crops; got {frames.shape}.") | |
| # Validation CenterCrop(88, 88), Normalize(0, 255), Normalize(.421, .165). | |
| x = frames[:, 4:92, 4:92].astype(np.float32) / 255.0 | |
| x = (x - 0.421) / 0.165 | |
| return torch.from_numpy(np.ascontiguousarray(x))[None, None] | |
| class FalconRuntime: | |
| """One model pair, serialized across requests; never returns the input as output.""" | |
| def __init__(self, variant: str = "small", checkpoint_path: str | Path | None = None): | |
| if variant not in VARIANTS: | |
| raise ValueError("variant must be 'small' or 'big'") | |
| self.variant = variant | |
| self.config_path = ROOT / "configs" / VARIANTS[variant]["config"] | |
| self.checkpoint_spec = checkpoint_path | |
| # Resolved (downloaded if needed) in load(); until then only the name is reported. | |
| local = checkpoint_path and not str(checkpoint_path).startswith("hf://") | |
| self.checkpoint_path = Path(checkpoint_path if local else VARIANTS[variant]["weights"]) | |
| self._lock = threading.RLock() | |
| self.audio_model = None | |
| self.video_model = None | |
| self.device = "unloaded" | |
| self.requested_device = "auto" | |
| self.ready = False | |
| self.error: str | None = None | |
| self.warning: str | None = None | |
| self.last_inference_ms: float | None = None | |
| self.warmup_ms: float | None = None | |
| self.audio_keys = 0 | |
| self.video_keys = 0 | |
| self.parameter_count = 0 | |
| def load(self, device: str = "auto") -> "FalconRuntime": | |
| if device not in {"auto", "cpu", "cuda", "mps"}: | |
| raise ValueError("device must be 'auto', 'cpu', 'cuda' or 'mps'") | |
| with self._lock: | |
| self.ready = False | |
| self.error = None | |
| self.warning = None | |
| self.requested_device = device | |
| try: | |
| self.checkpoint_path = resolve_checkpoint(self.variant, self.checkpoint_spec) | |
| if not self.checkpoint_path.is_file(): | |
| raise FileNotFoundError(f"Falcon checkpoint not found: {self.checkpoint_path}") | |
| # All repository modules are loaded locally; there are no network downloads. | |
| config = yaml.safe_load(self.config_path.read_text()) | |
| if "train_conf" in config: | |
| config = config["train_conf"] | |
| audio_config = copy.deepcopy(config["audionet"]["audionet_config"]) | |
| video_config = copy.deepcopy(config["videonet"]["videonet_config"]) | |
| video_config["pretrain"] = None | |
| architecture = config["audionet"]["audionet_name"] | |
| if architecture != VARIANTS[self.variant]["architecture"]: | |
| raise ValueError(f"Unexpected architecture for {self.variant}: {architecture}") | |
| if config["videonet"]["videonet_name"] != "FRCNNVideoModel": | |
| raise ValueError(f"Unsupported video architecture: {config['videonet']['videonet_name']}") | |
| self.audio_model = MODEL_CLASSES[architecture](sample_rate=SAMPLE_RATE, **audio_config) | |
| self.video_model = FRCNNVideoModel(**video_config) | |
| # weights_only avoids executing arbitrary checkpoint Python objects. | |
| ckpt = torch.load(self.checkpoint_path, map_location="cpu", weights_only=True) | |
| state = ckpt.get("state_dict", ckpt) | |
| audio_state = {k[len("audio_model."):]: v for k, v in state.items() if k.startswith("audio_model.")} | |
| video_state = {k[len("video_model."):]: v for k, v in state.items() if k.startswith("video_model.")} | |
| if not audio_state or not video_state: | |
| raise ValueError("Checkpoint must contain both audio_model and video_model weights.") | |
| # Missing shared/repeats config keys can silently disable almost all layers. | |
| # Strict loading rejects that failure rather than silently using partial weights. | |
| self.audio_model.load_state_dict(audio_state, strict=True) | |
| self.video_model.load_state_dict(video_state, strict=True) | |
| self.audio_keys, self.video_keys = len(audio_state), len(video_state) | |
| self.parameter_count = sum(p.numel() for m in (self.audio_model, self.video_model) for p in m.parameters()) | |
| del ckpt, state, audio_state, video_state | |
| # Small recurrent kernels generally slow down with excessive CPU threads. | |
| threads = max(1, min(8, int(os.environ.get("FALCON_CPU_THREADS", "4")))) | |
| torch.set_num_threads(threads) | |
| for model in (self.audio_model, self.video_model): | |
| # The original video model's train() override does not return self. | |
| model.eval() | |
| model.requires_grad_(False) | |
| self.device = pick_device(device) | |
| # Warm up real model paths before reporting ready; in auto mode a GPU that | |
| # cannot run the model (driver, memory, unsupported ops) falls back to the CPU. | |
| started = time.perf_counter() | |
| try: | |
| self._warmup() | |
| except Exception as exc: | |
| if self.device == "cpu" or device != "auto": | |
| raise | |
| name = {"cuda": "CUDA", "mps": "Metal"}[self.device] | |
| self.warning = f"{name} warmup failed; using CPU: {str(exc).splitlines()[0][:180]}" | |
| self.device = "cpu" | |
| self._warmup() | |
| self.warmup_ms = round((time.perf_counter() - started) * 1000, 1) | |
| self.ready = True | |
| return self | |
| except Exception as exc: | |
| self.error = str(exc) | |
| raise | |
| def _warmup(self) -> None: | |
| for model in (self.audio_model, self.video_model): | |
| model.to(self.device) | |
| self._forward(np.zeros(WINDOW_SAMPLES, dtype=np.float32), np.full((WINDOW_FRAMES, 96, 96), 107, np.uint8)) | |
| def from_pretrained(cls, name: str = "JusperLee/Falcon", device: str = "auto") -> "FalconRuntime": | |
| """Load by Hugging Face short link ("JusperLee/Falcon", "JusperLee/Falcon-Big") or "small"/"big". | |
| device: "auto" (CUDA, then Metal, then CPU), "cuda", "mps" or "cpu". | |
| """ | |
| return cls(variant_for(name)).load(device) | |
| def status(self) -> dict: | |
| # Do not acquire the inference lock: status remains responsive during processing. | |
| return { | |
| "ready": self.ready, | |
| "variant": self.variant, | |
| "architecture": VARIANTS[self.variant]["architecture"], | |
| "device": self.device, | |
| "requested_device": self.requested_device, | |
| "checkpoint": self.checkpoint_path.name, | |
| "sample_rate": SAMPLE_RATE, | |
| "video_fps": VIDEO_FPS, | |
| "window_seconds": 2.0, | |
| "audio_weight_keys": self.audio_keys, | |
| "video_weight_keys": self.video_keys, | |
| "parameters": self.parameter_count, | |
| "last_inference_ms": self.last_inference_ms, | |
| "warmup_ms": self.warmup_ms, | |
| "error": self.error, | |
| "warning": self.warning, | |
| "mode": "windowed_audio_visual", | |
| } | |
| def _forward(self, audio: np.ndarray, frames: np.ndarray) -> np.ndarray: | |
| mouth = preprocess_mouth(frames).to(self.device) | |
| waveform = torch.from_numpy(np.ascontiguousarray(audio)).unsqueeze(0).to(self.device) | |
| mouth_embedding = self.video_model(mouth) | |
| separated = self.audio_model(waveform, mouth_embedding) | |
| output = separated[0, 0].float().cpu().numpy().copy() | |
| if output.shape != (WINDOW_SAMPLES,) or not np.isfinite(output).all(): | |
| raise RuntimeError("Falcon produced invalid output; this window was not played.") | |
| return output | |
| def infer(self, audio: np.ndarray, frames: np.ndarray) -> np.ndarray: | |
| if not isinstance(audio, np.ndarray) or audio.dtype != np.float32 or audio.shape != (WINDOW_SAMPLES,): | |
| raise ValueError(f"Expected float32 mono audio shape ({WINDOW_SAMPLES},) at 16 kHz.") | |
| if not np.isfinite(audio).all(): | |
| raise ValueError("Audio contains non-finite samples.") | |
| with self._lock: | |
| if not self.ready: | |
| raise RuntimeError(self.error or "Falcon model has not finished loading.") | |
| started = time.perf_counter() | |
| output = self._forward(audio, frames) | |
| self.last_inference_ms = round((time.perf_counter() - started) * 1000, 1) | |
| return output | |