"""Cortex-0 — moteur d'inférence unifié (texte / image / audio / vidéo). Rassemble les modèles maison des Phases 6–9 derrière une API simple, utilisée par le serveur FastAPI (`src/api.py`) et la démo Gradio (`app/gradio_app.py`). Honnêteté : si aucun checkpoint entraîné n'est fourni, le moteur construit des modèles aux poids **aléatoires init** et renvoie ``trained=False``. Le pipeline (décodage des entrées, formatage, génération) reste réel et testable ; seule la qualité des sorties dépend de l'entraînement. Aucun poids pré-entraîné n'est chargé. """ from __future__ import annotations import io from dataclasses import dataclass from pathlib import Path import numpy as np import torch from src.data.audio_features import LogMelSpectrogram from src.model.audio_encoder import AudioConfig, AudioEncoder from src.model.multimodal import VideoEncoder from src.model.projectors import MLPProjector from src.model.text_model import CortexTextLM, TextConfig from src.model.vit import VisionEncoder, ViTConfig from src.tokenizer import CortexTokenizer, load_tokenizer from src.tokenizer.special_tokens import BOS_ID, EOS_ID, SPECIAL_TOKEN_IDS IMG_S, IMG_P, IMG_E = (SPECIAL_TOKEN_IDS["<|image_start|>"], SPECIAL_TOKEN_IDS["<|image_pad|>"], SPECIAL_TOKEN_IDS["<|image_end|>"]) AUD_S, AUD_P, AUD_E = (SPECIAL_TOKEN_IDS["<|audio_start|>"], SPECIAL_TOKEN_IDS["<|audio_pad|>"], SPECIAL_TOKEN_IDS["<|audio_end|>"]) VID_S, VID_P, VID_E = (SPECIAL_TOKEN_IDS["<|video_start|>"], SPECIAL_TOKEN_IDS["<|video_pad|>"], SPECIAL_TOKEN_IDS["<|video_end|>"]) @dataclass class CortexDemoConfig: """Tailles de la démo (alignées sur les phases jouets, tournent sur CPU).""" grid: int = 32 n_frames: int = 4 patch_size: int = 8 n_mels: int = 64 n_frames_audio: int = 96 patch_mels: int = 16 patch_frames: int = 16 text_layers: int = 2 text_embd: int = 96 text_heads: int = 4 vision_layers: int = 1 vision_embd: int = 96 audio_layers: int = 1 audio_embd: int = 96 ctx_len: int = 320 class CortexEngine: """Moteur multimodal : image, audio, vidéo -> texte (décodeur unique).""" def __init__(self, checkpoint: str | Path | None = None, demo: CortexDemoConfig | None = None, device: str = "cpu"): self.cfg = demo or CortexDemoConfig() self.device = torch.device(device) self.tokenizer: CortexTokenizer = load_tokenizer() self.trained = False self._build_models() if checkpoint is not None and Path(checkpoint).is_file(): self.load(checkpoint) # -- construction ------------------------------------------------------ def _text_cfg(self) -> TextConfig: return TextConfig(vocab_size=self.tokenizer.vocab_size, n_layer=self.cfg.text_layers, n_head=self.cfg.text_heads, n_embd=self.cfg.text_embd, ctx_len=self.cfg.ctx_len, tie_embeddings=True) def _vit_cfg(self) -> ViTConfig: return ViTConfig(image_size=self.cfg.grid, patch_size=self.cfg.patch_size, n_layer=self.cfg.vision_layers, n_head=4, n_embd=self.cfg.vision_embd) def _audio_cfg(self) -> AudioConfig: return AudioConfig(n_mels=self.cfg.n_mels, n_frames=self.cfg.n_frames_audio, patch_mels=self.cfg.patch_mels, patch_frames=self.cfg.patch_frames, n_layer=self.cfg.audio_layers, n_head=4, n_embd=self.cfg.audio_embd) def _build_models(self) -> None: self.text = CortexTextLM(self._text_cfg()) self.vision = VisionEncoder(self._vit_cfg()) self.audio = AudioEncoder(self._audio_cfg()) self.video = VideoEncoder(self._vit_cfg(), self.cfg.n_frames, temporal_layers=1) self.image_projector = MLPProjector(self.cfg.vision_embd, self.cfg.text_embd) self.audio_projector = MLPProjector(self.cfg.audio_embd, self.cfg.text_embd) self.video_projector = MLPProjector(self.cfg.vision_embd, self.cfg.text_embd) self.mel = LogMelSpectrogram(sr=16000, n_fft=256, hop_length=64, n_mels=self.cfg.n_mels) self.eval() def modules(self) -> list[torch.nn.Module]: return [self.text, self.vision, self.audio, self.video, self.image_projector, self.audio_projector, self.video_projector] def eval(self) -> None: for m in self.modules(): m.to(self.device).eval() @property def n_image_tokens(self) -> int: return self.vision.patch_embed.n_patches @property def n_audio_tokens(self) -> int: return self.audio.n_audio_tokens @property def n_video_tokens(self) -> int: return self.cfg.n_frames * self.video.n_patches # -- checkpoints ------------------------------------------------------- def load(self, path: str | Path) -> None: state = torch.load(path, map_location=self.device, weights_only=False) if "model" in state and isinstance(state["model"], dict): state = state["model"] mods = {"text.": self.text, "vision.": self.vision, "audio.": self.audio, "video.": self.video, "image_projector.": self.image_projector, "audio_projector.": self.audio_projector, "video_projector.": self.video_projector} loaded = 0 for mod, obj in mods.items(): sub = {k[len(mod):]: v for k, v in state.items() if k.startswith(mod)} if sub: obj.load_state_dict(sub, strict=False) loaded += 1 self.trained = loaded > 0 print(f"[CortexEngine] checkpoint chargé : {path} (modules={loaded})") # -- entrées -> tenseurs ---------------------------------------------- def _image_tensor(self, image) -> torch.Tensor: arr = _to_numpy_image(image) arr = _resize_nearest(arr, self.cfg.grid) if arr.ndim == 2: arr = np.stack([arr] * 3, axis=-1) if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[0] != arr.shape[1]: arr = np.transpose(arr, (1, 2, 0)) arr = arr[:, :, :3].astype(np.float32) arr = arr / 255.0 if arr.max() > 1.5 else arr return torch.from_numpy(np.transpose(arr, (2, 0, 1)))[None].to(self.device) def _mel_tensor(self, audio) -> torch.Tensor: wav = _to_waveform(audio) with torch.no_grad(): m = self.mel(torch.from_numpy(wav)[None]).squeeze(1) t = m.size(-1) m = m[:, :, :self.cfg.n_frames_audio] if t >= self.cfg.n_frames_audio else \ torch.nn.functional.pad(m, (0, self.cfg.n_frames_audio - t)) return m.unsqueeze(1).to(self.device) def _video_tensor(self, frames) -> torch.Tensor: if isinstance(frames, np.ndarray) and frames.ndim == 5: arr = frames else: arr = np.stack([_to_numpy_image(f) for f in frames]) arr = np.stack([_resize_nearest(f, self.cfg.grid) for f in arr]) if arr.ndim == 4 and arr.shape[-1] in (1, 3, 4) and arr.shape[-1] != self.cfg.grid: arr = np.transpose(arr, (0, 3, 1, 2)) arr = arr[:, :3].astype(np.float32) arr = arr / 255.0 if arr.max() > 1.5 else arr return torch.from_numpy(arr)[None].to(self.device) # -- génération -------------------------------------------------------- def _features_for(self, kind: str, payload) -> torch.Tensor: if kind == "image": return self.image_projector(self.vision(self._image_tensor(payload))[:, 1:]) if kind == "audio": return self.audio_projector(self.audio(self._mel_tensor(payload))) if kind == "video": return self.video_projector(self.video(self._video_tensor(payload))) return None @torch.no_grad() def _run(self, ids, mask, features, max_new_tokens) -> str: inp = torch.tensor([ids], dtype=torch.long, device=self.device) mk = torch.tensor([mask], dtype=torch.bool, device=self.device) for _ in range(max_new_tokens): emb = self.text.token_embedding(inp).clone() if features is not None and bool(mk.any()): emb[mk] = features.reshape(-1, features.size(-1)) logits, _ = self.text(inputs_embeds=emb) nxt = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True) inp = torch.cat([inp, nxt], dim=1) mk = torch.cat([mk, torch.zeros_like(nxt, dtype=torch.bool)], dim=1) if (nxt == EOS_ID).all(): break return self.tokenizer.decode(inp[0].tolist()[len(ids):], skip_special_tokens=True).strip() def _single(self, prompt, s_id, pad_id, n_tok, e_id, features, max_new_tokens) -> str: ids = [BOS_ID, *self.tokenizer.encode(prompt)] mask = [0] * len(ids) ids += [s_id] + [pad_id] * n_tok + [e_id] mask += [0] + [1] * n_tok + [0] return self._run(ids, mask, features, max_new_tokens) # -- API publique ------------------------------------------------------ def describe_image(self, image, max_new_tokens: int = 20) -> dict: feat = self._features_for("image", image) text = self._single("Décris l'image :", IMG_S, IMG_P, self.n_image_tokens, IMG_E, feat, max_new_tokens) return {"text": text, "trained": self.trained} def transcribe_audio(self, audio, max_new_tokens: int = 20) -> dict: feat = self._features_for("audio", audio) text = self._single("Transcris :", AUD_S, AUD_P, self.n_audio_tokens, AUD_E, feat, max_new_tokens) return {"text": text, "trained": self.trained} def describe_video(self, frames, audio=None, max_new_tokens: int = 24) -> dict: vfeat = self._features_for("video", frames) prompt = "Décris la vidéo :" if audio is None else "Que se passe-t-il ?" ids = [BOS_ID, *self.tokenizer.encode(prompt)] mask = [0] * len(ids) ids += [VID_S] + [VID_P] * self.n_video_tokens + [VID_E] mask += [0] + [1] * self.n_video_tokens + [0] feats = [vfeat] if audio is not None: ids += [AUD_S] + [AUD_P] * self.n_audio_tokens + [AUD_E] mask += [0] + [1] * self.n_audio_tokens + [0] feats.append(self._features_for("audio", audio)) fused = torch.cat(feats, dim=1) return {"text": self._run(ids, mask, fused, max_new_tokens), "trained": self.trained} def chat(self, prompt: str, max_new_tokens: int = 40) -> dict: ids = [BOS_ID, *self.tokenizer.encode(prompt)] text = self._run(ids, [0] * len(ids), None, max_new_tokens) return {"text": text, "trained": self.trained} def _to_numpy_image(image) -> np.ndarray: """Convertit `PIL.Image`, chemin, bytes ou ndarray en tableau numpy.""" if isinstance(image, np.ndarray): return image if isinstance(image, (str, Path)): from PIL import Image return np.asarray(Image.open(image).convert("RGB")) if isinstance(image, (bytes, bytearray)): from PIL import Image return np.asarray(Image.open(io.BytesIO(bytes(image))).convert("RGB")) if hasattr(image, "convert"): # PIL.Image return np.asarray(image.convert("RGB")) raise TypeError(f"image non supportée : {type(image)!r}") def _resize_nearest(arr: np.ndarray, size: int) -> np.ndarray: """Redimensionne une image HWC vers ``size × size`` (plus proche voisin, numpy).""" h, w = arr.shape[0], arr.shape[1] if h == size and w == size: return arr ys = (np.arange(size) * h // size).clip(0, h - 1) xs = (np.arange(size) * w // size).clip(0, w - 1) return arr[ys][:, xs] def _decode_wav(raw: bytes) -> np.ndarray: import wave with wave.open(io.BytesIO(raw), "rb") as w: n, sw, ch = w.getnframes(), w.getsampwidth(), w.getnchannels() data = w.readframes(n) dt = {1: np.int8, 2: np.int16, 4: np.int32}[sw] wav = np.frombuffer(data, dtype=dt).astype(np.float32) if ch > 1: wav = wav.reshape(-1, ch).mean(axis=1) return wav / max(1.0, float(np.abs(wav).max())) def _to_waveform(audio) -> np.ndarray: """Convertit un chemin WAV, des bytes, ou un ndarray en ``float32 (n,)`` mono.""" if isinstance(audio, np.ndarray): wav = audio.astype(np.float32) if wav.ndim > 1: wav = wav.mean(axis=1) return wav if isinstance(audio, (str, Path)): return _decode_wav(Path(audio).read_bytes()) if isinstance(audio, (bytes, bytearray)): return _decode_wav(bytes(audio)) raise TypeError(f"audio non supporté : {type(audio)!r}") def default_engine(checkpoint: str | Path | None = None) -> CortexEngine: return CortexEngine(checkpoint=checkpoint)