src/inference.py

"""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)