SolitaryThinker
431f4daddb
Three more capabilities on distinct untested surfaces (167 tests pass, no new runtime primitive). A — Adapter plane (§9.19): one base + swappable LoRA/ControlNet adapters, selected per request (DiffusionParams.adapters); AdapterDenoiseLoop applies each active adapter's velocity delta. Per-request selection changes output, multi-LoRA composes, ControlNet conditions on a control image, mixed-adapter requests interleave without smearing, hot-swap changes generation, cache key partitions by adapter stack (the adapter_versions field, previously declared-only). ToyLoRA/ToyControlNet; models/adapters/. 6 tests. B — Non-linear workflows (§9.17): ParallelWorkflow (fan-out: one input → N models → merged) and BestOfNWorkflow (generate N → score with the served reward card → return best; inference-time scaling). The shapes a linear chain can't express. 4 tests. D — RL→distill flywheel (§9.18): run_flywheel RL-improves the base (NFT), then distills FROM the RL'd model (DMD2 teacher = RL'd policy) into a faster card, recording the base→rl→distilled provenance chain in RecipeSpec.parents. The distilled student is measurably closer to the RL'd teacher than the base; the distilled card serves few-step. training/flywheel.py. 4 tests. designv4 §9.17–§9.19 + layout/closing/counts (167 tests, 29 files).
416 lines
20 KiB
Python
416 lines
20 KiB
Python
"""Toy numpy components — the CPU-testable backend (design_v3 §17 honesty note).
|
||
|
||
There is no GPU/torch/weights in this environment, so the heavy Wan/LTX neural forwards are
|
||
represented by small, deterministic numpy components. They exercise the *real* loop control
|
||
flow, CFG/flow-shift policies, scheduler steps, cache reuse, parity gates, and training-method
|
||
math with real numbers — just with a toy network instead of a 1.3B DiT.
|
||
|
||
The contract these toys honor is what matters:
|
||
* deterministic given (weights, inputs) → bit-reproducible (the interleave/parity gates rely on this);
|
||
* same prompt → same text embedding → feature-cache reuse is correct;
|
||
* a velocity ("flow") prediction the flow-match sampler consumes (Wan/LTX prediction_type).
|
||
|
||
On a GPU box, ``ComponentSpec.factory`` swaps these for lazy torch adapters wrapping the real
|
||
``fastvideo.models`` modules + weights (see each model's ``components.py``); the loops, policies,
|
||
scheduler, caches, parity, and training code are unchanged. That is the whole point of the
|
||
(recipe, runtime) separation.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
|
||
import numpy as np
|
||
|
||
LATENT_CHANNELS = 4
|
||
TEXT_SEQ = 4
|
||
TEXT_DIM = 8
|
||
|
||
|
||
def _seed_from(s: str) -> int:
|
||
return int(hashlib.sha256(s.encode()).hexdigest()[:8], 16)
|
||
|
||
|
||
class ToyTextEncoder:
|
||
"""Deterministic text→embedding (same text ⇒ same embedding ⇒ feature-cache reuse works)."""
|
||
|
||
def __init__(self, seq: int = TEXT_SEQ, dim: int = TEXT_DIM):
|
||
self.seq, self.dim = seq, dim
|
||
|
||
def encode(self, text: str) -> np.ndarray:
|
||
rng = np.random.default_rng(_seed_from("txt:" + (text or "<empty>")))
|
||
return (rng.standard_normal((self.seq, self.dim)) * 0.1).astype("float32")
|
||
|
||
|
||
class ToyDiT:
|
||
"""A tiny deterministic velocity predictor (stands in for the 1.3B DiT).
|
||
|
||
velocity = tanh( Wx·latent + Wt·σ + s·mean(text_embed) + c·mean(context) )
|
||
Channel-mixing via ``Wx`` makes the trajectory non-trivial; ``context`` carries causal
|
||
chunk history for the self-forcing chunk_rollout loop.
|
||
"""
|
||
|
||
def __init__(self, channels: int = LATENT_CHANNELS, seed: int = 0):
|
||
self.C = channels
|
||
rng = np.random.default_rng(seed)
|
||
self.w_x = (rng.standard_normal((channels, channels)) * 0.15).astype("float32")
|
||
self.w_t = (rng.standard_normal(channels) * 0.1).astype("float32")
|
||
self.s_text = 0.5
|
||
self.s_ctx = 0.3
|
||
|
||
def _pre_tanh(self, latent: np.ndarray, text_embed, sigma: float, context) -> np.ndarray:
|
||
mixed = np.tensordot(self.w_x, latent, axes=([1], [0])) # [C,...] channel mix
|
||
mixed = mixed + self.w_t.reshape((self.C,) + (1,) * (latent.ndim - 1)) * float(sigma)
|
||
cond = float(np.mean(text_embed)) if text_embed is not None else 0.0
|
||
mixed = mixed + self.s_text * cond
|
||
if context is not None:
|
||
mixed = mixed + self.s_ctx * float(np.mean(context))
|
||
return mixed
|
||
|
||
def __call__(self, latent: np.ndarray, text_embed: np.ndarray | None, sigma: float,
|
||
context: np.ndarray | None = None) -> np.ndarray:
|
||
latent = np.asarray(latent, dtype=np.float32)
|
||
return np.tanh(self._pre_tanh(latent, text_embed, sigma, context)).astype("float32")
|
||
|
||
# --- minimal trainable surface (so training methods do real optimizer steps) --------- #
|
||
def clone(self) -> "ToyDiT":
|
||
c = ToyDiT.__new__(ToyDiT)
|
||
c.C, c.s_text, c.s_ctx = self.C, self.s_text, self.s_ctx
|
||
c.w_x, c.w_t = self.w_x.copy(), self.w_t.copy()
|
||
return c
|
||
|
||
def blend_from(self, other: "ToyDiT", decay: float) -> None:
|
||
"""EMA / decay-blended-old-policy update: self ← decay·self + (1-decay)·other (design_v3 §10)."""
|
||
self.w_x = (decay * self.w_x + (1.0 - decay) * other.w_x).astype("float32")
|
||
self.w_t = (decay * self.w_t + (1.0 - decay) * other.w_t).astype("float32")
|
||
|
||
def copy_from(self, other: "ToyDiT") -> None:
|
||
self.w_x, self.w_t = other.w_x.copy(), other.w_t.copy()
|
||
|
||
def mse_grad_step(self, latent: np.ndarray, text_embed, sigma: float, target: np.ndarray,
|
||
lr: float, context=None, weight: float = 1.0) -> tuple[float, float]:
|
||
"""One exact SGD step minimizing ``weight·MSE(forward(latent), target)`` w.r.t. ``w_x``.
|
||
|
||
Returns (loss, grad_norm). The grad_norm is the #1396-style per-method regression signal.
|
||
This is a *real* gradient (chain rule through tanh), so the toy student actually learns —
|
||
the structural stand-in for the FSDP/optimizer step a GPU trainer would run."""
|
||
latent = np.asarray(latent, dtype=np.float32)
|
||
target = np.asarray(target, dtype=np.float32)
|
||
z = self._pre_tanh(latent, text_embed, sigma, context)
|
||
pred = np.tanh(z)
|
||
err = pred - target
|
||
loss = float(weight * np.mean(err ** 2))
|
||
n = err.size
|
||
g = weight * (2.0 * err / n) * (1.0 - pred ** 2)
|
||
axes = (list(range(1, g.ndim)), list(range(1, latent.ndim)))
|
||
grad_wx = np.tensordot(g, latent, axes=axes) # [C_out, C_in]
|
||
grad_norm = float(np.linalg.norm(grad_wx))
|
||
self.w_x = (self.w_x - lr * grad_wx).astype("float32")
|
||
return loss, grad_norm
|
||
|
||
|
||
class ToyTokenizer:
|
||
"""Deterministic byte-ish tokenizer for the omni AR pathway (phase 2)."""
|
||
EOS = 0
|
||
VOCAB = 256
|
||
|
||
def encode(self, text: str) -> list[int]:
|
||
toks = [(ord(c) % (self.VOCAB - 1)) + 1 for c in (text or "")[:16]]
|
||
return toks or [1]
|
||
|
||
def decode(self, tokens) -> str:
|
||
return "".join(chr(32 + (int(t) % 90)) for t in tokens)
|
||
|
||
|
||
class ToyMoTDiT(ToyDiT):
|
||
"""Mixture-of-Transformers stand-in: ONE resident module that runs BOTH an AR (understanding)
|
||
pathway and a diffusion (generation) pathway on shared weights (design_v3 §4.2, §16).
|
||
|
||
``ar_forward`` is the und pathway (next-token); ``__call__`` (inherited) is the gen pathway
|
||
(velocity). Binding both the ar_decode and diffusion_denoise loops to one instance of this
|
||
component is the MoT requirement no DAG-of-engines can express.
|
||
"""
|
||
VOCAB = 256
|
||
EOS = 0
|
||
|
||
def ar_forward(self, tokens) -> int:
|
||
# deterministic next-token from the context (greedy ⇒ trivially interleave-safe)
|
||
ctx = sum(int(t) for t in tokens)
|
||
nxt = (ctx * 7 + 13) % self.VOCAB
|
||
return int(nxt)
|
||
|
||
def reasoner_embed(self, tokens) -> np.ndarray:
|
||
"""Pack the und-pathway tokens into a conditioning embed the gen pathway consumes
|
||
(the Cosmos3 'prompt upsampling before diffusion in the same request')."""
|
||
rng = np.random.default_rng((sum(int(t) for t in tokens) + 1) % (1 << 31))
|
||
return (rng.standard_normal((4, 8)) * 0.1).astype("float32")
|
||
|
||
|
||
class ToyPromptRefiner:
|
||
"""Toy LM prompt-refiner — UniRL/PromptRL's Qwen role (a *separate expert*, not MoT-shared).
|
||
|
||
A categorical policy over a small set of refinement "actions" (stand-in for the LM emitting a
|
||
refined prompt). Sampling an action picks an embedding offset that shifts the diffusion
|
||
conditioning; the realized reward then drives the policy by REINFORCE. This is the *second
|
||
trainable expert* in the unified-RL stress test: one reward → a token-policy-gradient update
|
||
here AND a FlowGRPO update on the DiT, under one advantage (design_v3 §10; PromptRL §6).
|
||
|
||
Real gradient: loss = −A·logπ(a); ∇_logits = −A·(onehot(a) − softmax). So the toy LM actually
|
||
learns to prefer the reward-favored action — the LM-side analogue of ToyDiT.mse_grad_step.
|
||
"""
|
||
|
||
def __init__(self, n_actions: int = 8, dim: int = TEXT_DIM, seq: int = TEXT_SEQ, seed: int = 0):
|
||
self.n = int(n_actions)
|
||
self.dim, self.seq = int(dim), int(seq)
|
||
self.logits = np.zeros(self.n, dtype="float64") # the trainable policy params
|
||
rng = np.random.default_rng(seed)
|
||
# each action ⇒ a fixed conditioning offset (the "content" of that refined prompt)
|
||
self._emb = (rng.standard_normal((self.n, self.seq, self.dim)) * 0.1).astype("float32")
|
||
|
||
def _probs(self) -> np.ndarray:
|
||
z = self.logits - self.logits.max()
|
||
e = np.exp(z)
|
||
return e / e.sum()
|
||
|
||
def sample_refinement(self, rng) -> tuple[int, float]:
|
||
"""Sample an action (exploration) and return (action, logπ(action))."""
|
||
p = self._probs()
|
||
a = int(rng.choice(self.n, p=p))
|
||
return a, float(np.log(p[a] + 1e-12))
|
||
|
||
def argmax_refinement(self) -> tuple[int, float]:
|
||
a = int(np.argmax(self.logits))
|
||
return a, float(np.log(self._probs()[a] + 1e-12))
|
||
|
||
def action_logprob(self, action: int) -> float:
|
||
return float(np.log(self._probs()[int(action)] + 1e-12))
|
||
|
||
def ar_forward(self, tokens) -> int:
|
||
"""und-pathway hook so the declared ``ar_decode`` loop is runnable (greedy ⇒ interleave-safe)."""
|
||
return int(np.argmax(self.logits))
|
||
|
||
def refined_embed(self, base_embed, action: int) -> np.ndarray:
|
||
"""Apply the chosen refinement to the base text embedding → the diffusion conditioning."""
|
||
base = np.zeros((self.seq, self.dim), dtype="float32") if base_embed is None \
|
||
else np.asarray(base_embed, dtype="float32")
|
||
return (base + self._emb[int(action)]).astype("float32")
|
||
|
||
def pg_step(self, action: int, advantage: float, lr: float) -> float:
|
||
"""REINFORCE on the categorical policy. Returns grad_norm (the per-method regression signal)."""
|
||
p = self._probs()
|
||
onehot = np.zeros(self.n, dtype="float64")
|
||
onehot[int(action)] = 1.0
|
||
grad = -float(advantage) * (onehot - p) # ∇_logits (−A·logπ(a))
|
||
self.logits = self.logits - float(lr) * grad # descent on −A·logπ ⇒ ascent on A·logπ
|
||
return float(np.linalg.norm(grad))
|
||
|
||
def kl_to(self, other: "ToyPromptRefiner") -> float:
|
||
"""KL(π_self ‖ π_other) over the categorical — the LM-PG reference-KL penalty term."""
|
||
p, q = self._probs(), other._probs()
|
||
return float(np.sum(p * (np.log(p + 1e-12) - np.log(q + 1e-12))))
|
||
|
||
def clone(self) -> "ToyPromptRefiner":
|
||
c = ToyPromptRefiner.__new__(ToyPromptRefiner)
|
||
c.n, c.dim, c.seq = self.n, self.dim, self.seq
|
||
c.logits, c._emb = self.logits.copy(), self._emb.copy()
|
||
return c
|
||
|
||
def copy_from(self, other: "ToyPromptRefiner") -> None:
|
||
self.logits = other.logits.copy()
|
||
|
||
def blend_from(self, other: "ToyPromptRefiner", decay: float) -> None:
|
||
self.logits = decay * self.logits + (1.0 - decay) * other.logits
|
||
|
||
|
||
class ToyTalker(ToyMoTDiT):
|
||
"""Toy Talker — Qwen-Omni stage 1 (AR over a speech-codec vocab, conditioned on the Thinker).
|
||
|
||
A *separate expert* from the thinker (not weight-shared): its next-token depends on its OWN
|
||
weights (a seed salt) AND the prefilled thinker payload (tokens + hidden state). So the cascade
|
||
is real — change the thinker and the talker's tokens change; change the talker's weights and they
|
||
change too. The vllm-omni Talker is likewise a distinct model conditioned on Thinker hidden states.
|
||
"""
|
||
|
||
def __init__(self, channels: int = LATENT_CHANNELS, seed: int = 0):
|
||
super().__init__(channels=channels, seed=seed)
|
||
self.salt = (seed % 97) + 1
|
||
|
||
def ar_forward(self, tokens) -> int:
|
||
ctx = sum(int(t) for t in tokens)
|
||
return int((ctx * 7 + 13 + self.salt) % self.VOCAB) # weight-dependent (salt) + context
|
||
|
||
|
||
class ToyVocoder:
|
||
"""Toy streaming code2wav vocoder — Qwen-Omni stage 2 (BigVGAN/Code2Wav role).
|
||
|
||
Speech codec tokens → audio waveform, synthesized in chunks. Deterministic given the tokens so
|
||
the ``audio_decode`` loop is interleave-safe. Each token contributes a short, timbre-stamped
|
||
waveform segment — enough to exercise chunked synthesis, streaming emits, and the AudioArtifact.
|
||
"""
|
||
|
||
def __init__(self, samples_per_token: int = 16, vocab: int = 256, seed: int = 7):
|
||
self.spt = int(samples_per_token)
|
||
rng = np.random.default_rng(seed)
|
||
self.bank = (rng.standard_normal(vocab) * 0.3).astype("float32") # per-token timbre offset
|
||
|
||
def synthesize(self, token_chunk) -> np.ndarray:
|
||
"""One chunk of codec tokens → a waveform segment in [-1, 1], length spt·len(chunk)."""
|
||
segs = []
|
||
for t in token_chunk:
|
||
t = int(t) % self.bank.size
|
||
phase = np.linspace(0.0, np.pi * (1 + (t % 8)), self.spt, dtype="float32")
|
||
segs.append(np.tanh(self.bank[t] + np.sin(phase)).astype("float32"))
|
||
return np.concatenate(segs) if segs else np.zeros(0, dtype="float32")
|
||
|
||
|
||
class ToyLoRA:
|
||
"""Toy LoRA adapter — a low-rank velocity delta on a base DiT, selected and applied per request.
|
||
|
||
Many of these are served over ONE resident base (the adapter is tiny: a rank-r channel delta);
|
||
a request picks which to apply. Swappable/versioned independently (the cache key's
|
||
``adapter_versions``). Stand-in for a real LoRA / DoRA / IP-Adapter (design_v3 §9.19)."""
|
||
kind = "lora"
|
||
|
||
def __init__(self, adapter_id: str, channels: int = LATENT_CHANNELS, scale: float = 0.6,
|
||
rank: int = 2, seed: int = 0):
|
||
self.adapter_id = adapter_id
|
||
self.scale = float(scale)
|
||
rng = np.random.default_rng(seed)
|
||
self.a = (rng.standard_normal((channels, rank)) * 0.4).astype("float32")
|
||
self.b = (rng.standard_normal((rank, channels)) * 0.4).astype("float32")
|
||
|
||
def delta(self, latent, control=None) -> np.ndarray:
|
||
w = self.a @ self.b # [C, C] low-rank weight delta
|
||
d = np.tensordot(w, np.asarray(latent, dtype="float32"), axes=([1], [0])) # [C, ...]
|
||
return (self.scale * np.tanh(d)).astype("float32")
|
||
|
||
def update(self, seed: int) -> None: # hot-swap: replace the adapter's weights
|
||
rng = np.random.default_rng(seed)
|
||
self.a = (rng.standard_normal(self.a.shape) * 0.4).astype("float32")
|
||
|
||
|
||
class ToyControlNet:
|
||
"""Toy ControlNet adapter — conditions the velocity on a control signal (pose/depth/edge stand-in).
|
||
Different control ⇒ different generation; selected per request like a LoRA (design_v3 §9.19)."""
|
||
kind = "controlnet"
|
||
|
||
def __init__(self, adapter_id: str, channels: int = LATENT_CHANNELS, scale: float = 0.8, seed: int = 1):
|
||
self.adapter_id = adapter_id
|
||
self.scale = float(scale)
|
||
rng = np.random.default_rng(seed)
|
||
self.proj = (rng.standard_normal(channels) * 0.3).astype("float32")
|
||
|
||
def delta(self, latent, control=None) -> np.ndarray:
|
||
latent = np.asarray(latent, dtype="float32")
|
||
c = float(np.mean(control)) if control is not None else 0.0 # the control image's signal
|
||
shape = (latent.shape[0],) + (1,) * (latent.ndim - 1)
|
||
return (self.scale * c * self.proj.reshape(shape)).astype("float32")
|
||
|
||
def update(self, seed: int) -> None:
|
||
rng = np.random.default_rng(seed)
|
||
self.proj = (rng.standard_normal(self.proj.shape) * 0.3).astype("float32")
|
||
|
||
|
||
def _spec_target_next(tokens) -> int:
|
||
"""The target AR model's greedy next-token. Length-dependent so the sequence stays varied (no
|
||
trivial absorbing fixed point) — the substrate for a meaningful speculative-decoding accept rate."""
|
||
s = sum(int(t) for t in tokens)
|
||
n = len(tokens)
|
||
return int((s * 31 + 7 * n + 13) % 256)
|
||
|
||
|
||
class ToyTargetModel:
|
||
"""Toy speculative-decoding TARGET (the large, ground-truth AR model)."""
|
||
VOCAB = 256
|
||
EOS = 0
|
||
|
||
def __init__(self, seed: int = 0):
|
||
self.seed = seed
|
||
|
||
def ar_forward(self, tokens) -> int:
|
||
return _spec_target_next(tokens)
|
||
|
||
|
||
class ToyDraftModel:
|
||
"""Toy speculative-decoding DRAFT model — a cheap approximation of the target.
|
||
|
||
It replicates the target on a context-and-position-varying subset (~``agree`` fraction → those tokens
|
||
are accepted) and diverges otherwise (→ a rejection + a target correction). The accept rate varies
|
||
*within* a decode, so different rounds accept different lengths (a ragged loop). The demonstration is
|
||
exact regardless of draft quality: speculative decoding accepts only tokens the target would have
|
||
produced, so the output equals the target's own greedy decode (design_v3 §2.2)."""
|
||
VOCAB = 256
|
||
EOS = 0
|
||
|
||
def __init__(self, agree: float = 0.7, seed: int = 0):
|
||
self.agree = float(agree)
|
||
|
||
def ar_forward(self, tokens) -> int:
|
||
target = _spec_target_next(tokens)
|
||
s, n = sum(int(t) for t in tokens), len(tokens)
|
||
if ((s * 13 + n * 7) % 10) < int(round(self.agree * 10)): # agree on a varying subset
|
||
return int(target)
|
||
return int((target + 1) % self.VOCAB) # a cheap wrong guess (rejected)
|
||
|
||
|
||
class ToyRewardModel:
|
||
"""Toy learned reward model — a served scorer/verifier (PickScore / CLIP / a VLM judge stand-in).
|
||
|
||
Maps a media latent → a scalar reward in [-1, 1] via a fixed projection over per-channel features.
|
||
The point isn't the math: it's that the reward is computed by a *model* (served, scheduled as
|
||
``REWARD_BATCH`` work units, place-able on its own pool), not a numpy heuristic bolted onto the
|
||
method — the reward plane composing with the serving plane (design_v3 §10)."""
|
||
|
||
def __init__(self, dim: int = LATENT_CHANNELS, seed: int = 5):
|
||
rng = np.random.default_rng(seed)
|
||
self.w = (rng.standard_normal(dim) * 0.5).astype("float32")
|
||
|
||
def score(self, media) -> float:
|
||
a = np.asarray(media, dtype="float64")
|
||
feat = a.reshape(a.shape[0], -1).mean(axis=1) if a.ndim > 1 else np.atleast_1d(a)
|
||
n = self.w.size
|
||
feat = feat[:n] if feat.size >= n else np.pad(feat, (0, n - feat.size))
|
||
return float(np.tanh(float(np.dot(self.w, feat))))
|
||
|
||
|
||
class ToyAudioVAE:
|
||
"""Toy audio VAE — mel-like audio latent → mono waveform (the LTX-2 / Cosmos3 audio branch).
|
||
|
||
Decodes a small ``[C, A, 1, 1]`` audio latent (denoised jointly with the video latent) into a 1-D
|
||
waveform: channel-collapse via a fixed projection, then upsample along time. Deterministic, so the
|
||
joint A/V denoise loop stays interleave-safe. The video VAE and this share nothing — two decoders,
|
||
one synchronized latent pair (design_v3 §15b joint A/V)."""
|
||
|
||
def __init__(self, samples_per_frame: int = 16, seed: int = 2):
|
||
self.spf = int(samples_per_frame)
|
||
rng = np.random.default_rng(seed)
|
||
self.proj = (rng.standard_normal(LATENT_CHANNELS) * 0.3).astype("float32")
|
||
|
||
def decode(self, audio_latent: np.ndarray) -> np.ndarray:
|
||
a = np.asarray(audio_latent, dtype="float32")
|
||
flat = a.reshape(a.shape[0], -1) # [C, A·...]
|
||
mono = np.tensordot(self.proj, flat, axes=([0], [0])) # [A·...]
|
||
return np.tanh(np.repeat(mono, self.spf)).astype("float32") # [A·spf] mono waveform
|
||
|
||
|
||
class ToyVAE:
|
||
"""Tiny deterministic VAE. encode: video→latent (mean-pool + channel proj); decode: latent→video."""
|
||
|
||
def __init__(self, channels: int = LATENT_CHANNELS, spatial: int = 8, temporal: int = 4, seed: int = 1):
|
||
self.C = channels
|
||
self.spatial = spatial
|
||
self.temporal = temporal
|
||
rng = np.random.default_rng(seed)
|
||
self.dec_proj = (rng.standard_normal((3, channels)) * 0.2).astype("float32")
|
||
|
||
def decode(self, latent: np.ndarray) -> np.ndarray:
|
||
"""latent [C,T,H,W] -> video [3, T, H*spatial, W*spatial] (deterministic upsample)."""
|
||
latent = np.asarray(latent, dtype=np.float32)
|
||
rgb = np.tensordot(self.dec_proj, latent, axes=([1], [0])) # [3,T,H,W]
|
||
up = np.repeat(np.repeat(rgb, self.spatial, axis=2), self.spatial, axis=3)
|
||
return np.tanh(up).astype("float32")
|
||
|
||
def encode(self, video: np.ndarray) -> np.ndarray:
|
||
video = np.asarray(video, dtype=np.float32)
|
||
pooled = video[:self.C] # crude: take C channels
|
||
return pooled.astype("float32")
|