diff --git a/.env.example b/.env.example new file mode 100644 index 0000000000000000000000000000000000000000..7cc0425cc56194b9d6546b600c11c445da09a6b6 --- /dev/null +++ b/.env.example @@ -0,0 +1,6 @@ +# Optional cache-directory override for HuggingFace/torch downloads and run outputs. +# Cache resolution priority: +# 1. NEVO_CACHE_DIR below (if set) +# 2. the system/user-default HuggingFace cache (HF_HOME, else ~/.cache/huggingface) +# 3. /cache/ (only if no default is resolvable) +# NEVO_CACHE_DIR=/absolute/path/to/nevo_cache diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..3e37e39470f7475fbbe8378f056a156e4922c3da --- /dev/null +++ b/.gitattributes @@ -0,0 +1,41 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text +assets/gallery/EBA.gif filter=lfs diff=lfs merge=lfs -text +assets/gallery/FFA.gif filter=lfs diff=lfs merge=lfs -text +assets/gallery/MT.gif filter=lfs diff=lfs merge=lfs -text +assets/gallery/PPA.gif filter=lfs diff=lfs merge=lfs -text +assets/gallery/V1.gif filter=lfs diff=lfs merge=lfs -text +assets/gallery/pSTS.gif filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..2fd866726b0ca07f0859350c7f5244a585759e90 --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +__pycache__/ +*.py[cod] +.pytest_cache/ +.DS_Store +cache/ +.env diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..6c693537205fa345e680457fad41e346c47b0cb1 --- /dev/null +++ b/README.md @@ -0,0 +1,213 @@ +--- +tags: +- neuroscience +- fmri +- brain-decoding +- stimulus-synthesis +- v-jepa +- diffusers +library_name: diffusers +pipeline_tag: text-to-video +--- + +> 🚧 **Work in progress** β€” this model is still being transferred from its main development repository, so the model card and API are subject to change. + +# NEvo β€” Neural-Guided Evolutionary Video Synthesis + +**🌐 Project website: [nevo-project.epfl.ch](https://nevo-project.epfl.ch/) Β· πŸ“„ Paper: [arXiv:2607.02317](https://arxiv.org/abs/2607.02317)** + +NEvo is a self-contained Hugging Face custom [Diffusers](https://github.com/huggingface/diffusers) pipeline for **neural-response-guided visual stimulus synthesis**. Given a brain target (a set of voxels, or a target fMRI vector), it searches over prompts, generates images and short videos, scores each candidate with a differentiable image/videoβ†’fMRI encoder, and returns the ranked stimuli predicted to best drive that target. + +It orchestrates three frozen models. The models below are **placeholders / defaults** and can be swapped for any compatible models (weights are not bundled β€” they are pulled from their own repos): + +| Role | Default model | +|------|---------------| +| Encoder (image/video β†’ fMRI) | [`epfl-neuroai/vjepa2-encoder-basic`](https://huggingface.co/epfl-neuroai/vjepa2-encoder-basic) (`predict_fmri`) | +| Text β†’ image | [`stabilityai/sdxl-turbo`](https://huggingface.co/stabilityai/sdxl-turbo) | +| Image β†’ video | [`Lightricks/LTX-Video-0.9.8-13B-distilled`](https://huggingface.co/Lightricks/LTX-Video-0.9.8-13B-distilled) | + +## Gallery + +Each clip is from the **top results of a NEvo search targeting one visual region** β€” the model discovers, from scratch, stimuli that drive that region's known selectivity. + +| Region | Stimulus | Region | Stimulus | +|:------:|:--------:|:------:|:--------:| +| **FFA** Β· faces | ![FFA](assets/gallery/FFA.gif) | **PPA** Β· places | ![PPA](assets/gallery/PPA.gif) | +| **MT** Β· motion | ![MT](assets/gallery/MT.gif) | **EBA** Β· bodies | ![EBA](assets/gallery/EBA.gif) | +| **pSTS** Β· social motion | ![pSTS](assets/gallery/pSTS.gif) | **V1** Β· early visual | ![V1](assets/gallery/V1.gif) | + +Explore the full interactive gallery and 3D brain maps at **[nevo-project.epfl.ch](https://nevo-project.epfl.ch/)**. + +## Installation + +**Off-the-shelf β€” no install.** Load NEvo as a custom Diffusers pipeline; the package and its bundled data are fetched from the Hub automatically (you only need the usual dependencies below): + +```python +from diffusers import DiffusionPipeline + +pipe = DiffusionPipeline.from_pretrained( + "epfl-neuroai/NEvo", custom_pipeline="epfl-neuroai/NEvo", trust_remote_code=True, +) +``` + +**Or install the package** (for cleaner `from stimulus_synthesis import ...` imports / development): + +```bash +conda create -n nevo python=3.10 -y +conda activate nevo +pip install "git+https://huggingface.co/epfl-neuroai/NEvo" +# or from a local clone: +# git clone https://huggingface.co/epfl-neuroai/NEvo && pip install ./NEvo +# then: from stimulus_synthesis import NevoPipeline; pipe = NevoPipeline.from_pretrained("epfl-neuroai/NEvo") +``` + +Runtime dependencies (either way): `torch`, `diffusers`, `transformers`, `huggingface_hub`, `numpy`, `pillow`, `av` (`pytest` for tests). No `nilearn` / atlas downloads β€” ROI masks are shipped as small precomputed data files. + +## Quickstart + +Target a brain region by name β€” NEvo resolves its voxels and searches for a video predicted to drive it: + +```python +from diffusers import DiffusionPipeline + +# fetches the pipeline (and package) from the Hub; model weights are pulled on first use +pipe = DiffusionPipeline.from_pretrained( + "epfl-neuroai/NEvo", custom_pipeline="epfl-neuroai/NEvo", trust_remote_code=True, +) + +out = pipe(roi="FFA", progress=True) # omit seed (default) -> different result each run; pass seed= to reproduce (the seed used is in out.metadata["seed"]) +print(out.best_prompt, out.best_score) + +from stimulus_synthesis.media import save_video, video_to_t_c_h_w # importable once the pipeline has loaded +save_video(video_to_t_c_h_w(out.best.video), "best_stimulus.mp4") # save the synthesized video +out.best.image.save("best_stimulus.png") # and the stage-1 best image (PIL) +``` + +This runs the two-stage search with the defaults β€” up to 400 image evaluations then 200 video evaluations (population 20), using the fast distilled-model defaults (1-step 512Γ—512 SDXL-Turbo, 8-step 512Γ—512 LTX). A run takes a few minutes and a good amount of GPU memory. + +### Faster run + +For a quicker first result, shrink the search and the video: + +```python +out = pipe( + roi="FFA", + progress=True, + image_max_evals=80, # stage-1 (image) evaluation budget (default: 400) + video_max_evals=40, # stage-2 (video) evaluation budget (default: 200) + population_size=8, # GA population per generation (default: 20) + seed=0, # RNG seed, for reproducibility + video_kwargs={ # merged over the fast defaults (8 steps / 25 frames / 512Β²); override any key + "num_inference_steps": 8, # denoising steps β€” the distilled LTX model needs only a few + "num_frames": 25, # clip length; LTX requires 8*k + 1 frames + "height": 256, "width": 256, + }, +) +``` + +**Enhanced search space.** Selecting a region (`roi=...`) restricts the prompt search to the categories relevant to that region β€” a smaller space that converges faster. Pass `enforce_general_search_space=True` to search the full general space instead. + +Available ROI tokens (comma-separated tokens are unioned): + +- **Named ROIs:** `FFA`, `PPA`, `MT`, `EBA`, `LOC`, `RSC`, `pSTS`, `aSTS`, `V1`, `V2`, `V3`, `V4` β€” optionally hemisphere-suffixed (`FFA_lh`, `MT_rh`). +- **Searchlight regions:** `SL-` (both hemispheres), `SL-_lh`, `SL-_rh` (58 both / 28 lh / 30 rh). + +```python +from stimulus_synthesis.neuro import available_rois, searchlight_counts +available_rois() # ['EBA','FFA','LOC','MT','PPA','RSC','V1','V2','V3','V4','aSTS','pSTS'] +searchlight_counts() # {'both': 58, 'lh': 28, 'rh': 30} +``` + +> **fsaverage5 only.** The bundled ROI/searchlight masks are defined on the **fsaverage5** cortical surface (20 484 vertices). Targeting a region by name therefore requires an encoder whose `predict_fmri` output lives in that same space β€” the default `epfl-neuroai/vjepa2-encoder-basic`. A custom encoder with a different output space can still be driven with explicit `vector`/`indices` targets, but not with the named-ROI helper. + +### Custom targets + +Instead of a region name, pass raw voxel indices or a full target fMRI vector: + +```python +import numpy as np +from stimulus_synthesis import resolve_driving_voxels + +mask = resolve_driving_voxels("FFA") # boolean mask, length 20484 +out = pipe(target={"type": "indices", "indices": np.flatnonzero(mask).tolist()}) +``` + +### Target types & objectives + +| Target | Objective (default) | Meaning | +|--------|--------------------|---------| +| `{"type": "indices", "indices": [...]}` | `indices_mean` | mean predicted response over ROI voxels | +| `{"type": "vector", "vector": [...]}` (len 20484) | `target_vector_cosine` / `vector_dot` | match a full target fMRI vector | +| `{"type": "weights", "weights": [...]}` | `weighted_mean` | weighted voxel objective | + +## Search parameters (defaults) + +Set in `stimulus_synthesis_config.json`: + +| Param | Default | Notes | +|-------|---------|-------| +| `default_image_max_evals` | 400 | stage-1 (image) evaluation budget (GA `max_evals`) | +| `default_video_max_evals` | 200 | stage-2 (video) evaluation budget | +| `default_population_size` | 20 | GA population per generation (= `n_init`) | +| `default_score_frames` | 24 | number of frames the encoder scores (a still image is replicated to this) | +| `default_score_size` | 224 | resolution the clip is resized to for the encoder (call-time: `score_size=`) | +| `default_mutation_rate` | 0.25 | | +| `default_elite_frac` | 0.35 | | +| `default_objective` | `indices_mean` | | +| `default_score_transform` | disabled | robust augmentation off by default (clean single pass) | +| `default_image_kwargs` | `{num_inference_steps: 1, guidance_scale: 0, height: 512, width: 512}` | fast SDXL-Turbo settings (merged under call-time `image_kwargs`) | +| `default_video_kwargs` | `{num_inference_steps: 8, num_frames: 25, height: 512, width: 512}` | fast LTX settings (merged under call-time `video_kwargs`) | + +Each stage runs a genetic search with population `population_size` (default 20) until it hits its evaluation budget β€” `image_max_evals` (default 400) and `video_max_evals` (default 200) generateβ†’score passes. Image and video generation use fast distilled defaults out of the box (`default_image_kwargs` / `default_video_kwargs`); anything you pass as `image_kwargs` / `video_kwargs` is merged over them, so you only override the keys you care about. + +### Robust scoring + +By default each candidate is scored with a single clean encoder pass. An optional **robust mode** β€” the mean over 4 augmented draws (random crop `0.8`, Gaussian `Οƒ=0.1`) via `RobustTransformScorer` β€” reduces sensitivity to encoder artifacts; turn it on by setting `"enabled": true` in `default_score_transform`. + +## Cache configuration + +Model weights and outputs cache location resolves in priority order: + +1. `NEvo_CACHE_DIR` β€” set it in a repo-root `.env` file (see `.env.example`) or the environment. +2. Otherwise the **system/user-default HuggingFace cache** (`HF_HOME`, else `~/.cache/huggingface`) is used and left untouched. +3. Only if no default is resolvable, a repo-local `cache/` is used. + +`cache/` and `.env` are git-ignored. + +## Batch runners + +Two ROI-driven, two-stage (image-search β†’ video-search) runners are included: + +- **`run_roi_samples.py`** β€” genetic search per ROI/seed, scoring in-memory tensors; writes `best_image.png` / `best_video.mp4` / scores. +- **`run_regional_asset_pilot.py`** β€” same search but exports every candidate to a deterministically-encoded file (PNG/MP4), hashes it (sha256), and scores the *decoded file* β€” producing provenance-tracked, reproducible published assets with manifests. + +Both take `--rois`, `--seeds`, `--image-evals` / `--video-evals`, `--encoder-model`, `--out-dir`, etc., and default to the config's encoder and a cache-relative output directory. + +## Reproducibility + +The pipeline is deterministic for a fixed seed/config: the shipped ROI masks reproduce the original atlas masks bit-for-bit, and a fixed-seed run reproduces prior scores exactly. Encoder scores are a *target-matching* signal, not ground-truth reconstruction quality. + +## Intended use & limitations + +- **Research use** in visual neuroscience / brain-decoding. Outputs are *predicted* to drive a target region under a specific encoder β€” they are hypotheses to validate, not ground truth. +- Optimizing hard against a single encoder can exploit encoder artifacts; inspect images and use held-out validation. +- Requires a CUDA GPU with enough memory for the 13B video model; you must accept the license/access terms of the referenced upstream models. + +## Citation + +If you use NEvo, please cite: + +```bibtex +@article{tang2026nevo, + title={NEvo: Neural-Guided Evolutionary Video Synthesis for Dynamic Visual Selectivity}, + author={Tang, Yingtian and Salehi, Sogand and Zhou, Ming and Zamir, Amir and Isik, Leyla and Schrimpf, Martin}, + journal={arXiv preprint arXiv:2607.02317}, + year={2026} +} +``` + +Project website: [nevo-project.epfl.ch](https://nevo-project.epfl.ch/) + +## Acknowledgements + +Builds on BrainDiVE-style encoder-guided synthesis, vJEPA-2, SDXL-Turbo, and LTX-Video. ROI/searchlight definitions derive from an fsaverage-space group atlas (precomputed and bundled). diff --git a/assets/gallery/EBA.gif b/assets/gallery/EBA.gif new file mode 100644 index 0000000000000000000000000000000000000000..e8624187df7869f2bd8bdff888622c6e151f6fc8 --- /dev/null +++ b/assets/gallery/EBA.gif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:683bbc7d810684fbb404d40e33138b1bd006337604e4d3fe099ce2514fb1e13f +size 1005015 diff --git a/assets/gallery/FFA.gif b/assets/gallery/FFA.gif new file mode 100644 index 0000000000000000000000000000000000000000..fe2a7ac059a2bbe3c6eafe9cd9cdd7a9cd63c31e --- /dev/null +++ b/assets/gallery/FFA.gif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4fed64effefc8fc7f344f93e6b4d9f1e49a1cd7907d6844d9f4fff43eeacf4ab +size 845046 diff --git a/assets/gallery/MT.gif b/assets/gallery/MT.gif new file mode 100644 index 0000000000000000000000000000000000000000..1b88a60d1c0af076d8ee390428ceee89d901304e --- /dev/null +++ b/assets/gallery/MT.gif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5089ec8bb2513a5601896b585c9b10d96502b8ecb392fe4b9fd43889d37ecd0b +size 1214727 diff --git a/assets/gallery/PPA.gif b/assets/gallery/PPA.gif new file mode 100644 index 0000000000000000000000000000000000000000..a10d44604c993197cb78e169c1974d9a92ad417e --- /dev/null +++ b/assets/gallery/PPA.gif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3e0be6ceec22838fd2ade0f0c505f2136e87b70c9a90235dcb0a43eb7c5b51a1 +size 1071942 diff --git a/assets/gallery/V1.gif b/assets/gallery/V1.gif new file mode 100644 index 0000000000000000000000000000000000000000..667092ada4a2e40062394758b25f3f281d086e15 --- /dev/null +++ b/assets/gallery/V1.gif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d4ceee8ad5903bb98b602a231f5bd2b8c5fc1bf6cbb36998a39a3ae2b5be2d13 +size 1077496 diff --git a/assets/gallery/pSTS.gif b/assets/gallery/pSTS.gif new file mode 100644 index 0000000000000000000000000000000000000000..dbf537415bab7baae8134abb9cabaa2ba788c31b --- /dev/null +++ b/assets/gallery/pSTS.gif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:17f5f6992c66447468f896d6d5ba6f7334e67acbe8777163d27f53e9e44a4746 +size 830935 diff --git a/examples/override_models.py b/examples/override_models.py new file mode 100644 index 0000000000000000000000000000000000000000..677c0489a4643dd067f8f5c4e2f1e8ca0e9ff8f0 --- /dev/null +++ b/examples/override_models.py @@ -0,0 +1,13 @@ +from stimulus_synthesis import NevoPipeline + +pipe = NevoPipeline.from_pretrained( + "epfl-neuroai/NEvo", + synthesis_config={ + "text_to_image_model_id": "stabilityai/sdxl-turbo", + "image_to_video_model_id": "Lightricks/LTX-Video-0.9.8-13B-distilled", + "encoder_model_id": "epfl-neuroai/vjepa2-encoder-basic", + }, +) + +out = pipe(target={"type": "indices", "indices": [0, 1, 2]}, seed_prompts=["a moving person"]) +print(out.best_prompt, out.best_score) diff --git a/examples/quickstart.py b/examples/quickstart.py new file mode 100644 index 0000000000000000000000000000000000000000..368532c84b7acdab172520517618ffa867f050de --- /dev/null +++ b/examples/quickstart.py @@ -0,0 +1,11 @@ +from stimulus_synthesis import NevoPipeline + +pipe = NevoPipeline.from_pretrained("epfl-neuroai/NEvo") + +out = pipe( + target={"type": "indices", "indices": [0, 1, 2]}, + seed_prompts=["a person running through a crowded street"], + num_candidates=16, + num_rounds=50, +) +print(out.best_prompt, out.best_score) diff --git a/examples/synthesize_for_roi.py b/examples/synthesize_for_roi.py new file mode 100644 index 0000000000000000000000000000000000000000..e10d8450d4c73405f0018f1f563ca1bbc062fe44 --- /dev/null +++ b/examples/synthesize_for_roi.py @@ -0,0 +1,11 @@ +import numpy as np +from stimulus_synthesis import NevoPipeline, resolve_driving_voxels + +pipe = NevoPipeline.from_pretrained("epfl-neuroai/NEvo") + +mask = resolve_driving_voxels("FFA") # boolean mask, length 20484 +out = pipe( + target={"type": "indices", "indices": np.flatnonzero(mask).tolist()}, + seed_prompts=["a close-up of a person's face"], +) +print(out.best_prompt, out.best_score) diff --git a/examples/synthesize_for_vector.py b/examples/synthesize_for_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..d557b4ba42824fc2e8fd36359bf615115f3d14f4 --- /dev/null +++ b/examples/synthesize_for_vector.py @@ -0,0 +1,9 @@ +import torch +from stimulus_synthesis import NevoPipeline + +pipe = NevoPipeline.from_pretrained("epfl-neuroai/NEvo") + +target_vector = torch.zeros(20484) +target_vector[:10] = 1.0 +out = pipe(target={"type": "vector", "vector": target_vector.tolist()}, seed_prompts=["natural action video"]) +print(out.best_prompt, out.best_score) diff --git a/model_index.json b/model_index.json new file mode 100644 index 0000000000000000000000000000000000000000..acf815fc8b0a0136ea9cdcd37d6151bdd708d5bd --- /dev/null +++ b/model_index.json @@ -0,0 +1,5 @@ +{ + "_class_name": "NevoPipeline", + "_diffusers_version": "0.30.0", + "_module": "pipeline" +} diff --git a/pipeline.py b/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..6b0b91ed59dc915a1d811943dbb23def70064bb1 --- /dev/null +++ b/pipeline.py @@ -0,0 +1,20 @@ +"""NEvo custom Diffusers pipeline entry point. + +Works both when the ``stimulus_synthesis`` package is installed (``pip install``) and, +off-the-shelf, when the pipeline is loaded via +``DiffusionPipeline.from_pretrained(repo, custom_pipeline=repo, trust_remote_code=True)`` +without installing anything β€” in that case the package (and its bundled data) is +fetched from the Hub and put on the import path. +""" +try: + from stimulus_synthesis.pipeline import NevoPipeline +except ModuleNotFoundError: + import sys + from huggingface_hub import snapshot_download + + _pkg_root = snapshot_download("epfl-neuroai/NEvo", allow_patterns=["stimulus_synthesis/**"]) + if _pkg_root not in sys.path: + sys.path.insert(0, _pkg_root) + from stimulus_synthesis.pipeline import NevoPipeline + +__all__ = ["NevoPipeline"] diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..e447563dae5eb99eae1394944d68e66db5900a12 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,32 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "nevo-stimulus-synthesis" +version = "0.1.0" +description = "NEVO: neural-response-guided visual stimulus synthesis (custom Diffusers pipeline)." +readme = "README.md" +requires-python = ">=3.10" +dependencies = [ + "torch", + "diffusers", + "transformers", + "huggingface_hub", + "numpy", + "pillow", + "av", +] + +[project.optional-dependencies] +test = ["pytest"] + +[tool.setuptools.packages.find] +include = ["stimulus_synthesis*"] + +[tool.setuptools.package-data] +"stimulus_synthesis.data" = ["*.npz"] + +[tool.pytest.ini_options] +pythonpath = ["."] +testpaths = ["tests"] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..ec69dc94e087b27455b6839eccdd650b0dc242f5 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,17 @@ +# NEVO stimulus-synthesis β€” minimal runtime dependencies +torch +diffusers +transformers +huggingface_hub +numpy +pillow +av +tqdm +tiktoken +sentencepiece +protobuf +timm +einops + +# test-only +pytest diff --git a/run_regional_asset_pilot.py b/run_regional_asset_pilot.py new file mode 100644 index 0000000000000000000000000000000000000000..c7fec5d9f552a0a880a74ffe2ed6f0384f17f6d3 --- /dev/null +++ b/run_regional_asset_pilot.py @@ -0,0 +1,271 @@ +from __future__ import annotations + +import argparse +import gc +import json +import shutil +import sys +import time +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from PIL import Image + +_REPO = Path(__file__).resolve().parent +if str(_REPO) not in sys.path: + sys.path.insert(0, str(_REPO)) + +from stimulus_synthesis.spaces import StructuredArtPromptSpace, VideoMotionPromptSpace, make_t2v_art_data +from stimulus_synthesis.neuro import resolve_driving_voxels +from stimulus_synthesis.config import StimulusSynthesisConfig +from stimulus_synthesis.paths import get_cache_dir +from stimulus_synthesis.asset_manifest import write_asset_manifest +from stimulus_synthesis.generators.diffusers_i2v import DiffusersImageToVideoAdapter +from stimulus_synthesis.generators.diffusers_t2i import DiffusersTextToImageAdapter +from stimulus_synthesis.media import AssetExportRecord, ImageAssetSpec, VideoAssetSpec, sha256_file +from stimulus_synthesis.scoring import AssetScorer, EncoderPreprocessSpec +from stimulus_synthesis.scoring.encoder_scorer import EncoderScorer +from stimulus_synthesis.search.genetic import GeneticSearch +from stimulus_synthesis.spaces import SeededSearchSpace + + +class StaticImageToVideo: + def generate(self, image: Any, prompt: str, **kwargs) -> Any: + return image + + def generate_batch(self, images: list[Any], prompts: list[str], **kwargs) -> list[Any]: + return list(images) + + +class FixedImageT2I: + def __init__(self, image: Any): + self.image = image + + def generate(self, prompts: list[str], **kwargs) -> list[Any]: + return [self.image for _ in prompts] + + +def seed_values(run_seed: int, count: int, *, stream: int) -> list[int]: + rng = np.random.default_rng(int(run_seed) + 1_000_003 * int(stream)) + return [int(x) for x in rng.integers(0, 2**31 - 1, size=int(count), dtype=np.int64)] + + +def copy_exported_asset(record: AssetExportRecord, path: Path, spec: ImageAssetSpec | VideoAssetSpec) -> AssetExportRecord: + path.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(record.path, path) + return AssetExportRecord( + path=str(path), + asset_type=record.asset_type, + sha256=sha256_file(path), + bytes=path.stat().st_size, + spec=spec.to_dict(), + ) + + +def run_one(roi: str, seed: int, args, t2i_base, i2v_base, scorer) -> dict[str, Any]: + seed_dir = Path(args.out_dir) / roi / f'seed_{seed:06d}' + seed_dir.mkdir(parents=True, exist_ok=True) + result_path = seed_dir / 'result.json' + if result_path.exists() and not args.overwrite: + print(f'[skip] {roi} seed={seed}: {result_path} exists', flush=True) + return json.loads(result_path.read_text()) + + voxels = resolve_driving_voxels(roi) + indices = np.flatnonzero(voxels).astype(int).tolist() + target = {'type': 'indices', 'indices': indices} + print(f'[start] roi={roi} seed={seed} voxels={len(indices)}', flush=True) + + image_kwargs = { + 'height': args.image_height, + 'width': args.image_width, + 'num_inference_steps': 1, + 'guidance_scale': 0.0, + } + image_spec = ImageAssetSpec(width=args.image_width, height=args.image_height, format='png') + video_spec = VideoAssetSpec(width=args.video_width, height=args.video_height, fps=args.fps, num_frames=args.video_frames, crf=args.video_crf) + asset_scorer = AssetScorer(scorer, target, preprocess_spec=EncoderPreprocessSpec(size=args.score_size, num_frames=args.score_frames)) + + image_space = SeededSearchSpace( + StructuredArtPromptSpace(art_data=make_t2v_art_data(), roi=roi, option_embeddings=None), + seed_values(seed, args.seed_gene_count, stream=0), + ) + image_search = GeneticSearch( + max_evals=args.image_evals, + population_size=args.image_population, + n_init=args.image_population, + mutation_rate=args.mutation_rate, + crossover_rate=args.crossover_rate, + elite_frac=args.elite_frac, + image_kwargs=image_kwargs, + video_kwargs={}, + score_kwargs={}, + video_size=args.score_size, + num_frames=args.score_frames, + asset_scorer=asset_scorer, + asset_dir=seed_dir / 'candidate_images', + asset_type='image', + image_asset_spec=image_spec, + ) + t0 = time.time() + image_result = image_search.run(image_space, t2i_base, StaticImageToVideo(), scorer, target, seed=seed) + image_record = copy_exported_asset(image_result.best_export_record, seed_dir / 'best_image.png', image_spec) + final_image_score = asset_scorer.score_image( + image_record.path, + asset_spec=image_spec, + metadata={'roi': roi, 'run_seed': seed, 'generation_seed': image_result.best_seed, 'prompt': image_result.best_prompt}, + ) + best_image = Image.open(image_record.path).convert('RGB') + np.save(seed_dir / 'image_history_best.npy', np.asarray(image_result.history_best, dtype=np.float32)) + print(f'[image done] roi={roi} seed={seed} asset_score={image_result.best_score:.6f} seconds={time.time()-t0:.1f}', flush=True) + + video_kwargs = { + 'height': args.video_height, + 'width': args.video_width, + 'num_frames': args.video_frames, + 'frame_rate': args.fps, + 'num_inference_steps': args.video_steps, + 'guidance_scale': args.video_guidance_scale, + 'output_type': 'np', + } + video_space = SeededSearchSpace( + VideoMotionPromptSpace(roi=roi, option_embeddings=None), + seed_values(seed, args.seed_gene_count, stream=1), + ) + video_search = GeneticSearch( + max_evals=args.video_evals, + population_size=args.video_population, + n_init=min(args.video_population, args.video_evals), + mutation_rate=args.mutation_rate, + crossover_rate=args.crossover_rate, + elite_frac=args.elite_frac, + image_kwargs={}, + video_kwargs=video_kwargs, + score_kwargs={}, + video_size=args.score_size, + num_frames=args.score_frames, + asset_scorer=asset_scorer, + asset_dir=seed_dir / 'candidate_videos', + asset_type='video', + video_asset_spec=video_spec, + ) + t1 = time.time() + video_result = video_search.run(video_space, FixedImageT2I(best_image), i2v_base, scorer, target, seed=seed) + video_record = copy_exported_asset(video_result.best_export_record, seed_dir / 'best_video.mp4', video_spec) + final_video_score = asset_scorer.score_video( + video_record.path, + asset_spec=video_spec, + metadata={'roi': roi, 'run_seed': seed, 'generation_seed': video_result.best_seed, 'prompt': video_result.best_prompt}, + ) + np.save(seed_dir / 'video_history_best.npy', np.asarray(video_result.history_best, dtype=np.float32)) + print(f'[video done] roi={roi} seed={seed} asset_score={video_result.best_score:.6f} seconds={time.time()-t1:.1f}', flush=True) + + manifest_path = seed_dir / 'asset_manifest.json' + write_asset_manifest([image_record, video_record, final_image_score, final_video_score], manifest_path, metadata={ + 'roi': roi, + 'seed': seed, + 'num_voxels': len(indices), + 'text_to_image_model_id': args.text_to_image_model, + 'image_to_video_model_id': args.image_to_video_model, + 'encoder_model_id': args.encoder_model, + 'score_size': args.score_size, + 'score_frames': args.score_frames, + }) + + meta = { + 'roi': roi, + 'seed': seed, + 'num_voxels': len(indices), + 'image': { + 'max_evals': args.image_evals, + 'best_prompt': image_result.best_prompt, + 'optimization_score': image_result.best_score, + 'final_asset_score': final_image_score.score, + 'best_image': image_record.path, + 'sha256': image_record.sha256, + 'generation_seed': image_result.best_seed, + 'candidate_key': image_result.best_key, + 'score_source': image_result.best_metadata.get('score_source'), + }, + 'video': { + 'max_evals': args.video_evals, + 'best_prompt': video_result.best_prompt, + 'optimization_score': video_result.best_score, + 'final_asset_score': final_video_score.score, + 'best_video': video_record.path, + 'sha256': video_record.sha256, + 'generation_seed': video_result.best_seed, + 'candidate_key': video_result.best_key, + 'score_source': video_result.best_metadata.get('score_source'), + 'sampled_frame_indices': final_video_score.sampled_frame_indices, + }, + 'params': { + 'image_kwargs': image_kwargs, + 'video_kwargs': {k: v for k, v in video_kwargs.items() if k != 'output_type'}, + 'video_crf': args.video_crf, + 'score_size': args.score_size, + 'score_frames': args.score_frames, + }, + 'asset_manifest': str(manifest_path), + } + result_path.write_text(json.dumps(meta, indent=2)) + print(f'[final asset] roi={roi} seed={seed} image={final_image_score.score:.6f} video={final_video_score.score:.6f}', flush=True) + return meta + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument('--rois', nargs='+', default=['FFA', 'PPA', 'pSTS', 'MT']) + p.add_argument('--seeds', nargs='+', type=int, default=[101]) + p.add_argument('--out-dir', default=str(get_cache_dir() / 'results' / 'hf_nevo_regional_asset_pilot')) + p.add_argument('--overwrite', action='store_true') + p.add_argument('--device', default='cuda') + p.add_argument('--text-to-image-model', default='stabilityai/sdxl-turbo') + p.add_argument('--image-to-video-model', default='Lightricks/LTX-Video-0.9.8-13B-distilled') + p.add_argument('--encoder-model', default=StimulusSynthesisConfig().encoder_model_id) + p.add_argument('--image-evals', type=int, default=24) + p.add_argument('--video-evals', type=int, default=8) + p.add_argument('--image-population', type=int, default=8) + p.add_argument('--video-population', type=int, default=4) + p.add_argument('--mutation-rate', type=float, default=0.2) + p.add_argument('--crossover-rate', type=float, default=0.5) + p.add_argument('--elite-frac', type=float, default=0.3) + p.add_argument('--image-width', type=int, default=256) + p.add_argument('--image-height', type=int, default=256) + p.add_argument('--video-width', type=int, default=256) + p.add_argument('--video-height', type=int, default=256) + p.add_argument('--video-frames', type=int, default=17) + p.add_argument('--video-steps', type=int, default=4) + p.add_argument('--video-guidance-scale', type=float, default=1.0) + p.add_argument('--video-crf', type=int, default=10) + p.add_argument('--fps', type=int, default=24) + p.add_argument('--score-size', type=int, default=224) + p.add_argument('--score-frames', type=int, default=16) + p.add_argument('--seed-gene-count', type=int, default=64) + args = p.parse_args() + + device = args.device if torch.cuda.is_available() else 'cpu' + args.device = device + Path(args.out_dir).mkdir(parents=True, exist_ok=True) + print(json.dumps({'stage': 'load_components', 'device': device, 'out_dir': args.out_dir}), flush=True) + t0 = time.time() + t2i_base = DiffusersTextToImageAdapter(args.text_to_image_model, device=device) + i2v_base = DiffusersImageToVideoAdapter(args.image_to_video_model, device=device) + scorer = EncoderScorer(args.encoder_model, encoder_call='predict_fmri', objective='indices_mean', device=device) + print(json.dumps({'stage': 'components_loaded', 'seconds': round(time.time() - t0, 1)}), flush=True) + + all_meta = [] + for roi in args.rois: + for seed in args.seeds: + all_meta.append(run_one(roi, seed, args, t2i_base, i2v_base, scorer)) + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + summary_path = Path(args.out_dir) / 'summary.json' + summary_path.write_text(json.dumps(all_meta, indent=2)) + print(json.dumps({'stage': 'done', 'summary': str(summary_path), 'runs': len(all_meta)}), flush=True) + + +if __name__ == '__main__': + main() diff --git a/run_roi_samples.py b/run_roi_samples.py new file mode 100644 index 0000000000000000000000000000000000000000..aab2fe2b184e58d0466ab7bf9e1e6fd6a7132b60 --- /dev/null +++ b/run_roi_samples.py @@ -0,0 +1,273 @@ +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from PIL import Image + +_REPO = Path(__file__).resolve().parent +if str(_REPO) not in sys.path: + sys.path.insert(0, str(_REPO)) + +from stimulus_synthesis.spaces import ( # noqa: E402 + StructuredArtPromptSpace, + VideoMotionPromptSpace, + make_t2v_art_data, +) +from stimulus_synthesis.neuro import resolve_driving_voxels # noqa: E402 +from stimulus_synthesis.generators.diffusers_t2i import DiffusersTextToImageAdapter # noqa: E402 +from stimulus_synthesis.generators.diffusers_i2v import DiffusersImageToVideoAdapter # noqa: E402 +from stimulus_synthesis.media.normalize import video_to_t_c_h_w # noqa: E402 +from stimulus_synthesis.media.video_io import save_video # noqa: E402 +from stimulus_synthesis.scoring.encoder_scorer import EncoderScorer # noqa: E402 +from stimulus_synthesis.search.genetic import GeneticSearch # noqa: E402 +from stimulus_synthesis.config import StimulusSynthesisConfig # noqa: E402 +from stimulus_synthesis.paths import get_cache_dir # noqa: E402 + + +def _vkw(args): + kw = {"height": args.video_height, "width": args.video_width, "num_frames": args.video_frames} + if getattr(args, "video_steps", 0) and args.video_steps > 0: + kw["num_inference_steps"] = int(args.video_steps) + return kw + + +class SeededTextToImage: + def __init__(self, inner: DiffusersTextToImageAdapter, seed: int): + self.inner = inner + self.seed = int(seed) + + def generate(self, prompts: list[str], **kwargs) -> list[Any]: + torch.manual_seed(self.seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(self.seed) + return self.inner.generate(prompts, **kwargs) + + +class StaticImageToVideo: + def generate(self, image: Image.Image, prompt: str, **kwargs) -> Image.Image: + return image + + def generate_batch(self, images: list[Any], prompts: list[str], **kwargs) -> list[Any]: + return list(images) + + +class SeededImageToVideo: + def __init__(self, inner: DiffusersImageToVideoAdapter, seed: int): + self.inner = inner + self.seed = int(seed) + self.counter = 0 + + def generate(self, image: Any, prompt: str, **kwargs) -> Any: + kwargs.pop("generator", None) # override any caller-supplied generator with the seeded one + seed = self.seed + self.counter + self.counter += 1 + generator = None + if torch.cuda.is_available(): + generator = torch.Generator(device="cuda").manual_seed(seed) + else: + generator = torch.Generator().manual_seed(seed) + return self.inner.generate(image, prompt, generator=generator, **kwargs) + + def generate_batch(self, images: list[Any], prompts: list[str], **kwargs) -> list[Any]: + return [self.generate(image, prompt, **kwargs) for image, prompt in zip(images, prompts)] + + +def save_image(image: Any, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + if isinstance(image, Image.Image): + image.save(path) + return + if torch.is_tensor(image): + x = image.detach().cpu().float().clamp(0, 1) + if x.ndim == 4: + x = x[0] + if x.ndim == 3 and x.shape[0] in (1, 3): + arr = (x.permute(1, 2, 0).numpy() * 255).astype(np.uint8) + Image.fromarray(arr).save(path) + return + raise TypeError(f"Unsupported image type for saving: {type(image)!r}") + + +def save_any_video(video: Any, path: Path, fps: int = 24) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tensor = video_to_t_c_h_w(video).clamp(0, 1) + save_video(tensor, str(path), fps=fps) + + +def best_artifacts_from_search(search: GeneticSearch, space, t2i, i2v, scorer, target, seed: int, image_kwargs, video_kwargs, score_kwargs): + result = search.run(space, t2i, i2v, scorer, target, seed=seed) + image = t2i.generate([result.best_prompt], **image_kwargs)[0] + video = i2v.generate(image, result.best_prompt, **video_kwargs) + return result, image, video + + +def run_one(roi: str, seed: int, args, t2i_base, i2v_base, scorer) -> dict: + seed_dir = Path(args.out_dir) / roi / f"seed_{seed:06d}" + seed_dir.mkdir(parents=True, exist_ok=True) + done = seed_dir / "result.json" + if done.exists() and not args.overwrite: + print(f"[skip] {roi} seed={seed}: {done} exists", flush=True) + return json.loads(done.read_text()) + + voxels = resolve_driving_voxels(roi) + indices = np.flatnonzero(voxels).astype(int).tolist() + target = {"type": "indices", "indices": indices} + print(f"[start] {roi} seed={seed} voxels={len(indices)}", flush=True) + + image_space = StructuredArtPromptSpace(art_data=make_t2v_art_data(), roi=roi, option_embeddings=None) + image_search = GeneticSearch( + max_evals=args.image_evals, + population_size=args.image_population, + n_init=args.image_population, + mutation_rate=args.mutation_rate, + crossover_rate=args.crossover_rate, + elite_frac=args.elite_frac, + image_kwargs={"num_inference_steps": 1, "guidance_scale": 0.0}, + video_kwargs={}, + score_kwargs={}, + video_size=args.score_size, + num_frames=args.score_frames, + ) + + t2i = SeededTextToImage(t2i_base, seed) + image_result, best_image, _static_video = best_artifacts_from_search( + image_search, + image_space, + t2i, + StaticImageToVideo(), + scorer, + target, + seed, + {"num_inference_steps": 1, "guidance_scale": 0.0}, + {}, + {}, + ) + best_image_path = seed_dir / "best_image.png" + save_image(best_image, best_image_path) + np.save(seed_dir / "image_history_best.npy", np.asarray(image_result.history_best, dtype=np.float32)) + (seed_dir / "image_result.json").write_text(json.dumps({ + "roi": roi, + "seed": seed, + "num_voxels": len(indices), + "best_prompt": image_result.best_prompt, + "best_score": image_result.best_score, + "best_image": str(best_image_path), + }, indent=2)) + print(f"[image done] {roi} seed={seed} score={image_result.best_score:.6f}", flush=True) + + video_space = VideoMotionPromptSpace(roi=roi, option_embeddings=None) + video_search = GeneticSearch( + max_evals=args.video_evals, + population_size=args.video_population, + n_init=min(args.video_population, args.video_evals), + mutation_rate=args.mutation_rate, + crossover_rate=args.crossover_rate, + elite_frac=args.elite_frac, + image_kwargs={"num_inference_steps": 1, "guidance_scale": 0.0}, + video_kwargs=_vkw(args), + score_kwargs={}, + video_size=args.score_size, + num_frames=args.score_frames, + ) + + class FixedImageT2I: + def generate(self, prompts: list[str], **kwargs) -> list[Any]: + return [best_image for _ in prompts] + + i2v = SeededImageToVideo(i2v_base, seed) + video_result, _image, best_video = best_artifacts_from_search( + video_search, + video_space, + FixedImageT2I(), + i2v, + scorer, + target, + seed, + {}, + _vkw(args), + {}, + ) + best_video_path = seed_dir / "best_video.mp4" + save_any_video(best_video, best_video_path, fps=args.fps) + np.save(seed_dir / "video_history_best.npy", np.asarray(video_result.history_best, dtype=np.float32)) + + meta = { + "roi": roi, + "seed": seed, + "num_voxels": len(indices), + "image": { + "max_evals": args.image_evals, + "best_prompt": image_result.best_prompt, + "best_score": image_result.best_score, + "best_image": str(best_image_path), + }, + "video": { + "max_evals": args.video_evals, + "best_prompt": video_result.best_prompt, + "best_score": video_result.best_score, + "best_video": str(best_video_path), + }, + } + done.write_text(json.dumps(meta, indent=2)) + print(f"[video done] {roi} seed={seed} score={video_result.best_score:.6f}", flush=True) + return meta + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--rois", nargs="+", default=["FFA", "PPA", "pSTS", "MT"]) + p.add_argument("--seeds", nargs="+", type=int, default=[33, 34, 35]) + p.add_argument("--image-evals", type=int, default=StimulusSynthesisConfig().default_image_max_evals) + p.add_argument("--video-evals", type=int, default=StimulusSynthesisConfig().default_video_max_evals) + p.add_argument("--image-population", type=int, default=20) + p.add_argument("--video-population", type=int, default=20) + p.add_argument("--encoder-model", default=StimulusSynthesisConfig().encoder_model_id) + p.add_argument("--mutation-rate", type=float, default=0.2) + p.add_argument("--crossover-rate", type=float, default=0.5) + p.add_argument("--elite-frac", type=float, default=0.3) + p.add_argument("--out-dir", default=str(get_cache_dir() / "results" / "hf_nevo_roi_samples")) + p.add_argument("--device", default="cuda") + p.add_argument("--score-size", type=int, default=224) + p.add_argument("--score-frames", type=int, default=16) + p.add_argument("--video-width", type=int, default=256) + p.add_argument("--video-height", type=int, default=256) + p.add_argument("--video-frames", type=int, default=49) + p.add_argument("--video-steps", type=int, default=0, help="LTX num_inference_steps; 0 = model default") + p.add_argument("--fps", type=int, default=24) + p.add_argument("--overwrite", action="store_true") + args = p.parse_args() + + device = args.device if torch.cuda.is_available() else "cpu" + Path(args.out_dir).mkdir(parents=True, exist_ok=True) + print(f"device={device} out_dir={args.out_dir}", flush=True) + + t0 = time.time() + t2i_base = DiffusersTextToImageAdapter("stabilityai/sdxl-turbo", device=device) + i2v_base = DiffusersImageToVideoAdapter("Lightricks/LTX-Video-0.9.8-13B-distilled", device=device) + scorer = EncoderScorer( + args.encoder_model, + encoder_call="predict_fmri", + objective="indices_mean", + device=device, + ) + print(f"components loaded in {time.time() - t0:.1f}s", flush=True) + + all_meta = [] + for roi in args.rois: + for seed in args.seeds: + all_meta.append(run_one(roi, seed, args, t2i_base, i2v_base, scorer)) + summary_path = Path(args.out_dir) / "summary.json" + summary_path.write_text(json.dumps(all_meta, indent=2)) + print(f"[done] wrote {summary_path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/stimulus_synthesis/__init__.py b/stimulus_synthesis/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0071ead3c7db2e22016ecdc28f4c4d62771f1368 --- /dev/null +++ b/stimulus_synthesis/__init__.py @@ -0,0 +1,8 @@ +from .paths import configure_cache as _configure_cache +_configure_cache() + +from .pipeline import NevoPipeline +from .outputs import StimulusSynthesisOutput +from .neuro import resolve_driving_voxels + +__all__ = ["NevoPipeline", "StimulusSynthesisOutput", "resolve_driving_voxels"] diff --git a/stimulus_synthesis/asset_manifest.py b/stimulus_synthesis/asset_manifest.py new file mode 100644 index 0000000000000000000000000000000000000000..6dc203825bf99e42647662bccfb5b989bb038369 --- /dev/null +++ b/stimulus_synthesis/asset_manifest.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import json +from dataclasses import asdict, is_dataclass +from pathlib import Path +from typing import Any + + +def write_asset_manifest(records: list[Any], path: str | Path, *, metadata: dict[str, Any] | None = None) -> dict[str, Any]: + manifest = { + "metadata": metadata or {}, + "records": [_to_jsonable(record) for record in records], + } + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(manifest, indent=2, sort_keys=True)) + return manifest + + +def load_asset_manifest(path: str | Path) -> dict[str, Any]: + return json.loads(Path(path).read_text()) + + +def _to_jsonable(value: Any) -> Any: + if hasattr(value, "to_dict"): + return _to_jsonable(value.to_dict()) + if is_dataclass(value): + return _to_jsonable(asdict(value)) + if isinstance(value, dict): + return {str(k): _to_jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [_to_jsonable(v) for v in value] + if isinstance(value, Path): + return str(value) + return value diff --git a/stimulus_synthesis/config.py b/stimulus_synthesis/config.py new file mode 100644 index 0000000000000000000000000000000000000000..e6dadb9fa2213c5a14b1dd2ec939695bb7dc799d --- /dev/null +++ b/stimulus_synthesis/config.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + + +@dataclass +class StimulusSynthesisConfig: + encoder_model_id: str = "epfl-neuroai/vjepa2-encoder-basic" + encoder_call: str = "predict_fmri" + text_to_image_model_id: str = "stabilityai/sdxl-turbo" + image_to_video_model_id: str = "Lightricks/LTX-Video-0.9.8-13B-distilled" + default_objective: str = "indices_mean" + default_device: str = "cuda" + default_population_size: int = 20 + default_image_max_evals: int = 400 + default_video_max_evals: int = 200 + default_score_frames: int = 24 + default_score_size: int = 224 + default_image_batch_size: int = 16 + default_video_batch_size: int = 8 + default_mutation_rate: float = 0.25 + default_elite_frac: float = 0.35 + default_score_transform: dict[str, Any] | None = field( + default_factory=lambda: { + "enabled": False, + "crop_scale": 0.80, + "gaussian_sigma": 0.10, + "num_draws": 4, + "aggregate": "mean", + "seed": 0, + } + ) + default_image_kwargs: dict[str, Any] = field( + default_factory=lambda: {"num_inference_steps": 1, "guidance_scale": 0.0, "height": 512, "width": 512} + ) + default_video_kwargs: dict[str, Any] = field( + default_factory=lambda: {"num_inference_steps": 8, "num_frames": 25, "height": 512, "width": 512} + ) + + @classmethod + def from_json_file(cls, path: str | Path) -> "StimulusSynthesisConfig": + with open(path, "r") as f: + data = json.load(f) + return cls(**data) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "StimulusSynthesisConfig": + known = {field.name for field in cls.__dataclass_fields__.values()} + return cls(**{k: v for k, v in data.items() if k in known}) + + def to_dict(self) -> dict[str, Any]: + return { + "encoder_model_id": self.encoder_model_id, + "encoder_call": self.encoder_call, + "text_to_image_model_id": self.text_to_image_model_id, + "image_to_video_model_id": self.image_to_video_model_id, + "default_objective": self.default_objective, + "default_device": self.default_device, + "default_population_size": self.default_population_size, + "default_image_max_evals": self.default_image_max_evals, + "default_video_max_evals": self.default_video_max_evals, + "default_score_frames": self.default_score_frames, + "default_score_size": self.default_score_size, + "default_image_batch_size": self.default_image_batch_size, + "default_video_batch_size": self.default_video_batch_size, + "default_mutation_rate": self.default_mutation_rate, + "default_elite_frac": self.default_elite_frac, + "default_score_transform": self.default_score_transform, + "default_image_kwargs": self.default_image_kwargs, + "default_video_kwargs": self.default_video_kwargs, + } diff --git a/stimulus_synthesis/data/__init__.py b/stimulus_synthesis/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/stimulus_synthesis/data/roi_masks.npz b/stimulus_synthesis/data/roi_masks.npz new file mode 100644 index 0000000000000000000000000000000000000000..f03709cfdf6994878704a5b39ad1792c03e3c5d8 --- /dev/null +++ b/stimulus_synthesis/data/roi_masks.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3be584b579472777c7f41e3589e47b6cee5b608636c633bf605e1dd266f76914 +size 5758 diff --git a/stimulus_synthesis/data/searchlight_both.npz b/stimulus_synthesis/data/searchlight_both.npz new file mode 100644 index 0000000000000000000000000000000000000000..543134333d0323f4e276ec02141cbe7bada448a3 --- /dev/null +++ b/stimulus_synthesis/data/searchlight_both.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2c3e3fec1a88d56e182ccb1362f303f50fa24bae1d8796fee99f80b30e2bea76 +size 2585 diff --git a/stimulus_synthesis/data/searchlight_lh.npz b/stimulus_synthesis/data/searchlight_lh.npz new file mode 100644 index 0000000000000000000000000000000000000000..5d63abd2e891ddad9ec328d660e14c6767ca7d73 --- /dev/null +++ b/stimulus_synthesis/data/searchlight_lh.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a9ba6bc47b7d7075c0faf0cf6edd8acebaf9f2fe1f588899d6e7323b00f89979 +size 1379 diff --git a/stimulus_synthesis/data/searchlight_rh.npz b/stimulus_synthesis/data/searchlight_rh.npz new file mode 100644 index 0000000000000000000000000000000000000000..0423734f11e9869f9c6ed7ad3ce84a537115ec0e --- /dev/null +++ b/stimulus_synthesis/data/searchlight_rh.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:838c8ef226320a484c1776092d653d73d77721881b2c64113a88d13b83190b3a +size 1478 diff --git a/stimulus_synthesis/generators/__init__.py b/stimulus_synthesis/generators/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..abf69fcd5a1df8579550cb0c7df54d8ef322a3f1 --- /dev/null +++ b/stimulus_synthesis/generators/__init__.py @@ -0,0 +1,5 @@ +from .base import ImageToVideoGenerator, TextToImageGenerator +from .diffusers_i2v import DiffusersImageToVideoAdapter +from .diffusers_t2i import DiffusersTextToImageAdapter + +__all__ = ["TextToImageGenerator", "ImageToVideoGenerator", "DiffusersTextToImageAdapter", "DiffusersImageToVideoAdapter"] diff --git a/stimulus_synthesis/generators/base.py b/stimulus_synthesis/generators/base.py new file mode 100644 index 0000000000000000000000000000000000000000..842123db7e9c9fa5ce182dc961ce7ae2b11ed313 --- /dev/null +++ b/stimulus_synthesis/generators/base.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + + +class TextToImageGenerator(ABC): + @abstractmethod + def generate(self, prompts: list[str], *, generator: Any | None = None, **kwargs) -> list[Any]: + ... + + +class ImageToVideoGenerator(ABC): + @abstractmethod + def generate(self, image: Any, prompt: str, *, generator: Any | None = None, **kwargs) -> Any: + ... + + def generate_batch(self, images: list[Any], prompts: list[str], *, generators: list[Any] | None = None, **kwargs) -> list[Any]: + generators = generators or [None] * len(images) + return [ + self.generate(image, prompt, generator=gen, **kwargs) + for image, prompt, gen in zip(images, prompts, generators) + ] diff --git a/stimulus_synthesis/generators/diffusers_i2v.py b/stimulus_synthesis/generators/diffusers_i2v.py new file mode 100644 index 0000000000000000000000000000000000000000..fc4076a7d2a83e9173a41a9c92074ced90d5f799 --- /dev/null +++ b/stimulus_synthesis/generators/diffusers_i2v.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from typing import Any + +import torch +from diffusers import DiffusionPipeline + +from .base import ImageToVideoGenerator + + +class DiffusersImageToVideoAdapter(ImageToVideoGenerator): + def __init__(self, model_id: str, device: str = "cuda", torch_dtype: Any | None = None, pipeline: Any | None = None, **kwargs) -> None: + if pipeline is None: + dtype = torch_dtype + if dtype is None and str(device).startswith("cuda"): + dtype = torch.bfloat16 + pipeline_cls = kwargs.pop("pipeline_cls", None) or _default_i2v_pipeline_cls(model_id) + pipeline = pipeline_cls.from_pretrained(model_id, torch_dtype=dtype, **kwargs) + self.pipe = pipeline + self.device = device + if hasattr(self.pipe, "to"): + self.pipe.to(device) + if hasattr(self.pipe, "set_progress_bar_config"): + self.pipe.set_progress_bar_config(disable=True) + + def generate(self, image: Any, prompt: str, *, generator: Any | None = None, **kwargs) -> Any: + kwargs.setdefault("output_type", "pt") + try: + out = self.pipe(image=image, prompt=prompt, generator=generator, **kwargs) + except TypeError as exc: + raise TypeError( + "The configured image-to-video model does not support the default " + "`image=..., prompt=..., generator=..., **kwargs` signature. " + "Pass a custom ImageToVideoGenerator adapter." + ) from exc + return self._normalize_output(out) + + def generate_batch(self, images, prompts, *, generators=None, **kwargs): + images = list(images) + prompts = list(prompts) + if not prompts: + return [] + kwargs.setdefault("output_type", "pt") + try: + out = self.pipe(image=images, prompt=prompts, generator=generators, **kwargs) + frames = getattr(out, "frames", None) + if frames is None: + frames = getattr(out, "videos", None) + if torch.is_tensor(frames) and frames.ndim == 5 and frames.shape[0] == len(prompts): + return [frames[i] for i in range(len(prompts))] + if isinstance(frames, (list, tuple)) and len(frames) == len(prompts): + return list(frames) + except (RuntimeError, TypeError, ValueError): + pass + # Fall back to per-item generation (e.g. model can't batch, or OOM). + gens = generators if isinstance(generators, (list, tuple)) else [generators] * len(prompts) + return [self.generate(img, p, generator=g, **kwargs) for img, p, g in zip(images, prompts, gens)] + + @staticmethod + def _normalize_output(out: Any) -> Any: + frames = getattr(out, "frames", None) + if frames is None: + frames = getattr(out, "videos", None) + if frames is None: + return out + if torch.is_tensor(frames): + return frames[0] if frames.ndim == 5 else frames + if isinstance(frames, (list, tuple)) and frames: + return frames[0] + return frames + + +def _default_i2v_pipeline_cls(model_id: str): + if model_id == "Lightricks/LTX-Video": + from diffusers import LTXImageToVideoPipeline + + return LTXImageToVideoPipeline + return DiffusionPipeline diff --git a/stimulus_synthesis/generators/diffusers_t2i.py b/stimulus_synthesis/generators/diffusers_t2i.py new file mode 100644 index 0000000000000000000000000000000000000000..fbccb4ea5638d4e1634fdca580b89bca9b52d929 --- /dev/null +++ b/stimulus_synthesis/generators/diffusers_t2i.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +from typing import Any + +import torch +from diffusers import DiffusionPipeline + +from .base import TextToImageGenerator + + +class DiffusersTextToImageAdapter(TextToImageGenerator): + def __init__(self, model_id: str, device: str = "cuda", torch_dtype: Any | None = None, pipeline: Any | None = None, **kwargs) -> None: + if pipeline is None: + dtype = torch_dtype + if dtype is None and str(device).startswith("cuda"): + dtype = torch.float16 + pipeline = DiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype, **kwargs) + self.pipe = pipeline + self.device = device + if hasattr(self.pipe, "to"): + self.pipe.to(device) + if hasattr(self.pipe, "set_progress_bar_config"): + self.pipe.set_progress_bar_config(disable=True) + + def generate(self, prompts: list[str], *, generator: Any | None = None, **kwargs) -> list[Any]: + out = self.pipe(prompt=prompts, generator=generator, **kwargs) + if hasattr(out, "images"): + return list(out.images) + if isinstance(out, list): + return out + raise TypeError("Text-to-image pipeline output does not expose `.images`.") + + def generate_batch(self, prompts, *, generators=None, **kwargs): + out = self.pipe(prompt=list(prompts), generator=generators, **kwargs) + if hasattr(out, "images"): + return list(out.images) + if isinstance(out, list): + return out + raise TypeError("Text-to-image pipeline output does not expose `.images`.") diff --git a/stimulus_synthesis/media/__init__.py b/stimulus_synthesis/media/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..145f240226c1bdf5409b383f13371eb1391611e8 --- /dev/null +++ b/stimulus_synthesis/media/__init__.py @@ -0,0 +1,22 @@ +from .asset_decode import DecodedImage, DecodedVideo, decode_image, decode_video +from .asset_export import AssetExportRecord, export_image, export_video, sha256_file +from .asset_spec import ImageAssetSpec, VideoAssetSpec +from .normalize import video_to_t_c_h_w, videos_to_b_t_c_h_w +from .video_io import load_video_as_tensor, save_video + +__all__ = [ + "AssetExportRecord", + "DecodedImage", + "DecodedVideo", + "ImageAssetSpec", + "VideoAssetSpec", + "decode_image", + "decode_video", + "export_image", + "export_video", + "load_video_as_tensor", + "save_video", + "sha256_file", + "video_to_t_c_h_w", + "videos_to_b_t_c_h_w", +] diff --git a/stimulus_synthesis/media/asset_decode.py b/stimulus_synthesis/media/asset_decode.py new file mode 100644 index 0000000000000000000000000000000000000000..e77b6977eef851de3f11c3a01472fe497ea9433d --- /dev/null +++ b/stimulus_synthesis/media/asset_decode.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass +from fractions import Fraction +from pathlib import Path +from typing import Any + +import av +import numpy as np +import torch +from PIL import Image + +from .asset_export import sha256_file + + +@dataclass(frozen=True) +class DecodedImage: + path: str + sha256: str + image: torch.Tensor + width: int + height: int + + def metadata(self) -> dict[str, Any]: + data = asdict(self) + data.pop("image") + return data + + +@dataclass(frozen=True) +class DecodedVideo: + path: str + sha256: str + frames: torch.Tensor + width: int + height: int + num_frames: int + fps: float | None + + def metadata(self) -> dict[str, Any]: + data = asdict(self) + data.pop("frames") + return data + + +def decode_image(path: str | Path) -> DecodedImage: + path = Path(path) + image = Image.open(path).convert("RGB") + tensor = torch.from_numpy(np.asarray(image).copy()).float().permute(2, 0, 1) / 255.0 + return DecodedImage( + path=str(path), + sha256=sha256_file(path), + image=tensor.contiguous(), + width=image.width, + height=image.height, + ) + + +def decode_video(path: str | Path) -> DecodedVideo: + path = Path(path) + container = av.open(str(path)) + try: + stream = container.streams.video[0] + fps = _fraction_to_float(stream.average_rate or stream.base_rate) + frames = [ + torch.from_numpy(frame.to_ndarray(format="rgb24")).float().permute(2, 0, 1) / 255.0 + for frame in container.decode(video=0) + ] + finally: + container.close() + if not frames: + raise ValueError(f"Decoded video has no frames: {path}") + tensor = torch.stack(frames, dim=0).contiguous() + return DecodedVideo( + path=str(path), + sha256=sha256_file(path), + frames=tensor, + width=int(tensor.shape[-1]), + height=int(tensor.shape[-2]), + num_frames=int(tensor.shape[0]), + fps=fps, + ) + + +def _fraction_to_float(value: Fraction | None) -> float | None: + if value is None: + return None + return float(value) diff --git a/stimulus_synthesis/media/asset_export.py b/stimulus_synthesis/media/asset_export.py new file mode 100644 index 0000000000000000000000000000000000000000..33afe846c6066d0b2387eedaacf613aa4dac999b --- /dev/null +++ b/stimulus_synthesis/media/asset_export.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import hashlib +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import av +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +from .asset_spec import ImageAssetSpec, VideoAssetSpec +from .normalize import video_to_t_c_h_w + + +@dataclass(frozen=True) +class AssetExportRecord: + path: str + asset_type: str + sha256: str + bytes: int + spec: dict[str, Any] + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def sha256_file(path: str | Path) -> str: + h = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +def export_image(image: Any, path: str | Path, spec: ImageAssetSpec) -> AssetExportRecord: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + pil = _to_pil_rgb(image) + pil = _resize_pil(pil, (spec.width, spec.height), spec.resize_filter) + save_kwargs: dict[str, Any] = {} + fmt = spec.format.upper() + if spec.quality is not None and fmt in {"JPEG", "JPG", "WEBP"}: + save_kwargs["quality"] = int(spec.quality) + if fmt == "JPG": + fmt = "JPEG" + pil.save(path, format=fmt, **save_kwargs) + return AssetExportRecord( + path=str(path), + asset_type="image", + sha256=sha256_file(path), + bytes=path.stat().st_size, + spec=spec.to_dict(), + ) + + +def export_video(frames: Any, path: str | Path, spec: VideoAssetSpec) -> AssetExportRecord: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + video = video_to_t_c_h_w(frames).clamp(0.0, 1.0) + video = _match_frames(video, spec.num_frames) + video = _resize_video(video, (spec.height, spec.width), spec.resize_filter) + frames_np = (video.detach().cpu().permute(0, 2, 3, 1).numpy() * 255.0).round().clip(0, 255).astype(np.uint8) + + container = av.open(str(path), mode="w", format=spec.container) + try: + stream = container.add_stream(spec.codec, rate=spec.fps) + stream.width = spec.width + stream.height = spec.height + stream.pix_fmt = spec.pixel_format + options: dict[str, str] = {} + if spec.crf is not None: + options["crf"] = str(int(spec.crf)) + if spec.preset is not None: + options["preset"] = str(spec.preset) + stream.options = options + for frame_np in frames_np: + container.mux(stream.encode(av.VideoFrame.from_ndarray(frame_np, format="rgb24"))) + container.mux(stream.encode(None)) + finally: + container.close() + + return AssetExportRecord( + path=str(path), + asset_type="video", + sha256=sha256_file(path), + bytes=path.stat().st_size, + spec=spec.to_dict(), + ) + + +def _to_pil_rgb(image: Any) -> Image.Image: + if isinstance(image, Image.Image): + return image.convert("RGB") + if torch.is_tensor(image): + tensor = image.detach().cpu().float() + if tensor.ndim == 4 and tensor.shape[0] == 1: + tensor = tensor.squeeze(0) + if tensor.ndim != 3: + raise ValueError(f"Expected image tensor with 3 dims, got {tuple(tensor.shape)}") + if tensor.shape[0] == 3: + tensor = tensor.permute(1, 2, 0) + if tensor.max() <= 1.0: + tensor = tensor * 255.0 + arr = tensor.round().clamp(0, 255).byte().numpy() + return Image.fromarray(arr, mode="RGB") + if isinstance(image, np.ndarray): + arr = image + if arr.ndim != 3: + raise ValueError(f"Expected image array with 3 dims, got {arr.shape}") + if arr.shape[0] == 3 and arr.shape[-1] != 3: + arr = np.transpose(arr, (1, 2, 0)) + if arr.dtype != np.uint8: + arr = arr.astype(np.float32) + if arr.max() <= 1.0: + arr = arr * 255.0 + arr = np.rint(arr).clip(0, 255).astype(np.uint8) + return Image.fromarray(arr, mode="RGB") + raise TypeError(f"Unsupported image type: {type(image)!r}") + + +def _resize_pil(image: Image.Image, size: tuple[int, int], resize_filter: str) -> Image.Image: + if image.size == size: + return image + filters = { + "nearest": Image.Resampling.NEAREST, + "bilinear": Image.Resampling.BILINEAR, + "bicubic": Image.Resampling.BICUBIC, + "lanczos": Image.Resampling.LANCZOS, + } + return image.resize(size, filters.get(resize_filter.lower(), Image.Resampling.BICUBIC)) + + +def _match_frames(video: torch.Tensor, num_frames: int) -> torch.Tensor: + if video.shape[0] == num_frames: + return video + if video.shape[0] > num_frames: + idx = torch.linspace(0, video.shape[0] - 1, steps=num_frames).round().long() + return video[idx] + reps = int(np.ceil(num_frames / video.shape[0])) + return video.repeat((reps, 1, 1, 1))[:num_frames] + + +def _resize_video(video: torch.Tensor, size_hw: tuple[int, int], resize_filter: str) -> torch.Tensor: + if tuple(video.shape[-2:]) == size_hw: + return video.contiguous() + mode = resize_filter.lower() + if mode not in {"nearest", "bilinear", "bicubic"}: + mode = "bilinear" + kwargs = {} if mode == "nearest" else {"align_corners": False} + return F.interpolate(video, size=size_hw, mode=mode, **kwargs).contiguous() diff --git a/stimulus_synthesis/media/asset_spec.py b/stimulus_synthesis/media/asset_spec.py new file mode 100644 index 0000000000000000000000000000000000000000..1ed67e8e6728194d71cdd8682f6c508ec767fcee --- /dev/null +++ b/stimulus_synthesis/media/asset_spec.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass +from typing import Any + + +@dataclass(frozen=True) +class ImageAssetSpec: + width: int + height: int + format: str = "png" + quality: int | None = None + color_space: str = "srgb" + resize_filter: str = "bicubic" + resize_policy: str = "resize" + + def __post_init__(self) -> None: + _validate_positive_int("width", self.width) + _validate_positive_int("height", self.height) + if self.quality is not None and not (1 <= int(self.quality) <= 100): + raise ValueError("quality must be between 1 and 100 when set.") + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class VideoAssetSpec: + width: int + height: int + fps: int + num_frames: int + container: str = "mp4" + codec: str = "libx264" + crf: int | None = 10 + preset: str | None = "slow" + pixel_format: str = "yuv420p" + color_space: str = "srgb" + resize_filter: str = "bilinear" + resize_policy: str = "resize" + + def __post_init__(self) -> None: + _validate_positive_int("width", self.width) + _validate_positive_int("height", self.height) + _validate_positive_int("fps", self.fps) + _validate_positive_int("num_frames", self.num_frames) + if self.crf is not None and not (0 <= int(self.crf) <= 51): + raise ValueError("crf must be between 0 and 51 when set.") + if self.pixel_format == "yuv420p" and (self.width % 2 or self.height % 2): + raise ValueError("yuv420p video export requires even width and height.") + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def _validate_positive_int(name: str, value: int) -> None: + if int(value) <= 0: + raise ValueError(f"{name} must be a positive integer.") diff --git a/stimulus_synthesis/media/normalize.py b/stimulus_synthesis/media/normalize.py new file mode 100644 index 0000000000000000000000000000000000000000..0aac7afe39dfcb200c3361651965e9284da878a6 --- /dev/null +++ b/stimulus_synthesis/media/normalize.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + + +def _frame_to_tensor(frame: Any) -> torch.Tensor: + if isinstance(frame, Image.Image): + arr = np.asarray(frame.convert("RGB"), dtype=np.float32) / 255.0 + return torch.from_numpy(arr).permute(2, 0, 1) + if isinstance(frame, np.ndarray): + arr = frame.astype(np.float32, copy=False) + if arr.max() > 1.0: + arr = arr / 255.0 + tensor = torch.from_numpy(arr) + if tensor.ndim == 3 and tensor.shape[-1] == 3: + tensor = tensor.permute(2, 0, 1) + if tensor.ndim != 3: + raise ValueError(f"Expected frame array with 3 dims, got {arr.shape}") + return tensor.float().contiguous() + if torch.is_tensor(frame): + tensor = frame.detach().float() + if tensor.ndim == 3 and tensor.shape[-1] == 3: + tensor = tensor.permute(2, 0, 1) + if tensor.ndim != 3: + raise ValueError(f"Expected frame tensor with 3 dims, got {tuple(tensor.shape)}") + if tensor.max() > 1.0: + tensor = tensor / 255.0 + return tensor.contiguous() + raise TypeError(f"Unsupported frame type: {type(frame)!r}") + + +def video_to_t_c_h_w(video: Any) -> torch.Tensor: + if torch.is_tensor(video): + tensor = video.detach().float() + if tensor.ndim == 5 and tensor.shape[0] == 1: + tensor = tensor.squeeze(0) + if tensor.ndim == 3: + tensor = tensor.unsqueeze(0) + if tensor.ndim != 4: + raise ValueError(f"Expected video tensor with 4 dims, got {tuple(tensor.shape)}") + if tensor.shape[-1] == 3: + tensor = tensor.permute(0, 3, 1, 2) + if tensor.max() > 1.0: + tensor = tensor / 255.0 + return tensor.contiguous() + + if isinstance(video, np.ndarray): + arr = video.astype(np.float32, copy=False) + if arr.ndim == 5 and arr.shape[0] == 1: + arr = arr[0] + if arr.ndim == 3: + return _frame_to_tensor(arr).unsqueeze(0).contiguous() + if arr.ndim != 4: + raise ValueError(f"Expected video array with 4 dims, got {arr.shape}") + tensor = torch.from_numpy(arr) + if tensor.shape[-1] == 3: + tensor = tensor.permute(0, 3, 1, 2) + if tensor.max() > 1.0: + tensor = tensor.float() / 255.0 + return tensor.float().contiguous() + + if isinstance(video, (list, tuple)): + if not video: + raise ValueError("Video frame list is empty.") + return torch.stack([_frame_to_tensor(frame) for frame in video], dim=0).contiguous() + + if isinstance(video, Image.Image): + return _frame_to_tensor(video).unsqueeze(0).contiguous() + + raise TypeError(f"Unsupported video type: {type(video)!r}") + + +def videos_to_b_t_c_h_w(videos: list[Any], *, size: int | tuple[int, int] | None = None, num_frames: int | None = None) -> torch.Tensor: + tensors = [video_to_t_c_h_w(video) for video in videos] + if num_frames is not None: + tensors = [_match_frames(tensor, num_frames) for tensor in tensors] + if size is not None: + target_size = (size, size) if isinstance(size, int) else tuple(size) + tensors = [_resize_video(tensor, target_size) for tensor in tensors] + return torch.stack(tensors, dim=0).clamp(0.0, 1.0).contiguous() + + +def _match_frames(video: torch.Tensor, num_frames: int) -> torch.Tensor: + if video.shape[0] == num_frames: + return video + if video.shape[0] > num_frames: + idx = torch.linspace(0, video.shape[0] - 1, steps=num_frames).round().long() + return video[idx] + reps = int(np.ceil(num_frames / video.shape[0])) + return video.repeat((reps, 1, 1, 1))[:num_frames] + + +def _resize_video(video: torch.Tensor, size: tuple[int, int]) -> torch.Tensor: + return F.interpolate(video, size=size, mode="bilinear", align_corners=False) diff --git a/stimulus_synthesis/media/video_io.py b/stimulus_synthesis/media/video_io.py new file mode 100644 index 0000000000000000000000000000000000000000..82d3ecae9bb8e696874ffe5a8cb12dcb0a2f6f27 --- /dev/null +++ b/stimulus_synthesis/media/video_io.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import av +import numpy as np +import torch + + +def load_video_as_tensor(mp4_path: str) -> torch.Tensor: + container = av.open(mp4_path) + try: + frames = [ + torch.from_numpy(frame.to_ndarray(format="rgb24")).float() / 255.0 + for frame in container.decode(video=0) + ] + finally: + container.close() + return torch.stack(frames).permute(0, 3, 1, 2).contiguous() + + +def save_video(tensor: torch.Tensor, path: str, fps: int = 24) -> None: + frames_np = (tensor.detach().cpu().permute(0, 2, 3, 1).numpy() * 255).clip(0, 255).astype(np.uint8) + height, width = frames_np.shape[1], frames_np.shape[2] + height, width = height - (height % 2), width - (width % 2) + frames_np = frames_np[:, :height, :width, :] + + container = av.open(path, mode="w") + stream = container.add_stream("libx264", rate=fps) + stream.width, stream.height = width, height + stream.pix_fmt = "yuv420p" + stream.options = {"crf": "10", "preset": "slow"} + for frame_np in frames_np: + container.mux(stream.encode(av.VideoFrame.from_ndarray(frame_np, format="rgb24"))) + container.mux(stream.encode(None)) + container.close() diff --git a/stimulus_synthesis/neuro/__init__.py b/stimulus_synthesis/neuro/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f7f21f45d4550eb19c8fe02e09fb53e3855bedc3 --- /dev/null +++ b/stimulus_synthesis/neuro/__init__.py @@ -0,0 +1,3 @@ +from .roi import resolve_driving_voxels, available_rois, searchlight_counts, N_TOTAL_VOXELS + +__all__ = ["resolve_driving_voxels", "available_rois", "searchlight_counts", "N_TOTAL_VOXELS"] diff --git a/stimulus_synthesis/neuro/roi.py b/stimulus_synthesis/neuro/roi.py new file mode 100644 index 0000000000000000000000000000000000000000..892ea80cc96566554079038a2128e61a8e1a4f1a --- /dev/null +++ b/stimulus_synthesis/neuro/roi.py @@ -0,0 +1,110 @@ +"""Self-contained ROI -> voxel-mask resolver. + +Removes the need for an external ROI-atlas package (nilearn + fsaverage atlas +data). The ROI and +searchlight masks are precomputed and shipped as data under ``stimulus_synthesis/data``, +so resolution here needs only numpy. + +Token grammar (identical to the original): +- ``ROI`` -> both hemispheres (e.g. ``FFA``) +- ``ROI_lh`` / ``ROI_rh`` +- ``SL-`` -> both-hemisphere searchlight region (1-indexed) +- ``SL-_lh`` / ``SL-_rh`` +- comma-separated tokens are unioned. +""" +from __future__ import annotations + +import re +from pathlib import Path + +import numpy as np + +N_TOTAL_VOXELS = 20484 +N_HEMI_VOXELS = N_TOTAL_VOXELS // 2 + +_DATA_DIR = Path(__file__).parent.parent / "data" +_ROI_HEMI_PATTERN = re.compile(r"^(?P.+?)_(?Plh|rh)$", flags=re.IGNORECASE) +_SL_PATTERN = re.compile(r"^SL-(?P\d+)(?:_(?Plh|rh))?$", flags=re.IGNORECASE) + +_named_cache: dict[str, np.ndarray] | None = None +_sl_cache: dict[str, np.ndarray] = {} + + +def _load_named() -> dict[str, np.ndarray]: + global _named_cache + if _named_cache is None: + with np.load(_DATA_DIR / "roi_masks.npz") as z: + _named_cache = {k: z[k].astype(bool) for k in z.files} + return _named_cache + + +def _load_searchlight(key: str) -> np.ndarray: + if key not in _sl_cache: + with np.load(_DATA_DIR / f"searchlight_{key}.npz") as z: + _sl_cache[key] = z["regions"].astype(bool) + return _sl_cache[key] + + +def _hemi_slice(hemi: str) -> slice: + hemi = hemi.lower() + if hemi == "lh": + return slice(0, N_HEMI_VOXELS) + if hemi == "rh": + return slice(N_HEMI_VOXELS, N_TOTAL_VOXELS) + raise ValueError(f"Invalid hemisphere '{hemi}'. Expected 'lh' or 'rh'.") + + +def _restrict_to_hemi(mask: np.ndarray, hemi: str | None) -> np.ndarray: + if hemi is None: + return mask + restricted = np.zeros_like(mask) + restricted[_hemi_slice(hemi)] = mask[_hemi_slice(hemi)] + return restricted + + +def available_rois() -> list[str]: + """Named ROIs shipped with the package.""" + return sorted(_load_named()) + + +def searchlight_counts() -> dict[str, int]: + """Number of searchlight regions available per hemisphere key.""" + return {k: len(_load_searchlight(k)) for k in ("both", "lh", "rh")} + + +def resolve_driving_voxels(target_roi: str) -> np.ndarray: + """Resolve ROI token(s) to a full-brain boolean mask of length 20484.""" + rois = [x.strip() for x in str(target_roi).split(",") if x.strip()] + if not rois: + raise ValueError("target_roi must contain at least one ROI token.") + + combined = np.zeros(N_TOTAL_VOXELS, dtype=bool) + for token in rois: + sl_match = _SL_PATTERN.match(token) + if sl_match is not None: + idx = int(sl_match.group("idx")) + hemi = sl_match.group("hemi") + key = "both" if hemi is None else hemi.lower() + regions = _load_searchlight(key) + if idx < 1 or idx > len(regions): + raise ValueError(f"Invalid {token}: n must be in [1, {len(regions)}].") + combined |= regions[idx - 1] + continue + + hemi = None + name = token + roi_match = _ROI_HEMI_PATTERN.match(token) + if roi_match is not None: + name = roi_match.group("name") + hemi = roi_match.group("hemi").lower() + + named = _load_named() + if name not in named: + lookup = {k.lower(): k for k in named} + if name.lower() in lookup: + name = lookup[name.lower()] + else: + raise KeyError(f"Unknown ROI '{name}'. Available: {sorted(named)}") + combined |= _restrict_to_hemi(named[name].copy(), hemi) + + return combined diff --git a/stimulus_synthesis/outputs.py b/stimulus_synthesis/outputs.py new file mode 100644 index 0000000000000000000000000000000000000000..31f5cca9de69431f1f6f57ab535f54490e43dc8e --- /dev/null +++ b/stimulus_synthesis/outputs.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class StimulusCandidate: + prompt: str + score: float + image: Any | None = None + video: Any | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class StimulusSynthesisOutput: + candidates: list[StimulusCandidate] + best_prompt: str + best_score: float + history_best: list[float] + metadata: dict[str, Any] = field(default_factory=dict) + + @property + def best(self) -> StimulusCandidate: + return self.candidates[0] diff --git a/stimulus_synthesis/paths.py b/stimulus_synthesis/paths.py new file mode 100644 index 0000000000000000000000000000000000000000..5b5bdabb97d73104e479901dc1157113380558c4 --- /dev/null +++ b/stimulus_synthesis/paths.py @@ -0,0 +1,77 @@ +"""Repo-local cache/path configuration. + +Cache directory resolution order: +1. ``NEVO_CACHE_DIR`` environment variable (may be set via a ``.env`` file at the + repository root), +2. otherwise ``/cache``. + +``configure_cache()`` points HuggingFace/torch caches at that directory (using +``setdefault`` so any caller-provided env wins). No third-party dotenv dependency; +the ``.env`` parser here handles simple ``KEY=VALUE`` lines. +""" +from __future__ import annotations + +import os +from pathlib import Path + +_REPO_ROOT = Path(__file__).parent.parent + + +def _load_dotenv(root: Path) -> None: + env_path = root / ".env" + if not env_path.exists(): + return + for line in env_path.read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + key = key.strip() + value = value.strip().strip('"').strip("'") + if key: + os.environ.setdefault(key, value) + + +def get_cache_dir() -> Path: + """Resolve and create the cache directory.""" + _load_dotenv(_REPO_ROOT) + override = os.environ.get("NEVO_CACHE_DIR") + path = Path(override).expanduser() if override else (_REPO_ROOT / "cache") + path.mkdir(parents=True, exist_ok=True) + return path + + +def _point_env_at(cache: Path) -> Path: + cache.mkdir(parents=True, exist_ok=True) + os.environ.setdefault("HF_HOME", str(cache / "huggingface")) + os.environ.setdefault("HUGGINGFACE_HUB_CACHE", str(cache / "huggingface" / "hub")) + os.environ.setdefault("TORCH_HOME", str(cache / "torch")) + return cache + + +def configure_cache() -> Path | None: + """Configure the HuggingFace/torch cache location. + + Priority: + 1. ``NEVO_CACHE_DIR`` (from environment or ``.env``) β€” use it. + 2. Otherwise, respect the system/user-default HuggingFace cache + (``HF_HOME``/``HUGGINGFACE_HUB_CACHE`` if set, else ``~/.cache/huggingface``) + and leave the environment untouched. + 3. Only if no default is resolvable (no usable home directory) fall back to + ``/cache``. + + Returns the cache dir when this function set one, or ``None`` when the system + default is left in place. + """ + _load_dotenv(_REPO_ROOT) + override = os.environ.get("NEVO_CACHE_DIR") + if override: + return _point_env_at(Path(override).expanduser()) + # Respect an already-configured / default HuggingFace cache. + if os.environ.get("HF_HOME") or os.environ.get("HUGGINGFACE_HUB_CACHE"): + return None + home = os.path.expanduser("~") + if home and home != "~" and os.path.isdir(home): + return None # HuggingFace will use its own default (~/.cache/huggingface) + # No usable default: fall back to a repo-local cache. + return _point_env_at(_REPO_ROOT / "cache") diff --git a/stimulus_synthesis/pipeline.py b/stimulus_synthesis/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..b83f003c1890f50b349fabb0cebf3c53998e8e94 --- /dev/null +++ b/stimulus_synthesis/pipeline.py @@ -0,0 +1,420 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from diffusers import DiffusionPipeline + +from .config import StimulusSynthesisConfig +from .generators.diffusers_i2v import DiffusersImageToVideoAdapter +from .generators.diffusers_t2i import DiffusersTextToImageAdapter +from .outputs import StimulusCandidate, StimulusSynthesisOutput +from .scoring.encoder_scorer import EncoderScorer +from .scoring.robust_transform import RobustTransformScorer, RobustTransformSpec +from .spaces.structured_neuro_space import StructuredArtPromptSpace, VideoMotionPromptSpace, art_data, make_t2v_art_data +from .search.genetic import GeneticSearch +from .spaces.prompt_space import PromptSearchSpace + + +class NevoPipeline(DiffusionPipeline): + config_name = "stimulus_synthesis_config.json" + + def __init__( + self, + synthesis_config: StimulusSynthesisConfig | dict[str, Any] | None = None, + text_to_image=None, + image_to_video=None, + scorer=None, + search_space=None, + ) -> None: + super().__init__() + if isinstance(synthesis_config, dict): + synthesis_config = StimulusSynthesisConfig.from_dict(synthesis_config) + object.__setattr__(self, "synthesis_config", synthesis_config or StimulusSynthesisConfig()) + object.__setattr__(self, "text_to_image", text_to_image) + object.__setattr__(self, "image_to_video", image_to_video) + object.__setattr__(self, "scorer", scorer) + object.__setattr__(self, "search_space", search_space) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: str | Path, *args, **kwargs): + text_to_image = kwargs.pop("text_to_image", None) + image_to_video = kwargs.pop("image_to_video", None) + scorer = kwargs.pop("scorer", None) + search_space = kwargs.pop("search_space", None) + config_override = kwargs.pop("synthesis_config", None) + device = kwargs.pop("device", None) + + path = Path(pretrained_model_name_or_path) + if path.exists(): + config_path = path / cls.config_name + synthesis_config = StimulusSynthesisConfig.from_json_file(config_path) if config_path.exists() else StimulusSynthesisConfig() + if config_override is not None: + synthesis_config = StimulusSynthesisConfig.from_dict({**synthesis_config.to_dict(), **_as_dict(config_override)}) + return cls( + synthesis_config=synthesis_config, + text_to_image=text_to_image, + image_to_video=image_to_video, + scorer=scorer, + search_space=search_space, + ) + + from huggingface_hub import hf_hub_download + + config_file = hf_hub_download(str(pretrained_model_name_or_path), cls.config_name, repo_type="model", **_hub_kwargs(kwargs)) + synthesis_config = StimulusSynthesisConfig.from_json_file(config_file) + if config_override is not None: + synthesis_config = StimulusSynthesisConfig.from_dict({**synthesis_config.to_dict(), **_as_dict(config_override)}) + if device is not None: + synthesis_config.default_device = device + return cls( + synthesis_config=synthesis_config, + text_to_image=text_to_image, + image_to_video=image_to_video, + scorer=scorer, + search_space=search_space, + ) + + def _ensure_components(self, device: str | None = None) -> None: + cfg = self.synthesis_config + device = device or cfg.default_device + if device == "cuda" and not torch.cuda.is_available(): + device = "cpu" + if self.text_to_image is None: + object.__setattr__(self, "text_to_image", DiffusersTextToImageAdapter(cfg.text_to_image_model_id, device=device)) + if self.image_to_video is None: + object.__setattr__(self, "image_to_video", DiffusersImageToVideoAdapter(cfg.image_to_video_model_id, device=device)) + if self.scorer is None: + base_scorer = EncoderScorer( + cfg.encoder_model_id, + encoder_call=cfg.encoder_call, + objective=cfg.default_objective, + device=device, + ) + transform_spec = RobustTransformSpec.from_dict(dict(cfg.default_score_transform or {})) + scorer = RobustTransformScorer(base_scorer, transform_spec) if transform_spec is not None else base_scorer + object.__setattr__(self, "scorer", scorer) + + def make_search_space( + self, + roi: str | None = None, + *, + enforce_general_search_space: bool = False, + search_space=None, + prompt_banks: dict[str, Any] | None = None, + seed_prompts: list[str] | None = None, + ): + """Select the prompt search space for the single-stage fallback path. + + When an ``roi`` is given, an ROI-aware *enhanced* structured space is used: + only the prompt categories relevant to that region are searched (the rest are + locked). Pass ``enforce_general_search_space=True`` to search the general + (all-category) space instead. An explicit ``search_space``, a pipeline-level + ``search_space``, or ``prompt_banks`` take precedence. + """ + if search_space is not None: + return search_space + if self.search_space is not None: + return self.search_space + if prompt_banks is not None: + return PromptSearchSpace(prompt_banks=prompt_banks, seed_prompts=seed_prompts) + if roi is not None: + space_roi = None if enforce_general_search_space else roi + return StructuredArtPromptSpace(make_t2v_art_data(), roi=space_roi) + return PromptSearchSpace(prompt_banks=prompt_banks, seed_prompts=seed_prompts) + + def _run_two_stage( + self, + *, + target, + roi, + enforce_general_search_space, + progress, + image_max_evals, + video_max_evals, + image_batch_size, + video_batch_size, + population_size, + seed, + text_to_image, + image_to_video, + scorer, + image_kwargs, + video_kwargs, + score_kwargs, + score_size, + num_frames, + ) -> StimulusSynthesisOutput: + """Two-stage evolutionary search (matches the paper / batch runners): + + Stage 1 evolves the *image* prompt, scored on the generated image; + Stage 2 freezes the best image and evolves the *motion* prompt, scored on video. + """ + cfg = self.synthesis_config + space_roi = None if enforce_general_search_space else roi + pop = max(2, int(population_size or cfg.default_population_size)) + image_evals = max(2, int(image_max_evals or cfg.default_image_max_evals)) + video_evals = max(2, int(video_max_evals or cfg.default_video_max_evals)) + score_frames = num_frames if num_frames is not None else cfg.default_score_frames + + # ---- Stage 1: image prompt search ---- + image_search = GeneticSearch( + max_evals=image_evals, + population_size=max(2, min(pop, image_evals)), + n_init=max(2, min(pop, image_evals)), + mutation_rate=cfg.default_mutation_rate, + elite_frac=cfg.default_elite_frac, + image_kwargs=image_kwargs, + video_kwargs={}, + score_kwargs=score_kwargs, + score_size=score_size, + num_frames=score_frames, + image_batch_size=image_batch_size, + video_batch_size=video_batch_size, + show_progress=progress, + progress_desc="Stage 1 - image", + ) + image_result = image_search.run( + StructuredArtPromptSpace(art_data, roi=space_roi), + text_to_image, + _StaticImageToVideo(num_frames=score_frames), + scorer, + target, + seed=seed, + ) + best_image = image_result.best_image + if best_image is None: + best_image = text_to_image.generate([image_result.best_prompt], **(image_kwargs or {}))[0] + + # ---- Stage 2: motion prompt search on the fixed best image ---- + video_search = GeneticSearch( + max_evals=video_evals, + population_size=max(2, min(pop, video_evals)), + n_init=max(2, min(pop, video_evals)), + mutation_rate=cfg.default_mutation_rate, + elite_frac=cfg.default_elite_frac, + image_kwargs={}, + video_kwargs=video_kwargs, + score_kwargs=score_kwargs, + score_size=score_size, + num_frames=score_frames, + image_batch_size=image_batch_size, + video_batch_size=video_batch_size, + show_progress=progress, + progress_desc="Stage 2 - video", + ) + video_result = video_search.run( + VideoMotionPromptSpace(roi=space_roi), + _FixedImageT2I(best_image), + image_to_video, + scorer, + target, + seed=seed, + ) + best_video = video_result.best_video + if best_video is None: + best_video = image_to_video.generate(best_image, video_result.best_prompt, **(video_kwargs or {})) + + best_prompt = ", ".join(p for p in (image_result.best_prompt, video_result.best_prompt) if p) + candidate = StimulusCandidate( + prompt=best_prompt, + score=video_result.best_score, + image=best_image, + video=best_video, + metadata={ + "rank": 1, + "image_prompt": image_result.best_prompt, + "image_score": image_result.best_score, + "video_prompt": video_result.best_prompt, + }, + ) + return StimulusSynthesisOutput( + candidates=[candidate], + best_prompt=best_prompt, + best_score=video_result.best_score, + history_best=video_result.history_best, + metadata={ + "encoder_model_id": cfg.encoder_model_id, + "text_to_image_model_id": cfg.text_to_image_model_id, + "image_to_video_model_id": cfg.image_to_video_model_id, + "objective": cfg.default_objective, + "two_stage": True, + "image_max_evals": image_evals, + "video_max_evals": video_evals, + "population_size": pop, + "seed": seed, + }, + ) + + def __call__( + self, + target=None, + seed_prompts: list[str] | None = None, + *, + roi: str | None = None, + enforce_general_search_space: bool = False, + progress: bool = False, + image_batch_size: int | None = None, + video_batch_size: int | None = None, + image_max_evals: int | None = None, + video_max_evals: int | None = None, + population_size: int | None = None, + seed: int | None = None, + prompt_banks: dict[str, list[str]] | None = None, + search_space=None, + text_to_image=None, + image_to_video=None, + scorer=None, + device: str | None = None, + image_kwargs: dict[str, Any] | None = None, + video_kwargs: dict[str, Any] | None = None, + score_kwargs: dict[str, Any] | None = None, + score_size: int | tuple[int, int] | None = None, + num_frames: int | None = None, + ) -> StimulusSynthesisOutput: + cfg = self.synthesis_config + if seed is None: + import secrets + seed = int(secrets.randbelow(2**31)) + self._ensure_components(device=device) + text_to_image = text_to_image or self.text_to_image + image_to_video = image_to_video or self.image_to_video + scorer = scorer or self.scorer + img_bs = image_batch_size if image_batch_size is not None else cfg.default_image_batch_size + vid_bs = video_batch_size if video_batch_size is not None else cfg.default_video_batch_size + pop = max(2, population_size if population_size is not None else cfg.default_population_size) + image_evals = max(2, image_max_evals if image_max_evals is not None else cfg.default_image_max_evals) + video_evals = max(2, video_max_evals if video_max_evals is not None else cfg.default_video_max_evals) + image_kwargs = {**(cfg.default_image_kwargs or {}), **(image_kwargs or {})} + video_kwargs = {**(cfg.default_video_kwargs or {}), **(video_kwargs or {})} + score_size = score_size if score_size is not None else cfg.default_score_size + + if roi is not None and target is None: + from .neuro import resolve_driving_voxels + mask = resolve_driving_voxels(roi) + target = {"type": "indices", "indices": np.flatnonzero(mask).astype(int).tolist()} + if target is None: + raise ValueError("Provide either `target` or `roi`.") + + # Default: two-stage structured search (evolve the image prompt, then the motion + # prompt on the fixed best image). An explicit search space / prompt bank / seed + # prompts falls back to a single joint search over that space. + if search_space is None and self.search_space is None and prompt_banks is None and seed_prompts is None: + return self._run_two_stage( + target=target, + roi=roi, + enforce_general_search_space=enforce_general_search_space, + progress=progress, + image_max_evals=image_evals, + video_max_evals=video_evals, + image_batch_size=img_bs, + video_batch_size=vid_bs, + population_size=pop, + seed=seed, + text_to_image=text_to_image, + image_to_video=image_to_video, + scorer=scorer, + image_kwargs=image_kwargs, + video_kwargs=video_kwargs, + score_kwargs=score_kwargs, + score_size=score_size, + num_frames=num_frames, + ) + + # ---- single-stage fallback (explicit search space / prompt bank / seed prompts) ---- + space = self.make_search_space( + roi=roi, + enforce_general_search_space=enforce_general_search_space, + search_space=search_space, + prompt_banks=prompt_banks, + seed_prompts=seed_prompts, + ) + search = GeneticSearch( + max_evals=image_evals, + population_size=max(2, min(pop, image_evals)), + n_init=max(2, min(pop, image_evals)), + mutation_rate=cfg.default_mutation_rate, + elite_frac=cfg.default_elite_frac, + image_kwargs=image_kwargs, + video_kwargs=video_kwargs, + score_kwargs=score_kwargs, + score_size=score_size, + num_frames=num_frames, + image_batch_size=img_bs, + video_batch_size=vid_bs, + show_progress=progress, + ) + result = search.run(space, text_to_image, image_to_video, scorer, target, seed=seed) + best_image = result.best_image + if best_image is None: + best_image = text_to_image.generate([result.best_prompt], **(image_kwargs or {}))[0] + best_video = result.best_video + if best_video is None: + best_video = image_to_video.generate(best_image, result.best_prompt, **(video_kwargs or {})) + candidate = StimulusCandidate( + prompt=result.best_prompt, + score=result.best_score, + image=best_image, + video=best_video, + metadata={"rank": 1, **result.best_metadata}, + ) + return StimulusSynthesisOutput( + candidates=[candidate], + best_prompt=result.best_prompt, + best_score=result.best_score, + history_best=result.history_best, + metadata={ + "encoder_model_id": cfg.encoder_model_id, + "text_to_image_model_id": cfg.text_to_image_model_id, + "image_to_video_model_id": cfg.image_to_video_model_id, + "objective": cfg.default_objective, + "max_evals": image_evals, + "seed": seed, + }, + ) + + +class _StaticImageToVideo: + """Stage-1 image-to-video stand-in: replicate the still image into an ``num_frames`` + clip (no motion) so the video encoder can score it. ``num_frames`` is the value + passed to the pipeline (or its default).""" + + def __init__(self, num_frames: int = 16): + self._num_frames = max(2, int(num_frames)) + + def generate(self, image, prompt, *, generator=None, **kwargs): + from .media.normalize import _frame_to_tensor + frame = _frame_to_tensor(image) # (C,H,W) in [0,1] β€” convert the still ONCE + return frame.unsqueeze(0).expand(self._num_frames, -1, -1, -1) + + def generate_batch(self, images, prompts, *, generators=None, **kwargs): + return [self.generate(img, prompt) for img, prompt in zip(images, prompts)] + + +class _FixedImageT2I: + """Stage-2 text-to-image stand-in: always returns the fixed best image from stage 1.""" + + def __init__(self, image): + self._image = image + + def generate(self, prompts, *, generator=None, **kwargs): + return [self._image for _ in prompts] + + +def _as_dict(value: Any) -> dict[str, Any]: + if isinstance(value, StimulusSynthesisConfig): + return value.to_dict() + if isinstance(value, dict): + return value + if isinstance(value, (str, Path)): + with open(value, "r") as f: + return json.load(f) + raise TypeError(f"Unsupported synthesis_config override: {type(value)!r}") + + +def _hub_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: + allowed = {"revision", "token", "cache_dir", "local_files_only"} + return {k: kwargs[k] for k in list(kwargs.keys()) if k in allowed} diff --git a/stimulus_synthesis/scoring/__init__.py b/stimulus_synthesis/scoring/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..90d019c83f10b570b6d1859a1edee12146d85e67 --- /dev/null +++ b/stimulus_synthesis/scoring/__init__.py @@ -0,0 +1,26 @@ +from .asset_scorer import AssetScoreRecord, AssetScorer, score_image_asset, score_video_asset +from .base import Scorer +from .encoder_preprocess import EncoderPreparedInput, EncoderPreprocessSpec, prepare_image_for_encoder, prepare_video_for_encoder +from .encoder_scorer import EncoderScorer +from .objectives import build_objective +from .robust_transform import RobustTransformScorer, RobustTransformSpec, apply_robust_transform +from .targets import TargetSpec, parse_target + +__all__ = [ + "AssetScoreRecord", + "AssetScorer", + "EncoderPreparedInput", + "EncoderPreprocessSpec", + "Scorer", + "EncoderScorer", + "RobustTransformScorer", + "RobustTransformSpec", + "TargetSpec", + "apply_robust_transform", + "build_objective", + "parse_target", + "prepare_image_for_encoder", + "prepare_video_for_encoder", + "score_image_asset", + "score_video_asset", +] diff --git a/stimulus_synthesis/scoring/asset_scorer.py b/stimulus_synthesis/scoring/asset_scorer.py new file mode 100644 index 0000000000000000000000000000000000000000..f6a74f05a9ee61df5c76865f212b3094d8f72b61 --- /dev/null +++ b/stimulus_synthesis/scoring/asset_scorer.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +from stimulus_synthesis.media.asset_decode import decode_image, decode_video +from stimulus_synthesis.scoring.encoder_preprocess import EncoderPreprocessSpec, prepare_image_for_encoder, prepare_video_for_encoder + + +@dataclass(frozen=True) +class AssetScoreRecord: + path: str + asset_type: str + sha256: str + score: float + sampled_frame_indices: list[int] + decoded: dict[str, Any] + preprocess: dict[str, Any] + asset_spec: dict[str, Any] | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +class AssetScorer: + def __init__( + self, + scorer: Any, + target: Any, + *, + preprocess_spec: EncoderPreprocessSpec | None = None, + score_kwargs: dict[str, Any] | None = None, + ) -> None: + self.scorer = scorer + self.target = target + self.preprocess_spec = preprocess_spec or EncoderPreprocessSpec() + self.score_kwargs = score_kwargs or {} + + def score_image(self, path: str | Path, *, asset_spec: Any | None = None, metadata: dict[str, Any] | None = None) -> AssetScoreRecord: + decoded = decode_image(path) + prepared = prepare_image_for_encoder(decoded.image, self.preprocess_spec) + score = self.scorer.score(prepared.videos, self.target, **self.score_kwargs)[0] + return AssetScoreRecord( + path=decoded.path, + asset_type="image", + sha256=decoded.sha256, + score=float(score), + sampled_frame_indices=prepared.frame_indices, + decoded=decoded.metadata(), + preprocess=prepared.spec, + asset_spec=_spec_to_dict(asset_spec), + metadata=metadata or {}, + ) + + def score_video(self, path: str | Path, *, asset_spec: Any | None = None, metadata: dict[str, Any] | None = None) -> AssetScoreRecord: + decoded = decode_video(path) + prepared = prepare_video_for_encoder(decoded.frames, self.preprocess_spec) + score = self.scorer.score(prepared.videos, self.target, **self.score_kwargs)[0] + return AssetScoreRecord( + path=decoded.path, + asset_type="video", + sha256=decoded.sha256, + score=float(score), + sampled_frame_indices=prepared.frame_indices, + decoded=decoded.metadata(), + preprocess=prepared.spec, + asset_spec=_spec_to_dict(asset_spec), + metadata=metadata or {}, + ) + + +def score_image_asset( + path: str | Path, + scorer: Any, + target: Any, + *, + preprocess_spec: EncoderPreprocessSpec | None = None, + asset_spec: Any | None = None, + score_kwargs: dict[str, Any] | None = None, + metadata: dict[str, Any] | None = None, +) -> AssetScoreRecord: + return AssetScorer(scorer, target, preprocess_spec=preprocess_spec, score_kwargs=score_kwargs).score_image( + path, asset_spec=asset_spec, metadata=metadata + ) + + +def score_video_asset( + path: str | Path, + scorer: Any, + target: Any, + *, + preprocess_spec: EncoderPreprocessSpec | None = None, + asset_spec: Any | None = None, + score_kwargs: dict[str, Any] | None = None, + metadata: dict[str, Any] | None = None, +) -> AssetScoreRecord: + return AssetScorer(scorer, target, preprocess_spec=preprocess_spec, score_kwargs=score_kwargs).score_video( + path, asset_spec=asset_spec, metadata=metadata + ) + + +def _spec_to_dict(spec: Any | None) -> dict[str, Any] | None: + if spec is None: + return None + if hasattr(spec, "to_dict"): + return spec.to_dict() + if isinstance(spec, dict): + return dict(spec) + return asdict(spec) diff --git a/stimulus_synthesis/scoring/base.py b/stimulus_synthesis/scoring/base.py new file mode 100644 index 0000000000000000000000000000000000000000..155f21cb4d8996f381e8810eaac3911958afbf74 --- /dev/null +++ b/stimulus_synthesis/scoring/base.py @@ -0,0 +1,10 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + + +class Scorer(ABC): + @abstractmethod + def score(self, videos: Any, target: Any, **kwargs) -> list[float]: + ... diff --git a/stimulus_synthesis/scoring/encoder_preprocess.py b/stimulus_synthesis/scoring/encoder_preprocess.py new file mode 100644 index 0000000000000000000000000000000000000000..1083cc195606fec39f741426704eca8e3b715526 --- /dev/null +++ b/stimulus_synthesis/scoring/encoder_preprocess.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +from stimulus_synthesis.media.normalize import video_to_t_c_h_w + + +@dataclass(frozen=True) +class EncoderPreprocessSpec: + size: int | tuple[int, int] | None = 224 + num_frames: int | None = None + frame_sampling: str = "uniform" + normalize_mean: tuple[float, float, float] | None = None + normalize_std: tuple[float, float, float] | None = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class EncoderPreparedInput: + videos: torch.Tensor + frame_indices: list[int] + spec: dict[str, Any] + + +def prepare_image_for_encoder(image: Any, spec: EncoderPreprocessSpec | None = None) -> EncoderPreparedInput: + spec = spec or EncoderPreprocessSpec(num_frames=1) + frame = _image_to_c_h_w(image) + num_frames = int(spec.num_frames or 1) + video = frame.unsqueeze(0).repeat(num_frames, 1, 1, 1) + video = _resize_video(video, spec.size) + video = _normalize(video, spec) + return EncoderPreparedInput(videos=video.unsqueeze(0).contiguous(), frame_indices=[0] * num_frames, spec=spec.to_dict()) + + +def prepare_video_for_encoder(frames: Any, spec: EncoderPreprocessSpec | None = None) -> EncoderPreparedInput: + spec = spec or EncoderPreprocessSpec() + video = video_to_t_c_h_w(frames) + indices = sample_frame_indices(video.shape[0], spec.num_frames, spec.frame_sampling) + if indices: + video = video[torch.as_tensor(indices, dtype=torch.long)] + video = _resize_video(video, spec.size) + video = _normalize(video, spec) + return EncoderPreparedInput(videos=video.unsqueeze(0).contiguous(), frame_indices=indices, spec=spec.to_dict()) + + +def sample_frame_indices(total_frames: int, num_frames: int | None, policy: str = "uniform") -> list[int]: + if total_frames <= 0: + raise ValueError("total_frames must be positive.") + if num_frames is None: + return list(range(total_frames)) + if num_frames <= 0: + raise ValueError("num_frames must be positive when set.") + if policy != "uniform": + raise ValueError(f"Unsupported frame sampling policy: {policy!r}") + if total_frames == num_frames: + return list(range(total_frames)) + if total_frames > num_frames: + return torch.linspace(0, total_frames - 1, steps=num_frames).round().long().tolist() + reps = int(np.ceil(num_frames / total_frames)) + return (list(range(total_frames)) * reps)[:num_frames] + + +def _image_to_c_h_w(image: Any) -> torch.Tensor: + if isinstance(image, Image.Image): + arr = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0 + return torch.from_numpy(arr).permute(2, 0, 1).contiguous() + if isinstance(image, np.ndarray): + arr = image.astype(np.float32, copy=False) + if arr.max() > 1.0: + arr = arr / 255.0 + tensor = torch.from_numpy(arr) + if tensor.ndim != 3: + raise ValueError(f"Expected image array with 3 dims, got {arr.shape}") + if tensor.shape[-1] == 3: + tensor = tensor.permute(2, 0, 1) + return tensor.float().contiguous() + if torch.is_tensor(image): + tensor = image.detach().float() + if tensor.ndim == 4: + if tensor.shape[0] != 1: + raise ValueError(f"Expected single-frame image tensor, got {tuple(tensor.shape)}") + tensor = tensor.squeeze(0) + if tensor.ndim != 3: + raise ValueError(f"Expected image tensor with 3 dims, got {tuple(tensor.shape)}") + if tensor.shape[-1] == 3: + tensor = tensor.permute(2, 0, 1) + if tensor.max() > 1.0: + tensor = tensor / 255.0 + return tensor.contiguous() + raise TypeError(f"Unsupported image type: {type(image)!r}") + + +def _resize_video(video: torch.Tensor, size: int | tuple[int, int] | None) -> torch.Tensor: + if size is None: + return video.float().clamp(0.0, 1.0).contiguous() + size_hw = (int(size), int(size)) if isinstance(size, int) else (int(size[0]), int(size[1])) + if tuple(video.shape[-2:]) == size_hw: + return video.float().clamp(0.0, 1.0).contiguous() + return F.interpolate(video.float(), size=size_hw, mode="bilinear", align_corners=False).clamp(0.0, 1.0).contiguous() + + +def _normalize(video: torch.Tensor, spec: EncoderPreprocessSpec) -> torch.Tensor: + if spec.normalize_mean is None and spec.normalize_std is None: + return video + mean = torch.tensor(spec.normalize_mean or (0.0, 0.0, 0.0), dtype=video.dtype, device=video.device).view(1, 3, 1, 1) + std = torch.tensor(spec.normalize_std or (1.0, 1.0, 1.0), dtype=video.dtype, device=video.device).view(1, 3, 1, 1) + return (video - mean) / std diff --git a/stimulus_synthesis/scoring/encoder_scorer.py b/stimulus_synthesis/scoring/encoder_scorer.py new file mode 100644 index 0000000000000000000000000000000000000000..11d3ed98e41400a657c1dff610b91bf7b9c19e96 --- /dev/null +++ b/stimulus_synthesis/scoring/encoder_scorer.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from typing import Any + +import torch +from transformers import AutoModel + +from .base import Scorer +from .objectives import build_objective + + +class EncoderScorer(Scorer): + def __init__( + self, + encoder_model_id: str | None = None, + *, + encoder: Any | None = None, + encoder_call: str = "predict_fmri", + objective: str | Any = "indices_mean", + device: str = "cuda", + trust_remote_code: bool = True, + **encoder_kwargs, + ) -> None: + if encoder is None: + if encoder_model_id is None: + raise ValueError("encoder_model_id is required when encoder is not provided.") + encoder = AutoModel.from_pretrained( + encoder_model_id, + trust_remote_code=trust_remote_code, + **encoder_kwargs, + ) + self.encoder = encoder + self.encoder_call = encoder_call + self.objective = build_objective(objective) + self.device = device + if hasattr(self.encoder, "to"): + self.encoder.to(device) + if hasattr(self.encoder, "eval"): + self.encoder.eval() + + def score(self, videos: torch.Tensor, target: Any, **kwargs) -> list[float]: + videos = videos.to(self.device) + with torch.no_grad(): + if self.encoder_call: + fn = getattr(self.encoder, self.encoder_call) + predictions = fn(videos, **kwargs) + else: + predictions = self.encoder(videos, **kwargs) + scores = self.objective(predictions, target) + return [float(x) for x in scores.detach().cpu().reshape(-1)] diff --git a/stimulus_synthesis/scoring/objectives.py b/stimulus_synthesis/scoring/objectives.py new file mode 100644 index 0000000000000000000000000000000000000000..332350afdee53c6dde3039285c1e6acfe73934d9 --- /dev/null +++ b/stimulus_synthesis/scoring/objectives.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from collections.abc import Callable + +import torch +import torch.nn.functional as F + +from .targets import TargetSpec, parse_target + + +def indices_mean(predictions: torch.Tensor, target) -> torch.Tensor: + spec = parse_target(target) + if spec.type != "indices": + raise ValueError("indices_mean objective requires an indices target.") + idx = spec.value.to(predictions.device) + return predictions.index_select(dim=1, index=idx).mean(dim=1) + + +def vector_dot(predictions: torch.Tensor, target) -> torch.Tensor: + spec = parse_target(target) + weights = _target_vector(spec, predictions).to(predictions.device) + return predictions @ weights + + +def vector_cosine(predictions: torch.Tensor, target) -> torch.Tensor: + spec = parse_target(target) + vector = _target_vector(spec, predictions).to(predictions.device) + return F.cosine_similarity(predictions, vector.unsqueeze(0), dim=1) + + +def weighted_mean(predictions: torch.Tensor, target) -> torch.Tensor: + spec = parse_target(target) + weights = _target_vector(spec, predictions).to(predictions.device) + denom = weights.abs().sum().clamp_min(1e-8) + return (predictions * weights.unsqueeze(0)).sum(dim=1) / denom + + +def build_objective(name: str | Callable) -> Callable: + if callable(name): + return name + objectives = { + "indices_mean": indices_mean, + "target_vector_dot": vector_dot, + "vector_dot": vector_dot, + "target_vector_cosine": vector_cosine, + "vector_cosine": vector_cosine, + "weighted_mean": weighted_mean, + } + if name not in objectives: + raise ValueError(f"Unknown objective: {name}") + return objectives[name] + + +def _target_vector(spec: TargetSpec, predictions: torch.Tensor) -> torch.Tensor: + if spec.type in {"vector", "weights"}: + vector = spec.value.float() + if vector.numel() != predictions.shape[1]: + raise ValueError(f"Target vector has {vector.numel()} values, expected {predictions.shape[1]}.") + return vector.reshape(-1) + if spec.type == "indices": + vector = torch.zeros(predictions.shape[1], dtype=predictions.dtype) + vector[spec.value.long()] = 1.0 + return vector + raise ValueError(f"Unsupported target type for vector objective: {spec.type}") diff --git a/stimulus_synthesis/scoring/robust_transform.py b/stimulus_synthesis/scoring/robust_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..fa581c36208888d2e624223c01a8d714ebd4934c --- /dev/null +++ b/stimulus_synthesis/scoring/robust_transform.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import hashlib +from dataclasses import asdict, dataclass +from typing import Any + +import torch +import torch.nn.functional as F + + +@dataclass(frozen=True) +class RobustTransformSpec: + crop_scale: float = 0.80 + gaussian_sigma: float = 0.10 + num_draws: int = 4 + aggregate: str = "mean" + seed: int = 0 + + @classmethod + def from_dict(cls, data: dict[str, Any] | None) -> "RobustTransformSpec | None": + if data is None: + return None + enabled = bool(data.pop("enabled", True)) if "enabled" in data else True + return cls(**data) if enabled else None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +class RobustTransformScorer: + """Apply deterministic robust scoring draws before delegating to a scorer.""" + + def __init__(self, scorer: Any, spec: RobustTransformSpec | None = None) -> None: + self.scorer = scorer + self.spec = spec or RobustTransformSpec() + + def score(self, videos: torch.Tensor, target: Any, **kwargs) -> list[float]: + transformed = apply_robust_transform(videos, self.spec) + raw_scores = self.scorer.score(transformed, target, **kwargs) + score_tensor = torch.as_tensor(raw_scores, dtype=torch.float32).reshape(videos.shape[0], self.spec.num_draws) + if self.spec.aggregate != "mean": + raise ValueError(f"Unsupported robust score aggregate: {self.spec.aggregate!r}") + return score_tensor.mean(dim=1).tolist() + + +def apply_robust_transform(videos: torch.Tensor, spec: RobustTransformSpec | None = None) -> torch.Tensor: + spec = spec or RobustTransformSpec() + if videos.ndim != 5: + raise ValueError(f"Expected videos shaped (B,T,C,H,W), got {tuple(videos.shape)}") + if spec.num_draws <= 0: + raise ValueError("num_draws must be positive.") + if not (0.0 < spec.crop_scale <= 1.0): + raise ValueError("crop_scale must be in (0, 1].") + if spec.gaussian_sigma < 0.0: + raise ValueError("gaussian_sigma must be non-negative.") + + videos = videos.float().clamp(0.0, 1.0) + out = [] + for batch_idx in range(videos.shape[0]): + base_seed = _content_seed(videos[batch_idx], spec.seed) + for draw_idx in range(spec.num_draws): + out.append(_transform_one(videos[batch_idx], spec, base_seed + draw_idx * 7919)) + return torch.stack(out, dim=0).contiguous() + + +def _transform_one(video: torch.Tensor, spec: RobustTransformSpec, seed: int) -> torch.Tensor: + generator = torch.Generator(device=video.device).manual_seed(int(seed) % (2**63 - 1)) + transformed = _random_resized_crop(video, spec.crop_scale, generator) + if spec.gaussian_sigma: + noise = torch.randn( + transformed.shape, + generator=generator, + device=transformed.device, + dtype=transformed.dtype, + ) + transformed = transformed + float(spec.gaussian_sigma) * noise + return transformed.clamp(0.0, 1.0) + + +def _random_resized_crop(video: torch.Tensor, crop_scale: float, generator: torch.Generator) -> torch.Tensor: + if crop_scale == 1.0: + return video + _t, _c, h, w = video.shape + crop_h = max(1, int(round(h * crop_scale))) + crop_w = max(1, int(round(w * crop_scale))) + max_y = h - crop_h + max_x = w - crop_w + y0 = int(torch.randint(max_y + 1, (1,), generator=generator, device=video.device).item()) if max_y else 0 + x0 = int(torch.randint(max_x + 1, (1,), generator=generator, device=video.device).item()) if max_x else 0 + crop = video[:, :, y0 : y0 + crop_h, x0 : x0 + crop_w] + return F.interpolate(crop, size=(h, w), mode="bilinear", align_corners=False) + + +def _content_seed(video: torch.Tensor, seed: int) -> int: + quantized = (video.detach().cpu().clamp(0.0, 1.0) * 255).round().to(torch.uint8).numpy().tobytes() + digest = hashlib.blake2b(quantized, digest_size=8, person=b"nevo-rbt").digest() + return (int.from_bytes(digest, "little") + int(seed)) % (2**63 - 1) diff --git a/stimulus_synthesis/scoring/targets.py b/stimulus_synthesis/scoring/targets.py new file mode 100644 index 0000000000000000000000000000000000000000..eb20b7b4be618992f138f743c13ba3819636d4eb --- /dev/null +++ b/stimulus_synthesis/scoring/targets.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass(frozen=True) +class TargetSpec: + type: str + value: Any + + +def parse_target(target: Any) -> TargetSpec: + if isinstance(target, TargetSpec): + return target + if torch.is_tensor(target): + return TargetSpec("vector", target.float()) + if isinstance(target, dict): + kind = target.get("type") + if kind == "indices": + return TargetSpec("indices", torch.as_tensor(target["indices"], dtype=torch.long)) + if kind == "vector": + return TargetSpec("vector", torch.as_tensor(target["vector"], dtype=torch.float32)) + if kind == "weights": + return TargetSpec("weights", torch.as_tensor(target["weights"], dtype=torch.float32)) + if isinstance(target, (list, tuple)) and all(isinstance(x, int) for x in target): + return TargetSpec("indices", torch.as_tensor(target, dtype=torch.long)) + raise TypeError( + "Unsupported target. Use a TargetSpec, tensor vector, list of indices, or " + "{'type': 'indices'|'vector'|'weights', ...}." + ) diff --git a/stimulus_synthesis/search/__init__.py b/stimulus_synthesis/search/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e31c6da2af5af2169cbd400d7d8506f0add8f224 --- /dev/null +++ b/stimulus_synthesis/search/__init__.py @@ -0,0 +1,5 @@ +from .base import SearchAlgorithm, SearchResult +from .evaluation import EvaluationManager +from .genetic import GeneticSearch + +__all__ = ["SearchAlgorithm", "SearchResult", "EvaluationManager", "GeneticSearch"] diff --git a/stimulus_synthesis/search/base.py b/stimulus_synthesis/search/base.py new file mode 100644 index 0000000000000000000000000000000000000000..55790d3c9718b0705766c8219d9dca130339eb0c --- /dev/null +++ b/stimulus_synthesis/search/base.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any + +from stimulus_synthesis.generators.base import TextToImageGenerator +from stimulus_synthesis.scoring.base import Scorer +from stimulus_synthesis.spaces.base import Candidate, SearchSpace + + +@dataclass +class SearchResult: + best_candidate: Candidate + best_prompt: str + best_score: float + history_best: list[float] + history_overall: list[float] + history_gmean: list[float] = field(default_factory=list) + best_key: str | None = None + best_seed: int | None = None + best_image: Any | None = None + best_video: Any | None = None + best_asset_path: str | None = None + best_export_record: Any | None = None + best_score_record: Any | None = None + best_metadata: dict[str, Any] = field(default_factory=dict) + + +class SearchAlgorithm(ABC): + @abstractmethod + def run( + self, + space: SearchSpace, + text_to_image: TextToImageGenerator, + image_to_video, + scorer: Scorer, + target, + seed: int, + ) -> SearchResult: + ... diff --git a/stimulus_synthesis/search/evaluation.py b/stimulus_synthesis/search/evaluation.py new file mode 100644 index 0000000000000000000000000000000000000000..175b7b0b8e1aefe1e0333a23e1628d62af0e88eb --- /dev/null +++ b/stimulus_synthesis/search/evaluation.py @@ -0,0 +1,222 @@ +from __future__ import annotations + +import re +from pathlib import Path +from typing import Any, Sequence + +import numpy as np +import torch + +from stimulus_synthesis.media import ImageAssetSpec, VideoAssetSpec, export_image, export_video +from stimulus_synthesis.media.normalize import videos_to_b_t_c_h_w +from stimulus_synthesis.scoring.base import Scorer +from stimulus_synthesis.spaces.base import Candidate, SearchSpace + + +class EvaluationManager: + def __init__(self) -> None: + self.score_cache: dict[str, list[float]] = {} + self.image_cache: dict[str, Any] = {} + self.video_cache: dict[str, Any] = {} + self.ind_by_key: dict[str, list[int]] = {} + self.prompt_by_key: dict[str, str] = {} + self.seed_by_key: dict[str, int] = {} + self.asset_export_cache: dict[str, Any] = {} + self.asset_score_cache: dict[str, Any] = {} + self.call_log: list[dict[str, Any]] = [] + + @property + def eval_count(self) -> int: + return sum(len(v) for v in self.score_cache.values()) + + def evaluate( + self, + space: SearchSpace, + text_to_image, + image_to_video, + scorer: Scorer, + target, + inds: Sequence[Sequence[int]], + *, + seed: int, + image_kwargs: dict[str, Any] | None = None, + video_kwargs: dict[str, Any] | None = None, + score_kwargs: dict[str, Any] | None = None, + score_size: int | tuple[int, int] | None = None, + num_frames: int | None = None, + image_batch_size: int | None = None, + video_batch_size: int | None = None, + asset_scorer: Any | None = None, + asset_dir: str | Path | None = None, + asset_type: str | None = None, + image_asset_spec: ImageAssetSpec | None = None, + video_asset_spec: VideoAssetSpec | None = None, + progress=None, + ) -> list[float]: + image_kwargs = image_kwargs or {} + video_kwargs = video_kwargs or {} + score_kwargs = score_kwargs or {} + + entries = [self._entry(space, ind, seed) for ind in inds] + for entry in entries: + self.ind_by_key[entry["key"]] = list(entry["genes"]) + self.prompt_by_key[entry["key"]] = entry["prompt"] + self.seed_by_key[entry["key"]] = int(entry["seed"]) + + to_generate_images = [entry for entry in entries if entry["key"] not in self.image_cache] + img_bs = image_batch_size or len(to_generate_images) or 1 + t2i_batch = hasattr(text_to_image, "generate_batch") + for start in range(0, len(to_generate_images), img_bs): + chunk = to_generate_images[start : start + img_bs] + gens = [_torch_generator(int(e["seed"])) for e in chunk] + if progress is not None: + progress.refresh() + if t2i_batch: + images = text_to_image.generate_batch( + [e["prompt"] for e in chunk], generators=gens, **image_kwargs + ) + else: + images = [ + text_to_image.generate([e["prompt"]], generator=g, **image_kwargs)[0] + for e, g in zip(chunk, gens) + ] + for e, image in zip(chunk, images): + self.image_cache[e["key"]] = image + + to_generate_videos = [entry for entry in entries if entry["key"] not in self.video_cache] + vid_bs = video_batch_size or len(to_generate_videos) or 1 + i2v_batch = hasattr(image_to_video, "generate_batch") + for start in range(0, len(to_generate_videos), vid_bs): + chunk = to_generate_videos[start : start + vid_bs] + imgs = [self.image_cache[e["key"]] for e in chunk] + gens = [_torch_generator(int(e["seed"]) + 10_000_000) for e in chunk] + if progress is not None: + progress.refresh() + if i2v_batch: + videos = image_to_video.generate_batch( + imgs, [e["prompt"] for e in chunk], generators=gens, **video_kwargs + ) + else: + videos = [ + image_to_video.generate(img, e["prompt"], generator=g, **video_kwargs) + for img, e, g in zip(imgs, chunk, gens) + ] + for e, video in zip(chunk, videos): + self.video_cache[e["key"]] = video + + if progress is not None: + progress.refresh() + before = {key: len(values) for key, values in self.score_cache.items()} + if asset_scorer is not None: + raw_scores = self._score_saved_assets(entries, asset_scorer, asset_dir, asset_type, image_asset_spec, video_asset_spec) + else: + videos = [self.video_cache[entry["key"]] for entry in entries] + video_tensor = videos_to_b_t_c_h_w(videos, size=score_size, num_frames=num_frames) + raw_scores = scorer.score(video_tensor, target, **score_kwargs) + + scores = [] + for entry, raw_score in zip(entries, raw_scores): + key = entry["key"] + self.score_cache.setdefault(key, []).append(float(raw_score)) + scores.append(float(np.mean(self.score_cache[key]))) + self._log_new_calls(before) + return scores + + def best_prompt(self) -> tuple[str | None, float | None]: + best_key, best_score = self.best_key() + if best_key is None: + return None, None + return self.prompt_by_key.get(best_key), best_score + + def best_key(self) -> tuple[str | None, float | None]: + best_key, best_score = None, None + for key, values in self.score_cache.items(): + score = float(np.mean(values)) + if best_score is None or score > best_score: + best_key, best_score = key, score + return best_key, best_score + + def _entry(self, space: SearchSpace, ind: Sequence[int], default_seed: int) -> dict[str, Any]: + cand = Candidate(tuple(ind)) + prompt = space.decode(cand) + if hasattr(space, "decode_seed"): + cand_seed = int(space.decode_seed(cand, default=default_seed)) + else: + cand_seed = int(default_seed) + return { + "genes": list(ind), + "prompt": prompt, + "seed": cand_seed, + "key": candidate_key(prompt, cand_seed), + } + + def _score_saved_assets( + self, + entries: Sequence[dict[str, Any]], + asset_scorer: Any, + asset_dir: str | Path | None, + asset_type: str | None, + image_asset_spec: ImageAssetSpec | None, + video_asset_spec: VideoAssetSpec | None, + ) -> list[float]: + if asset_type not in {"image", "video"}: + raise ValueError("asset_type must be 'image' or 'video' when asset_scorer is set.") + if asset_dir is None: + raise ValueError("asset_dir is required when asset_scorer is set.") + asset_root = Path(asset_dir) + asset_root.mkdir(parents=True, exist_ok=True) + scores = [] + for entry in entries: + key = entry["key"] + if key not in self.asset_score_cache: + stem = _safe_stem(key) + metadata = {"prompt": entry["prompt"], "seed": int(entry["seed"]), "genes": list(entry["genes"]), "key": key} + if asset_type == "image": + spec = image_asset_spec + if spec is None: + raise ValueError("image_asset_spec is required for image asset scoring.") + path = asset_root / f"{stem}.png" + export_record = export_image(self.image_cache[key], path, spec) + score_record = asset_scorer.score_image(path, asset_spec=spec, metadata=metadata) + else: + spec = video_asset_spec + if spec is None: + raise ValueError("video_asset_spec is required for video asset scoring.") + path = asset_root / f"{stem}.mp4" + export_record = export_video(self.video_cache[key], path, spec) + score_record = asset_scorer.score_video(path, asset_spec=spec, metadata=metadata) + self.asset_export_cache[key] = export_record + self.asset_score_cache[key] = score_record + scores.append(float(self.asset_score_cache[key].score)) + return scores + + def _log_new_calls(self, before: dict[str, int]) -> None: + for key, values in self.score_cache.items(): + prev = before.get(key, 0) + for j in range(prev, len(values)): + self.call_log.append( + { + "call_index": len(self.call_log), + "key": key, + "prompt": self.prompt_by_key.get(key, ""), + "seed": self.seed_by_key.get(key), + "genes": [int(g) for g in self.ind_by_key.get(key, [])], + "raw_score": float(values[j]), + "cached_mean_score": float(np.mean(values[: j + 1])), + "n_evals_for_candidate": j + 1, + } + ) + + +def candidate_key(prompt: str, seed: int | None) -> str: + return f"{prompt}::seed={int(seed) if seed is not None else 'none'}" + + +def _torch_generator(seed: int): + device = "cuda" if torch.cuda.is_available() else "cpu" + return torch.Generator(device=device).manual_seed(int(seed)) + + +def _safe_stem(key: str) -> str: + stem = re.sub(r"[^A-Za-z0-9_.-]+", "_", key).strip("._") + return stem[:160] or "candidate" diff --git a/stimulus_synthesis/search/genetic.py b/stimulus_synthesis/search/genetic.py new file mode 100644 index 0000000000000000000000000000000000000000..21fdf9b71f73b825de62e922a36a8c7cf7bece0f --- /dev/null +++ b/stimulus_synthesis/search/genetic.py @@ -0,0 +1,271 @@ +from __future__ import annotations + +import random +from pathlib import Path +from typing import Any + +import numpy as np + +from stimulus_synthesis.media import ImageAssetSpec, VideoAssetSpec +from stimulus_synthesis.spaces.base import Candidate, SearchSpace +from stimulus_synthesis.utils.seed import set_all_seeds + +from .base import SearchAlgorithm, SearchResult +from .evaluation import EvaluationManager, candidate_key + + +class GeneticSearch(SearchAlgorithm): + def __init__( + self, + *, + max_evals: int = 200, + population_size: int = 16, + n_init: int | None = None, + mutation_rate: float = 0.25, + crossover_rate: float = 0.5, + elite_frac: float = 0.35, + image_kwargs: dict[str, Any] | None = None, + video_kwargs: dict[str, Any] | None = None, + score_kwargs: dict[str, Any] | None = None, + score_size: int | tuple[int, int] | None = 224, + num_frames: int | None = None, + asset_scorer: Any | None = None, + asset_dir: str | Path | None = None, + asset_type: str | None = None, + image_asset_spec: ImageAssetSpec | None = None, + video_asset_spec: VideoAssetSpec | None = None, + image_batch_size: int | None = None, + video_batch_size: int | None = None, + show_progress: bool = False, + progress_desc: str = "NEvo search", + ) -> None: + self.max_evals = int(max_evals) + self.population_size = int(population_size) + self.n_init = int(n_init or population_size) + self.mutation_rate = float(mutation_rate) + self.crossover_rate = float(crossover_rate) + self.elite_frac = float(elite_frac) + self.image_kwargs = image_kwargs or {} + self.video_kwargs = video_kwargs or {} + self.score_kwargs = score_kwargs or {} + self.score_size = score_size + self.num_frames = num_frames + self.asset_scorer = asset_scorer + self.asset_dir = asset_dir + self.asset_type = asset_type + self.image_asset_spec = image_asset_spec + self.video_asset_spec = video_asset_spec + self.show_progress = bool(show_progress) + self.progress_desc = str(progress_desc) + self.image_batch_size = image_batch_size + self.video_batch_size = video_batch_size + + def run(self, space: SearchSpace, text_to_image, image_to_video, scorer, target, seed: int) -> SearchResult: + set_all_seeds(seed) + evaluator = EvaluationManager() + seen: set[str] = set() + history_best, history_mean = [], [] + + progress = None + if self.show_progress: + try: + from tqdm.auto import tqdm + progress = tqdm(total=self.max_evals, initial=0, desc=self.progress_desc, unit="eval") + except Exception: + progress = None + + population = self._make_random(space, seen, self.n_init, seed) + self._evaluate(evaluator, space, text_to_image, image_to_video, scorer, target, population, seed, progress=progress) + + _best_key, best_score = evaluator.best_key() + running_best = float(best_score if best_score is not None else -np.inf) + init_mean = self._mean_score(evaluator) + history_best = [running_best] * evaluator.eval_count + history_mean = [init_mean] * evaluator.eval_count + if progress is not None: + progress.update(evaluator.eval_count - progress.n) + progress.set_postfix(best=f"{running_best:.4f}") + + while evaluator.eval_count < self.max_evals: + remaining = self.max_evals - evaluator.eval_count + if remaining <= 0: + break + + scored = self._score_population(space, population, evaluator, seed) + n_elite = max(1, int(round(len(scored) * self.elite_frac))) + elites = [ind for _score, ind in scored[:n_elite]] + + next_gen = list(elites) + attempts = 0 + while len(next_gen) < self.population_size and attempts < self.population_size * 50: + attempts += 1 + if evaluator.eval_count >= self.max_evals: + break + child = self._breed_one(space, elites, seen, evaluator, seed) + if child is not None: + next_gen.append(child) + + to_eval = self._unevaluated(space, next_gen, evaluator, seed)[:remaining] + if not to_eval: + break + + self._evaluate(evaluator, space, text_to_image, image_to_video, scorer, target, to_eval, seed, progress=progress) + population = next_gen + + _current_key, current_best = evaluator.best_key() + if current_best is not None: + running_best = max(running_best, float(current_best)) + new_evals = evaluator.eval_count - len(history_best) + gen_mean = self._mean_score_for_inds(space, to_eval, evaluator, seed) + history_best.extend([running_best] * max(new_evals, 1)) + history_mean.extend([gen_mean] * max(new_evals, 1)) + if progress is not None: + _pk, _pb = evaluator.best_key() + progress.update(evaluator.eval_count - progress.n) + progress.set_postfix(best=f"{float(_pb):.4f}" if _pb is not None else "n/a") + + if progress is not None: + progress.close() + + best_key, best_score = evaluator.best_key() + if best_key is None or best_score is None: + raise RuntimeError("No candidate was evaluated.") + best_prompt = evaluator.prompt_by_key[best_key] + best_genes = evaluator.ind_by_key[best_key] + export_record = evaluator.asset_export_cache.get(best_key) + score_record = evaluator.asset_score_cache.get(best_key) + return SearchResult( + best_candidate=Candidate(tuple(best_genes)), + best_prompt=best_prompt, + best_score=float(best_score), + history_best=history_best, + history_overall=history_best, + history_gmean=history_mean, + best_key=best_key, + best_seed=evaluator.seed_by_key.get(best_key), + best_image=evaluator.image_cache.get(best_key), + best_video=evaluator.video_cache.get(best_key), + best_asset_path=getattr(export_record, "path", None), + best_export_record=export_record, + best_score_record=score_record, + best_metadata={ + "key": best_key, + "seed": evaluator.seed_by_key.get(best_key), + "genes": [int(g) for g in best_genes], + "score_source": "saved_asset" if score_record is not None else "memory", + }, + ) + + def _evaluate(self, evaluator, space, text_to_image, image_to_video, scorer, target, inds, seed, progress=None): + evaluator.evaluate( + space, + text_to_image, + image_to_video, + scorer, + target, + inds, + seed=seed, + image_kwargs=self.image_kwargs, + video_kwargs=self.video_kwargs, + score_kwargs=self.score_kwargs, + score_size=self.score_size, + num_frames=self.num_frames, + image_batch_size=self.image_batch_size, + video_batch_size=self.video_batch_size, + asset_scorer=self.asset_scorer, + asset_dir=self.asset_dir, + asset_type=self.asset_type, + image_asset_spec=self.image_asset_spec, + video_asset_spec=self.video_asset_spec, + progress=progress, + ) + + def _make_random(self, space: SearchSpace, seen: set[str], n: int, default_seed: int) -> list[list[int]]: + out = [] + attempts = 0 + while len(out) < n and attempts < n * 50: + attempts += 1 + ind = list(space.random_candidate().genes) + key = self._key(space, ind, default_seed) + if key in seen: + continue + seen.add(key) + out.append(ind) + return out + + def _breed_one( + self, + space: SearchSpace, + elites: list[list[int]], + seen: set[str], + evaluator: EvaluationManager, + default_seed: int, + ) -> list[int] | None: + if not elites: + return None + parent_a = random.choice(elites) + if random.random() < self.crossover_rate and len(elites) > 1: + parent_b = random.choice(elites) + child = [parent_a[i] if random.random() < 0.5 else parent_b[i] for i in range(len(parent_a))] + else: + child = list(parent_a) + + for i in range(len(child)): + if random.random() < self.mutation_rate: + if hasattr(space, "active_mask") and not space.active_mask[i]: + child[i] = space.empty_indices[i] + else: + child[i] = random.randrange(len(space.options[i])) + + key = self._key(space, child, default_seed) + if key in seen or key in evaluator.score_cache: + return None + seen.add(key) + evaluator.ind_by_key[key] = list(child) + return child + + @staticmethod + def _key(space: SearchSpace, ind: list[int] | tuple[int, ...], default_seed: int) -> str: + cand = Candidate(tuple(ind)) + prompt = space.decode(cand) + cand_seed = int(space.decode_seed(cand, default=default_seed)) if hasattr(space, "decode_seed") else int(default_seed) + return candidate_key(prompt, cand_seed) + + @classmethod + def _score_population( + cls, + space: SearchSpace, + population: list[list[int]], + evaluator: EvaluationManager, + default_seed: int, + ) -> list[tuple[float, list[int]]]: + scored = [] + for ind in population: + key = cls._key(space, ind, default_seed) + score = float(np.mean(evaluator.score_cache.get(key, [0.0]))) + scored.append((score, ind)) + return sorted(scored, key=lambda x: x[0], reverse=True) + + @classmethod + def _unevaluated( + cls, + space: SearchSpace, + population: list[list[int]], + evaluator: EvaluationManager, + default_seed: int, + ) -> list[list[int]]: + return [ind for ind in population if cls._key(space, ind, default_seed) not in evaluator.score_cache] + + @staticmethod + def _mean_score(evaluator: EvaluationManager) -> float: + values = [score for scores in evaluator.score_cache.values() for score in scores] + return float(np.mean(values)) if values else float("nan") + + @classmethod + def _mean_score_for_inds(cls, space: SearchSpace, inds: list[list[int]], evaluator: EvaluationManager, default_seed: int) -> float: + scores = [] + for ind in inds: + key = cls._key(space, ind, default_seed) + if key in evaluator.score_cache: + scores.append(float(np.mean(evaluator.score_cache[key]))) + return float(np.mean(scores)) if scores else float("nan") diff --git a/stimulus_synthesis/spaces/__init__.py b/stimulus_synthesis/spaces/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0fac7d6659e157f322a568966f9eb83ff620b8bf --- /dev/null +++ b/stimulus_synthesis/spaces/__init__.py @@ -0,0 +1,6 @@ +from .base import Candidate, SearchSpace +from .prompt_space import PromptSearchSpace +from .seeded import SeededSearchSpace +from .structured_neuro_space import StructuredArtPromptSpace, VideoMotionPromptSpace, make_t2v_art_data + +__all__ = ["Candidate", "SearchSpace", "PromptSearchSpace", "SeededSearchSpace", "StructuredArtPromptSpace", "VideoMotionPromptSpace", "make_t2v_art_data"] diff --git a/stimulus_synthesis/spaces/base.py b/stimulus_synthesis/spaces/base.py new file mode 100644 index 0000000000000000000000000000000000000000..f0a770ce8429ea53f18e173b7463cb81ec4a1570 --- /dev/null +++ b/stimulus_synthesis/spaces/base.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, Sequence + + +@dataclass(frozen=True) +class Candidate: + genes: Sequence[int] + + +class SearchSpace(ABC): + @abstractmethod + def random_candidate(self) -> Candidate: + ... + + @abstractmethod + def decode(self, cand: Candidate) -> Any: + ... + + @abstractmethod + def mutate(self, cand: Candidate, mutation_rate: float) -> Candidate: + ... + + @abstractmethod + def crossover(self, a: Candidate, b: Candidate, crossover_rate: float) -> tuple[Candidate, Candidate]: + ... diff --git a/stimulus_synthesis/spaces/prompt_banks.py b/stimulus_synthesis/spaces/prompt_banks.py new file mode 100644 index 0000000000000000000000000000000000000000..b02d2754134e3111b6bb7bca6a878bd8ba3813af --- /dev/null +++ b/stimulus_synthesis/spaces/prompt_banks.py @@ -0,0 +1,36 @@ +DEFAULT_PROMPT_BANKS = { + "subject": [ + "a person", + "two people", + "a face", + "a moving animal", + "a vehicle", + "geometric shapes", + "", + ], + "action": [ + "running", + "walking", + "turning around", + "jumping", + "reaching toward the camera", + "moving in slow motion", + "", + ], + "scene": [ + "in a crowded street", + "in a forest", + "inside a room", + "on a beach", + "against a simple background", + "", + ], + "style": [ + "photorealistic", + "cinematic", + "natural video", + "high contrast", + "soft daylight", + "", + ], +} diff --git a/stimulus_synthesis/spaces/prompt_space.py b/stimulus_synthesis/spaces/prompt_space.py new file mode 100644 index 0000000000000000000000000000000000000000..1e7638c98cd4ee265cf1c46daa14976d8d07d7ed --- /dev/null +++ b/stimulus_synthesis/spaces/prompt_space.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import random +from collections.abc import Mapping, Sequence + +from .base import Candidate, SearchSpace +from .prompt_banks import DEFAULT_PROMPT_BANKS + + +class PromptSearchSpace(SearchSpace): + def __init__(self, prompt_banks: Mapping[str, Sequence[str]] | None = None, seed_prompts: Sequence[str] | None = None): + banks = dict(prompt_banks or DEFAULT_PROMPT_BANKS) + if seed_prompts: + banks = {"seed": list(seed_prompts), **banks} + self.categories = list(banks.keys()) + self.options = [list(banks[c]) for c in self.categories] + if any(len(opts) == 0 for opts in self.options): + raise ValueError("Prompt banks must not contain empty option lists.") + + def random_candidate(self) -> Candidate: + return Candidate(tuple(random.randrange(len(opts)) for opts in self.options)) + + def decode(self, cand: Candidate) -> str: + parts = [] + for gene, opts in zip(cand.genes, self.options): + text = opts[int(gene)] + if text: + parts.append(text) + return " ".join(parts).strip() + + def mutate(self, cand: Candidate, mutation_rate: float) -> Candidate: + genes = list(cand.genes) + for i, opts in enumerate(self.options): + if random.random() < mutation_rate: + genes[i] = random.randrange(len(opts)) + return Candidate(tuple(genes)) + + def crossover(self, a: Candidate, b: Candidate, crossover_rate: float) -> tuple[Candidate, Candidate]: + if len(a.genes) != len(b.genes): + raise ValueError("Candidates must have the same gene length.") + if random.random() >= crossover_rate or len(a.genes) < 2: + return a, b + cut = random.randrange(1, len(a.genes)) + child_a = Candidate(tuple(a.genes[:cut]) + tuple(b.genes[cut:])) + child_b = Candidate(tuple(b.genes[:cut]) + tuple(a.genes[cut:])) + return child_a, child_b diff --git a/stimulus_synthesis/spaces/seeded.py b/stimulus_synthesis/spaces/seeded.py new file mode 100644 index 0000000000000000000000000000000000000000..9e8f2f3edc0cc9006f335f960185a47a33db4dea --- /dev/null +++ b/stimulus_synthesis/spaces/seeded.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import random +from collections.abc import Sequence + +from .base import Candidate, SearchSpace + + +class SeededSearchSpace(SearchSpace): + """Append a generation seed as the last categorical gene of a search space.""" + + def __init__(self, base_space: SearchSpace, seed_values: Sequence[int]): + if not seed_values: + raise ValueError("seed_values must not be empty.") + self.base_space = base_space + self.seed_values = [int(seed) for seed in seed_values] + self.options = list(getattr(base_space, "options")) + [list(range(len(self.seed_values)))] + if hasattr(base_space, "categories"): + self.categories = list(getattr(base_space, "categories")) + ["seed"] + if hasattr(base_space, "active_mask"): + self.active_mask = list(getattr(base_space, "active_mask")) + [True] + if hasattr(base_space, "empty_indices"): + self.empty_indices = list(getattr(base_space, "empty_indices")) + [0] + + def random_candidate(self) -> Candidate: + base = list(self.base_space.random_candidate().genes) + base.append(random.randrange(len(self.seed_values))) + return Candidate(tuple(base)) + + def decode(self, cand: Candidate) -> str: + return self.base_space.decode(Candidate(tuple(cand.genes[:-1]))) + + def decode_seed(self, cand: Candidate, default: int | None = None) -> int: + if len(cand.genes) == 0: + if default is None: + raise ValueError("Cannot decode a seed from an empty candidate.") + return int(default) + idx = int(cand.genes[-1]) + if idx < 0 or idx >= len(self.seed_values): + raise IndexError(f"Seed gene index {idx} is outside seed_values.") + return self.seed_values[idx] + + def base_candidate(self, cand: Candidate) -> Candidate: + return Candidate(tuple(cand.genes[:-1])) + + def mutate(self, cand: Candidate, mutation_rate: float) -> Candidate: + genes = list(cand.genes) + for i, opts in enumerate(self.options): + if random.random() < mutation_rate: + if hasattr(self, "active_mask") and not self.active_mask[i]: + genes[i] = self.empty_indices[i] + else: + genes[i] = random.randrange(len(opts)) + return Candidate(tuple(genes)) + + def crossover(self, a: Candidate, b: Candidate, crossover_rate: float) -> tuple[Candidate, Candidate]: + if len(a.genes) != len(b.genes): + raise ValueError("Candidates must have the same gene length.") + if random.random() >= crossover_rate or len(a.genes) < 2: + return a, b + child_a = [ga if random.random() < 0.5 else gb for ga, gb in zip(a.genes, b.genes)] + child_b = [gb if random.random() < 0.5 else ga for ga, gb in zip(a.genes, b.genes)] + return Candidate(tuple(child_a)), Candidate(tuple(child_b)) diff --git a/stimulus_synthesis/spaces/structured_neuro_space.py b/stimulus_synthesis/spaces/structured_neuro_space.py new file mode 100644 index 0000000000000000000000000000000000000000..7da4a6b761ad72a80bc569f5b50ff8ef657366b4 --- /dev/null +++ b/stimulus_synthesis/spaces/structured_neuro_space.py @@ -0,0 +1,832 @@ +from __future__ import annotations +import random +import numpy as np +import torch +from typing import Dict, List, Optional + +from .base import Candidate, SearchSpace + +# ── art_data: original aesthetic dims + neuroscience dims ──────────────────── +# "" = don't specify (always last entry per category). +# Aesthetic dims: SDXL-fluent language for image quality anchoring +# Neuro dims: content-specific for ROI targeting + +art_data = { + + # ── Original aesthetic dims (SDXL-fluent) ──────────────────────────────── + + "RenderType": [ + "photorealistic", + "cinematic still", + "national geographic photo", + "extreme macro photography", + "drone orthographic view", + "35mm film", + "digital concept art", + "dark fantasy concept art", + "unreal engine render", + "double exposure", + "fashion editorial", + "infrared photography", + "CCTV security footage", + "underwater photography", + "medical anatomy illustration", + "", + ], + + "Subject": [ + "portrait face", + "direct eye contact", + "smiling face", + "fearful wide eyes", + "angry bared teeth", + "embracing couple", + "mother and newborn", + "hostile crowd", + "two figures fighting", + "dancer mid-motion", + "sprinting athlete", + "hands reaching toward viewer", + "falling figure", + "body silhouette", + "predator eyes close-up", + "snake striking", + "long dark corridor", + "spiral staircase", + "crashing ocean wave", + "volcanic eruption", + "eye iris extreme macro", + "neural pathways brain", + "high contrast checkerboard", + "radial starburst pattern", + "sacred geometry mandala", + "burning fire", + "", + ], + + "Lighting": [ + "dramatic chiaroscuro", + "golden hour warm", + "harsh midday shadow", + "rim lighting edge", + "backlit silhouette", + "volumetric god rays", + "neon light", + "strobe freeze frame", + "softbox studio", + "single spotlight", + "moonlight cold blue", + "underlighting eerie", + "natural window light", + "dappled forest light", + "emergency red light", + "", + ], + + "Mood": [ + "sublime overwhelming awe", + "uncanny valley dread", + "liminal eerie unease", + "euphoric transcendent joy", + "primal visceral fear", + "serene meditative calm", + "epic cinematic grandeur", + "deep melancholy sorrow", + "ominous threatening", + "suspenseful tense", + "disgust revulsion", + "acute danger urgent", + "hypnotic trance", + "grotesque body horror", + "overwhelming sensory overload", + "", + ], + + "Color": [ + "maximum vivid saturation", + "muted desaturated tones", + "monochromatic single hue", + "sepia aged warm", + "teal and orange cinematic", + "complementary color contrast", + "ultra-black deep shadows", + "neon electric palette", + "earthy ochre brown", + "dominant red", + "dominant blue", + "dominant green", + "dominant yellow", + "black and white", + "pastel soft", + "jewel tones rich", + "", + ], + + "Style": [ + "surrealist dreamscape", + "high renaissance realism", + "cyberpunk noir", + "vaporwave aesthetic", + "Bauhaus geometric", + "baroque theatrical", + "psychedelic", + "hyperrealism", + "pop art", + "street art graffiti", + "impressionist painterly", + "minimalist clean", + "gothic dark", + "futurist dynamic", + "", + ], + + "Framing": [ + "centered subject", + "extreme close-up", + "wide establishing shot", + "dutch angle tilt", + "vanishing point perspective", + "bird's eye view", + "worm's eye view", + "first-person subjective", + "over the shoulder", + "reflection in water", + "fisheye distortion", + "", + ], + + "Pattern": [ + "tessellated hexagons", + "recursive mandala", + "circuitry PCB traces", + "concentric ripple rings", + "fractal tree branching", + "sine wave interference", + "geometric grid tiling", + "high contrast checkerboard", + "vertical bar grating", + "horizontal bar grating", + "diagonal stripe pattern", + "random dot noise", + "op art illusion", + "", + ], + + "Surface": [ + "wet glistening skin", + "cold brushed metal", + "rough weathered stone", + "soft velvet fabric", + "cracked dry earth", + "gossamer silk sheer", + "translucent membrane", + "heavy film grain", + "liquid mercury chrome", + "flowing molten lava", + "wispy smoke ethereal", + "subsurface scatter skin", + "polished white bone", + "rusted corroded metal", + "", + ], + + "ScenePhysics": [ + "completely frozen still", + "slow drift motion", + "explosive sudden burst", + "slow motion", + "frozen mid-air suspended", + "swirling fluid vortex", + "impact collision moment", + "falling descent", + "chaotic turbulence", + "rhythmic pulsing", + "time lapse compressed", + "gravity defying float", + "", + ], + + # ── Neuroscience dims ───────────────────────────────────────────────────── + + "FaceContent": [ + "close-up portrait face", + "side profile face", + "three-quarter face", + "upside-down face", + "eyes closed face", + "direct eye contact", + "looking away", + "baby face", + "elderly face", + "dog face", + "cat face", + "monkey face", + "cartoon face", + "mannequin face", + "neutral expression", + "angry face", + "fearful face", + "smiling face", + "surprised face", + "disgusted face", + "crying face", + "multiple faces", + "no face", + "", + ], + + "FaceViewpoint": [ + "direct eye contact", + "slightly looking away", + "looking to the side", + "looking down", + "looking up", + "eyes closed", + "wide open eyes", + "squinting eyes", + "", + ], + + "BodyContent": [ + "full body portrait", + "full body side view", + "full body from behind", + "torso and arms", + "legs and feet", + "hands close-up", + "feet close-up", + "stick figure", + "body outline", + "headless body", + "body in action", + "resting body", + "body in shadow", + "no body", + "", + ], + + "BodyAction": [ + "standing still", + "reaching toward viewer", + "reaching for object", + "walking", + "running", + "jumping", + "falling", + "crouching", + "punching or kicking", + "hugging", + "pointing", + "waving", + "throwing", + "catching", + "pushing", + "pulling", + "dancing", + "crawling", + "sitting", + "lying down", + "", + ], + + "AgentType": [ + "adult person", + "baby or toddler", + "chimpanzee or monkey", + "dog", + "cat", + "bird", + "moving triangle", + "moving circle", + "moving square", + "stick figure", + "robot", + "hand only", + "masked person", + "silhouette person", + "no agent", + "", + ], + + "SocialCue": [ + "eye contact", + "both looking at same object", + "pointing at something", + "reaching toward person", + "turning toward person", + "turning away", + "mirroring posture", + "open arms welcoming", + "crossed arms closed", + "protective gesture", + "beckoning come here", + "pushing away rejection", + "bowing submissive", + "looming dominant", + "no social cue", + "", + ], + + "SocialDynamics": [ + "working together", + "competing against each other", + "one person dominant", + "threatening aggressive", + "playing together", + "parent and child", + "one chasing other", + "taking turns", + "moving in sync", + "ignoring each other", + "hugging affectionate", + "fighting", + "teaching showing", + "imitating copying", + "greeting handshake", + "walking away leaving", + "alone single person", + "", + ], + + "IntentionSignal": [ + "clearly reaching for goal", + "unclear ambiguous action", + "trying and failing", + "accidental slip", + "deliberately misleading", + "helping another person", + "blocking another person", + "asking requesting", + "offering giving", + "no agent", + "", + ], + + "AgentNumber": [ + "one person", + "two people", + "three people", + "small group", + "crowd of people", + "no people", + "", + ], + + "SceneCategory": [ + "living room", + "outdoor landscape", + "city street", + "forest", + "beach", + "kitchen", + "bedroom", + "office", + "hallway corridor", + "rooftop", + "cave", + "mountain", + "desert", + "underwater", + "parking lot", + "subway station", + "no scene", + "", + ], + + "SceneGeometry": [ + "deep perspective", + "flat surface view", + "low ceiling room", + "wide open panorama", + "bird eye view", + "looking up view", + "cluttered room", + "empty sparse room", + "narrow corridor", + "doorway framing", + "window framing", + "symmetrical room", + "staircase", + "", + ], + + "ObjectForm": [ + "recognizable everyday object", + "abstract geometric object", + "organic natural object", + "tool or instrument", + "household item", + "food item", + "vehicle", + "furniture", + "clothing item", + "plant or flower", + "transparent glass object", + "reflective shiny object", + "broken damaged object", + "single isolated object", + "multiple objects", + "", + ], + + "Texture": [ + "fine detailed texture", + "rough coarse surface", + "smooth glossy surface", + "human skin", + "flat solid color", + "fur", + "fabric", + "wood grain", + "stone surface", + "brushed metal", + "water ripples", + "sand", + "tree bark", + "reptile scales", + "feathers", + "velvet", + "", + ], + + "ColorOpponency": [ + "red and green", + "blue and yellow", + "black and white", + "vivid saturated colors", + "dominant red", + "dominant blue", + "dominant green", + "dominant yellow", + "dominant orange", + "dominant purple", + "soft pastel colors", + "warm tones", + "cool tones", + "complementary colors", + "neon colors", + "earthy muted tones", + "", + ], +} + + +# ── ROI β†’ active categories ─────────────────────────────────────────────────── +# Each ROI gets ~8-10 active cats: neuro content dims + aesthetic anchors. +# Aesthetic anchors (RenderType, Lighting, Mood, Style) ensure SDXL generates +# high-quality images regardless of which content dims fire. + +ROI_ACTIVE_CATEGORIES: Dict[str, List[str]] = { + "V1": [ + "Pattern", "Color", "Surface", "ScenePhysics", + "RenderType", "Style", + ], + "V2": [ + "Pattern", "Color", "Surface", "ScenePhysics", + "RenderType", "Style", + ], + "V4": [ + "ColorOpponency", "Texture", "ObjectForm", "Pattern", + "Color", "RenderType", "Style", "Lighting", + ], + "hV4": [ + "ColorOpponency", "Texture", "ObjectForm", "Pattern", + "Color", "RenderType", "Style", "Lighting", + ], + "LO": [ + "ObjectForm", "Texture", "Color", + "RenderType", "Lighting", "Style", "Framing", + ], + "OFA": [ + "FaceContent", "FaceViewpoint", "ObjectForm", + "RenderType", "Lighting", "Mood", "Framing", + ], + "FFA": [ + "FaceContent", "FaceViewpoint", "SocialCue", "AgentNumber", + "RenderType", "Lighting", "Mood", "Framing", + ], + "FFA-1": [ + "FaceContent", "FaceViewpoint", "SocialCue", "AgentNumber", + "RenderType", "Lighting", "Mood", "Framing", + ], + "FFA-2": [ + "FaceContent", "FaceViewpoint", "SocialCue", "AgentNumber", + "RenderType", "Lighting", "Mood", "Framing", + ], + "EBA": [ + "BodyContent", "BodyAction", "AgentNumber", + "RenderType", "Lighting", "Mood", "Framing", "Subject", + ], + "PPA": [ + "SceneCategory", "SceneGeometry", "ObjectForm", + "RenderType", "Lighting", "Mood", "Style", + ], + "OPA": [ + "SceneGeometry", "SceneCategory", + "RenderType", "Lighting", "Framing", "Style", + ], + "pSTS": [ + "AgentType", "SocialCue", "SocialDynamics", "IntentionSignal", + "AgentNumber", "BodyContent", "BodyAction", + "RenderType", "Lighting", "Mood", + ], + "aSTS": [ + "AgentType", "SocialDynamics", "IntentionSignal", "SocialCue", + "AgentNumber", "FaceContent", + "RenderType", "Lighting", "Mood", + ], + "STS": [ + "AgentType", "SocialCue", "SocialDynamics", "IntentionSignal", + "AgentNumber", "BodyContent", "FaceContent", + "RenderType", "Lighting", "Mood", + ], +} + +_ALL_CATEGORIES = list(art_data.keys()) + +# ── Video motion search space (for I2V mode) ────────────────────────────────── +# Used when an anchor image provides content β€” only motion/interaction dims needed. + +VIDEO_MOTION_ART_DATA: Dict[str, List[str]] = { + "MotionProfile": [ + "biological motion", "locomotor motion", "fine motor motion", + "mechanical motion", "fluid motion", "particle motion", + "oscillatory motion", "turbulent motion", "erratic motion", + "elastic motion", "drifting motion", "accelerating motion", + "intermittent motion", "crowd motion", "", + ], + "MotionStrength": [ + "very subtle motion", "gentle motion", "moderate motion", + "strong motion", "very fast intense motion", "", + ], + "MotionScale": [ + "local body part motion", "single entity motion", "two entities moving", + "whole scene motion", "camera dominant motion", "", + ], + "PrimaryAction": [ + "walking", "running", "jumping", "dancing", "fighting", "reaching", + "grasping", "rotating", "bouncing", "sliding", "flowing", "scattering", + "transforming", "vibrating", "falling", "rising", "spinning", + "crawling", "swaying", "", + ], + "Interaction": [ + "mutual attention", "joint attention toward object", "imitation mirroring", + "cooperation joint action", "competition opposing", "caregiving nurturing", + "confrontation threatening", "pursuit chasing", "avoidance fleeing", + "touching contact", "passing object between", "approaching", "departing", + "orbiting circling", "colliding", "no interaction", "", + ], + "SocialContingency": [ + "no contingency", "weak contingency", "clear contingency", + "reciprocal contingency", "turn taking", "synchronized movement", "", + ], + "EventStructure": [ + "ongoing continuous process", "single event", "repeated cyclic event", + "approach contact depart", "pursuit sequence", "cause effect chain", + "escalating sequence", "parallel simultaneous events", + "transformation sequence", "", + ], + "CameraMotion": [ + "static camera", "slow pan", "tracking follow", "zoom in", + "zoom out", "orbiting", "handheld shaky", "", + ], + "TemporalRhythm": [ + "smooth continuous", "staccato sudden bursts", "rhythmic periodic", + "chaotic irregular", "slow then fast", "fast then slow", "", + ], +} + +VIDEO_MOTION_ROI_ACTIVE: Dict[str, List[str]] = { + "pSTS": [ + "MotionProfile", "PrimaryAction", "Interaction", + "SocialContingency", "EventStructure", "MotionStrength", "TemporalRhythm", + ], + "aSTS": [ + "Interaction", "SocialContingency", "EventStructure", + "PrimaryAction", "MotionStrength", "TemporalRhythm", + ], + "EBA": [ + "MotionProfile", "PrimaryAction", "MotionStrength", + "MotionScale", "TemporalRhythm", "CameraMotion", + ], + "MT": [ + "MotionProfile", "MotionStrength", "MotionScale", + "TemporalRhythm", "CameraMotion", "EventStructure", + ], + "FFA": [ + "PrimaryAction", "Interaction", "SocialContingency", + "MotionStrength", "TemporalRhythm", + ], +} + + +# ── Combined T2V space (content + motion) ──────────────────────────────────── +# Only used for unconstrained T2V search (no anchor image). +# art_data and ROI_ACTIVE_CATEGORIES are unchanged for image search. + +def make_t2v_art_data() -> Dict: + combined = dict(art_data) + combined.update(VIDEO_MOTION_ART_DATA) + return combined + + +T2V_ROI_ACTIVE_CATEGORIES: Dict[str, List[str]] = { + roi: list(dict.fromkeys( + ROI_ACTIVE_CATEGORIES.get(roi, []) + VIDEO_MOTION_ROI_ACTIVE.get(roi, []) + )) + for roi in set(list(ROI_ACTIVE_CATEGORIES.keys()) + list(VIDEO_MOTION_ROI_ACTIVE.keys())) +} + + +def flatten_art_data(ad: Dict) -> tuple[list[str], list[list[str]]]: + categories = list(ad.keys()) + options: list[list[str]] = [] + for k in categories: + v = ad[k] + if isinstance(v, dict): + merged = [] + for s in v.values(): + merged.extend(s) + options.append(list(merged)) + else: + options.append(list(v)) + return categories, options + + +class StructuredArtPromptSpace(SearchSpace): + """ + ROI-aware prompt search space. + Only categories relevant to the target ROI are searched; + all others are locked to their empty-string slot. + """ + + def __init__( + self, + art_data: Dict, + roi: Optional[str] = None, + option_embeddings: Optional[list[torch.Tensor]] = None, + semantic_temperature: float = 0.2, + ): + self.categories, self.options = flatten_art_data(art_data) + self.option_embeddings = option_embeddings + self.semantic_temperature = semantic_temperature + self.sizes = np.array([len(o) for o in self.options], dtype=np.int32) + self.offsets = np.concatenate([[0], np.cumsum(self.sizes)[:-1]]).astype(np.int32) + self.feat_dim = int(np.sum(self.sizes)) + + active_cats = ROI_ACTIVE_CATEGORIES.get(roi, _ALL_CATEGORIES) if roi else _ALL_CATEGORIES + self.active_mask = [cat in active_cats for cat in self.categories] + self.empty_indices = [len(opts) - 1 for opts in self.options] + + def _empty_gene(self, cat_idx: int) -> int: + return self.empty_indices[cat_idx] + + def random_candidate(self) -> Candidate: + genes = [ + random.randrange(len(self.options[i])) if self.active_mask[i] + else self._empty_gene(i) + for i in range(len(self.options)) + ] + return Candidate(genes=tuple(genes)) + + def decode(self, cand: Candidate) -> str: + parts = [self.options[i][cand.genes[i]] for i in range(len(self.options))] + parts = [p for p in parts if p and p.strip()] + return ", ".join(parts) + + def encode(self, prompt: str) -> Optional[List[int]]: + tokens = [t.strip() for t in prompt.split(", ") if t.strip()] + genes = list(self.empty_indices) + token_idx = 0 + for cat_idx, opts in enumerate(self.options): + if token_idx >= len(tokens): + break + tok = tokens[token_idx] + if tok in opts: + genes[cat_idx] = opts.index(tok) + token_idx += 1 + if token_idx < len(tokens): + return None + return genes + + def mutate(self, cand: Candidate, mutation_rate: float) -> Candidate: + genes = list(cand.genes) + for i in range(len(genes)): + if self.active_mask[i] and random.random() < mutation_rate: + genes[i] = random.randrange(len(self.options[i])) + return Candidate(genes=tuple(genes)) + + def semantic_mutate(self, cand: Candidate, mutation_rate: float) -> Candidate: + if self.option_embeddings is None: + return self.mutate(cand, mutation_rate) + genes = list(cand.genes) + for i in range(len(genes)): + if self.active_mask[i] and random.random() < mutation_rate: + emb = self.option_embeddings[i] + if emb is None: + genes[i] = random.randrange(len(self.options[i])) + continue + sims = emb @ emb[genes[i]] + probs = torch.softmax(sims / self.semantic_temperature, dim=0).cpu().numpy() + genes[i] = int(np.random.choice(len(probs), p=probs)) + return Candidate(genes=tuple(genes)) + + def crossover(self, a: Candidate, b: Candidate, crossover_rate: float) -> tuple[Candidate, Candidate]: + if random.random() > crossover_rate: + return a, b + point = random.randrange(1, len(a.genes)) + c1 = a.genes[:point] + b.genes[point:] + c2 = b.genes[:point] + a.genes[point:] + return Candidate(tuple(c1)), Candidate(tuple(c2)) + + def hamming(self, a, b) -> int: + return int(sum(int(x != y) for x, y in zip(a, b))) + + def featurize(self, ind) -> np.ndarray: + x = np.zeros((self.feat_dim,), dtype=np.float32) + for ci, opt_idx in enumerate(ind): + x[int(self.offsets[ci] + opt_idx)] = 1.0 + return x + + +# ── VideoMotionPromptSpace ──────────────────────────────────────────────────── +# Motion-only search space for I2V mode. +# The anchor image handles content; this space searches over motion/interaction. + +class VideoMotionPromptSpace(SearchSpace): + def __init__(self, roi: Optional[str] = None, + option_embeddings: Optional[list] = None, + semantic_temperature: float = 0.2): + self.categories = list(VIDEO_MOTION_ART_DATA.keys()) + self.options = [VIDEO_MOTION_ART_DATA[c] for c in self.categories] + self.option_embeddings = option_embeddings + self.semantic_temperature = semantic_temperature + self.sizes = np.array([len(o) for o in self.options], dtype=np.int32) + self.offsets = np.concatenate([[0], np.cumsum(self.sizes)[:-1]]).astype(np.int32) + self.feat_dim = int(np.sum(self.sizes)) + + active_cats = VIDEO_MOTION_ROI_ACTIVE.get(roi, self.categories) if roi else self.categories + self.active_mask = [cat in active_cats for cat in self.categories] + self.empty_indices = [len(opts) - 1 for opts in self.options] + + def random_candidate(self) -> Candidate: + genes = [ + random.randrange(len(self.options[i])) if self.active_mask[i] + else self.empty_indices[i] + for i in range(len(self.options)) + ] + return Candidate(genes=tuple(genes)) + + def decode(self, cand: Candidate) -> str: + parts = [self.options[i][cand.genes[i]] for i in range(len(self.options))] + return ", ".join(p for p in parts if p and p.strip()) + + def encode(self, prompt: str) -> Optional[List[int]]: + tokens = [t.strip() for t in prompt.split(", ") if t.strip()] + genes = list(self.empty_indices) + token_idx = 0 + for cat_idx, opts in enumerate(self.options): + if token_idx >= len(tokens): + break + if tokens[token_idx] in opts: + genes[cat_idx] = opts.index(tokens[token_idx]) + token_idx += 1 + return genes if token_idx == len(tokens) else None + + def mutate(self, cand: Candidate, mutation_rate: float) -> Candidate: + genes = list(cand.genes) + for i in range(len(genes)): + if self.active_mask[i] and random.random() < mutation_rate: + genes[i] = random.randrange(len(self.options[i])) + return Candidate(genes=tuple(genes)) + + def semantic_mutate(self, cand: Candidate, mutation_rate: float) -> Candidate: + if self.option_embeddings is None: + return self.mutate(cand, mutation_rate) + genes = list(cand.genes) + for i in range(len(genes)): + if self.active_mask[i] and random.random() < mutation_rate: + emb = self.option_embeddings[i] + if emb is None: + genes[i] = random.randrange(len(self.options[i])) + continue + sims = emb @ emb[genes[i]] + probs = torch.softmax(sims / self.semantic_temperature, dim=0).cpu().numpy() + genes[i] = int(np.random.choice(len(probs), p=probs)) + return Candidate(genes=tuple(genes)) + + def crossover(self, a: Candidate, b: Candidate, crossover_rate: float): + if random.random() > crossover_rate: + return a, b + point = random.randrange(1, len(a.genes)) + return (Candidate(tuple(a.genes[:point] + b.genes[point:])), + Candidate(tuple(b.genes[:point] + a.genes[point:]))) + + def hamming(self, a, b) -> int: + return sum(x != y for x, y in zip(a, b)) + + def featurize(self, ind) -> np.ndarray: + x = np.zeros((self.feat_dim,), dtype=np.float32) + for ci, opt_idx in enumerate(ind): + x[int(self.offsets[ci] + opt_idx)] = 1.0 + return x \ No newline at end of file diff --git a/stimulus_synthesis/utils/__init__.py b/stimulus_synthesis/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/stimulus_synthesis/utils/__init__.py @@ -0,0 +1 @@ + diff --git a/stimulus_synthesis/utils/seed.py b/stimulus_synthesis/utils/seed.py new file mode 100644 index 0000000000000000000000000000000000000000..e8cb118b4af8f6544b6234d61d796e6cfb6f69f8 --- /dev/null +++ b/stimulus_synthesis/utils/seed.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +import random + +import numpy as np +import torch + + +def set_all_seeds(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) diff --git a/stimulus_synthesis_config.json b/stimulus_synthesis_config.json new file mode 100644 index 0000000000000000000000000000000000000000..a3e37cfb8ceb379042b5376b409a169b2f8f99d7 --- /dev/null +++ b/stimulus_synthesis_config.json @@ -0,0 +1,37 @@ +{ + "encoder_model_id": "epfl-neuroai/vjepa2-encoder-basic", + "encoder_call": "predict_fmri", + "text_to_image_model_id": "stabilityai/sdxl-turbo", + "image_to_video_model_id": "Lightricks/LTX-Video-0.9.8-13B-distilled", + "default_objective": "indices_mean", + "default_device": "cuda", + "default_mutation_rate": 0.25, + "default_elite_frac": 0.35, + "default_score_transform": { + "enabled": false, + "crop_scale": 0.8, + "gaussian_sigma": 0.1, + "num_draws": 4, + "aggregate": "mean", + "seed": 0 + }, + "default_score_frames": 24, + "default_image_batch_size": 16, + "default_video_batch_size": 8, + "default_population_size": 20, + "default_image_kwargs": { + "num_inference_steps": 1, + "guidance_scale": 0.0, + "height": 512, + "width": 512 + }, + "default_video_kwargs": { + "num_inference_steps": 8, + "num_frames": 25, + "height": 512, + "width": 512 + }, + "default_image_max_evals": 400, + "default_video_max_evals": 200, + "default_score_size": 224 +} diff --git a/tests/test_asset_scoring.py b/tests/test_asset_scoring.py new file mode 100644 index 0000000000000000000000000000000000000000..45ac2b282c50a152a87aed87deec5b8a7299e5b4 --- /dev/null +++ b/tests/test_asset_scoring.py @@ -0,0 +1,58 @@ +import json + +import torch +from PIL import Image + +from stimulus_synthesis.asset_manifest import load_asset_manifest, write_asset_manifest +from stimulus_synthesis.media import ImageAssetSpec, VideoAssetSpec, decode_image, decode_video, export_image, export_video +from stimulus_synthesis.scoring import AssetScorer, EncoderPreprocessSpec, prepare_video_for_encoder + + +class MeanScorer: + def score(self, videos, target, **kwargs): + return videos.mean(dim=(1, 2, 3, 4)).tolist() + + +def test_image_export_decode_and_asset_score(tmp_path): + image = Image.new("RGB", (6, 6), color=(128, 64, 32)) + spec = ImageAssetSpec(width=12, height=10, format="png") + path = tmp_path / "stimulus.png" + + export_record = export_image(image, path, spec) + decoded = decode_image(path) + scorer = AssetScorer(MeanScorer(), target=None, preprocess_spec=EncoderPreprocessSpec(size=8, num_frames=3)) + score_record = scorer.score_image(path, asset_spec=spec, metadata={"prompt": "test prompt"}) + + assert export_record.sha256 == decoded.sha256 == score_record.sha256 + assert decoded.image.shape == (3, 10, 12) + assert score_record.sampled_frame_indices == [0, 0, 0] + assert score_record.preprocess["size"] == 8 + assert score_record.asset_spec["width"] == 12 + assert isinstance(score_record.score, float) + + +def test_video_export_decode_preprocess_and_manifest(tmp_path): + frames = torch.zeros(3, 3, 8, 8) + frames[1] = 0.5 + frames[2] = 1.0 + spec = VideoAssetSpec(width=16, height=16, fps=24, num_frames=5, crf=18) + path = tmp_path / "stimulus.mp4" + + export_record = export_video(frames, path, spec) + decoded = decode_video(path) + prepared = prepare_video_for_encoder(decoded.frames, EncoderPreprocessSpec(size=(8, 8), num_frames=4)) + scorer = AssetScorer(MeanScorer(), target=None, preprocess_spec=EncoderPreprocessSpec(size=(8, 8), num_frames=4)) + score_record = scorer.score_video(path, asset_spec=spec) + + assert export_record.sha256 == decoded.sha256 == score_record.sha256 + assert decoded.frames.shape[1:] == (3, 16, 16) + assert decoded.num_frames == 5 + assert prepared.videos.shape == (1, 4, 3, 8, 8) + assert score_record.sampled_frame_indices == [0, 1, 3, 4] + + manifest_path = tmp_path / "manifest.json" + manifest = write_asset_manifest([export_record, score_record], manifest_path, metadata={"model": "mock"}) + loaded = load_asset_manifest(manifest_path) + assert loaded == manifest + assert loaded["metadata"]["model"] == "mock" + assert len(loaded["records"]) == 2 diff --git a/tests/test_output_schema.py b/tests/test_output_schema.py new file mode 100644 index 0000000000000000000000000000000000000000..c7785e45f3cb801e0a02ac608b546fd5765e9e45 --- /dev/null +++ b/tests/test_output_schema.py @@ -0,0 +1,25 @@ +import torch + +from stimulus_synthesis.outputs import StimulusCandidate, StimulusSynthesisOutput +from stimulus_synthesis.scoring.objectives import indices_mean, vector_cosine, vector_dot + + +def test_objectives(): + predictions = torch.tensor([[1.0, 2.0, 3.0], [3.0, 2.0, 1.0]]) + + assert torch.allclose(indices_mean(predictions, {"type": "indices", "indices": [0, 2]}), torch.tensor([2.0, 2.0])) + assert torch.allclose(vector_dot(predictions, {"type": "vector", "vector": [1.0, 0.0, 0.0]}), torch.tensor([1.0, 3.0])) + assert vector_cosine(predictions, {"type": "vector", "vector": [1.0, 0.0, 0.0]}).shape == (2,) + + +def test_output_schema(): + candidate = StimulusCandidate(prompt="person running", score=1.5, image="image", video="video") + output = StimulusSynthesisOutput( + candidates=[candidate], + best_prompt=candidate.prompt, + best_score=candidate.score, + history_best=[1.5], + ) + + assert output.best.prompt == "person running" + assert output.best_score == 1.5 diff --git a/tests/test_prompt_space.py b/tests/test_prompt_space.py new file mode 100644 index 0000000000000000000000000000000000000000..b1e4cb325b1d67f8056c0884df12f57d48dcdf5a --- /dev/null +++ b/tests/test_prompt_space.py @@ -0,0 +1,18 @@ +from stimulus_synthesis.spaces import Candidate, PromptSearchSpace + + +def test_prompt_space_decode_mutate_and_crossover(): + space = PromptSearchSpace( + prompt_banks={ + "subject": ["person", "face"], + "action": ["running", "walking"], + } + ) + cand = Candidate((0, 1)) + + assert space.decode(cand) == "person walking" + assert len(space.mutate(cand, mutation_rate=1.0).genes) == 2 + + child_a, child_b = space.crossover(Candidate((0, 0)), Candidate((1, 1)), crossover_rate=1.0) + assert len(child_a.genes) == 2 + assert len(child_b.genes) == 2 diff --git a/tests/test_robust_transform.py b/tests/test_robust_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..e54ba4ad3898b035211cf2c28296ba7f06a0228b --- /dev/null +++ b/tests/test_robust_transform.py @@ -0,0 +1,53 @@ +import torch + +from stimulus_synthesis.scoring.robust_transform import RobustTransformScorer, RobustTransformSpec, apply_robust_transform + + +class MeanScorer: + def score(self, videos, target, **kwargs): + return videos.mean(dim=(1, 2, 3, 4)).tolist() + + +def test_robust_transform_expands_draws_and_is_reproducible(): + videos = torch.linspace(0, 1, steps=2 * 3 * 8 * 8).reshape(1, 2, 3, 8, 8) + spec = RobustTransformSpec(num_draws=4, crop_scale=0.8, gaussian_sigma=0.1) + + first = apply_robust_transform(videos, spec) + second = apply_robust_transform(videos, spec) + + assert first.shape == (4, 2, 3, 8, 8) + assert torch.equal(first, second) + + +def test_robust_transform_scorer_is_independent_of_batch_order(): + a = torch.zeros(2, 3, 8, 8) + b = torch.ones(2, 3, 8, 8) * 0.5 + videos = torch.stack([a, b], dim=0) + reversed_videos = torch.stack([b, a], dim=0) + scorer = RobustTransformScorer(MeanScorer(), RobustTransformSpec(num_draws=4, crop_scale=0.8, gaussian_sigma=0.1)) + + scores = scorer.score(videos, None) + reversed_scores = scorer.score(reversed_videos, None) + + assert torch.allclose(torch.tensor(scores), torch.tensor(list(reversed(reversed_scores)))) + + +def test_pipeline_default_score_transform_is_clean(monkeypatch): + """Default scoring matches the canonical clean single pass (no robust augmentation).""" + from stimulus_synthesis.pipeline import NevoPipeline + from stimulus_synthesis.scoring.robust_transform import RobustTransformScorer + + class DummyEncoderScorer: + def __init__(self, *args, **kwargs): + pass + + def score(self, videos, target, **kwargs): + return [0.0] * videos.shape[0] + + monkeypatch.setattr("stimulus_synthesis.pipeline.EncoderScorer", DummyEncoderScorer) + pipe = NevoPipeline(text_to_image=object(), image_to_video=object()) + pipe._ensure_components(device="cpu") + + # Robust augmentation is disabled by default -> the encoder scorer is used directly. + assert not isinstance(pipe.scorer, RobustTransformScorer) + assert isinstance(pipe.scorer, DummyEncoderScorer) diff --git a/tests/test_roi_search_space.py b/tests/test_roi_search_space.py new file mode 100644 index 0000000000000000000000000000000000000000..e0e649316f6c9e4acffe63f55f6db438967edbc8 --- /dev/null +++ b/tests/test_roi_search_space.py @@ -0,0 +1,49 @@ +import torch +from PIL import Image + +from stimulus_synthesis import NevoPipeline +from stimulus_synthesis.spaces import StructuredArtPromptSpace, PromptSearchSpace + + +class MockTextToImage: + def generate(self, prompts, *, generator=None, **kwargs): + return [Image.new("RGB", (8, 8)) for _ in prompts] + + +class MockImageToVideo: + def generate(self, image, prompt, *, generator=None, **kwargs): + return torch.full((2, 3, 8, 8), len(prompt) / 100.0) + + def generate_batch(self, images, prompts, *, generators=None, **kwargs): + return [self.generate(i, p) for i, p in zip(images, prompts)] + + +class MockScorer: + def score(self, videos, target, **kwargs): + return videos.mean(dim=(1, 2, 3, 4)).tolist() + + +def _pipe(): + return NevoPipeline(text_to_image=MockTextToImage(), image_to_video=MockImageToVideo(), scorer=MockScorer()) + + +def test_roi_selects_enhanced_structured_space(): + pipe = _pipe() + enhanced = pipe.make_search_space(roi="FFA") + general = pipe.make_search_space(roi="FFA", enforce_general_search_space=True) + assert isinstance(enhanced, StructuredArtPromptSpace) + assert isinstance(general, StructuredArtPromptSpace) + # ROI-enhanced space searches strictly fewer categories than the general one + assert sum(enhanced.active_mask) < sum(general.active_mask) + + +def test_no_roi_uses_general_prompt_space(): + assert isinstance(_pipe().make_search_space(seed_prompts=["a cat"]), PromptSearchSpace) + + +def test_pipeline_runs_with_roi_and_general_override(): + pipe = _pipe() + out = pipe(roi="FFA", image_max_evals=2, video_max_evals=2, population_size=2, seed=0, score_size=None) + assert out.best_prompt and isinstance(out.best_score, float) + out2 = pipe(roi="FFA", enforce_general_search_space=True, image_max_evals=2, video_max_evals=2, population_size=2, seed=0, score_size=None) + assert out2.best_prompt is not None diff --git a/tests/test_seeded_asset_search.py b/tests/test_seeded_asset_search.py new file mode 100644 index 0000000000000000000000000000000000000000..12fdf862afe4f54e9238756bc5f7aa396ac8ab07 --- /dev/null +++ b/tests/test_seeded_asset_search.py @@ -0,0 +1,178 @@ +from pathlib import Path + +import torch +from PIL import Image + +from stimulus_synthesis.media import ImageAssetSpec, VideoAssetSpec +from stimulus_synthesis.scoring import AssetScorer, EncoderPreprocessSpec +from stimulus_synthesis.search.genetic import GeneticSearch +from stimulus_synthesis.spaces import PromptSearchSpace, SeededSearchSpace + + +class MeanScorer: + def score(self, videos, target, **kwargs): + return videos.mean(dim=(1, 2, 3, 4)).tolist() + + +class SeedColorTextToImage: + def __init__(self): + self.calls = [] + + def generate(self, prompts, *, generator=None, **kwargs): + seed = int(generator.initial_seed()) if generator is not None else -1 + self.calls.append((prompts[0], seed)) + value = seed % 256 + return [Image.new("RGB", (8, 8), color=(value, value, value)) for _ in prompts] + + +class StaticImageToVideo: + def generate(self, image, prompt, *, generator=None, **kwargs): + return image + + +class SeedValueVideo: + def generate(self, image, prompt, *, generator=None, **kwargs): + seed = int(generator.initial_seed()) if generator is not None else 0 + value = (seed % 256) / 255.0 + return torch.full((2, 3, 8, 8), value) + + +def test_seeded_search_space_appends_seed_gene(): + base = PromptSearchSpace(prompt_banks={"subject": ["face"], "style": ["photo"]}) + space = SeededSearchSpace(base, [11, 22]) + cand = space.random_candidate() + + assert len(cand.genes) == 3 + assert space.decode(cand) == "face photo" + assert space.decode_seed(cand) in {11, 22} + assert len(space.options) == 3 + + +def test_asset_scored_search_keys_by_prompt_and_seed(tmp_path): + base = PromptSearchSpace(prompt_banks={"subject": ["same prompt"]}) + space = SeededSearchSpace(base, [0, 255]) + scorer = AssetScorer(MeanScorer(), target=None, preprocess_spec=EncoderPreprocessSpec(size=8, num_frames=1)) + t2i = SeedColorTextToImage() + search = GeneticSearch( + max_evals=2, + population_size=2, + n_init=2, + mutation_rate=1.0, + image_kwargs={}, + video_kwargs={}, + score_size=8, + num_frames=1, + asset_scorer=scorer, + asset_dir=tmp_path / "candidates", + asset_type="image", + image_asset_spec=ImageAssetSpec(width=8, height=8), + ) + + result = search.run(space, t2i, StaticImageToVideo(), MeanScorer(), target=None, seed=123) + + assert result.best_seed == 255 + assert result.best_score == result.best_score_record.score + assert result.best_export_record.sha256 == result.best_score_record.sha256 + assert result.best_asset_path and Path(result.best_asset_path).exists() + assert result.best_metadata["score_source"] == "saved_asset" + assert len({seed for _prompt, seed in t2i.calls}) == 2 + + +def test_video_search_uses_candidate_seed_for_video_generator(tmp_path): + base = PromptSearchSpace(prompt_banks={"motion": ["move"]}) + space = SeededSearchSpace(base, [0, 127]) + scorer = AssetScorer(MeanScorer(), target=None, preprocess_spec=EncoderPreprocessSpec(size=8, num_frames=2)) + fixed_image = Image.new("RGB", (8, 8), color=(0, 0, 0)) + + class FixedImageT2I: + def generate(self, prompts, *, generator=None, **kwargs): + return [fixed_image for _ in prompts] + + search = GeneticSearch( + max_evals=2, + population_size=2, + n_init=2, + mutation_rate=1.0, + image_kwargs={}, + video_kwargs={}, + score_size=8, + num_frames=2, + asset_scorer=scorer, + asset_dir=tmp_path / "videos", + asset_type="video", + video_asset_spec=VideoAssetSpec(width=8, height=8, fps=24, num_frames=2, crf=0), + ) + + result = search.run(space, FixedImageT2I(), SeedValueVideo(), MeanScorer(), target=None, seed=123) + + assert result.best_seed == 127 + assert result.best_asset_path and Path(result.best_asset_path).exists() + assert result.best_export_record.sha256 == result.best_score_record.sha256 + + +def test_same_prompt_seed_and_params_give_same_saved_image_hash(tmp_path): + base = PromptSearchSpace(prompt_banks={"subject": ["same prompt"]}) + space = SeededSearchSpace(base, [12345]) + scorer = AssetScorer(MeanScorer(), target=None, preprocess_spec=EncoderPreprocessSpec(size=8, num_frames=1)) + spec = ImageAssetSpec(width=8, height=8) + + def run_once(out_dir): + search = GeneticSearch( + max_evals=1, + population_size=1, + n_init=1, + mutation_rate=0.0, + image_kwargs={}, + video_kwargs={}, + score_size=8, + num_frames=1, + asset_scorer=scorer, + asset_dir=out_dir, + asset_type="image", + image_asset_spec=spec, + ) + return search.run(space, SeedColorTextToImage(), StaticImageToVideo(), MeanScorer(), target=None, seed=999) + + first = run_once(tmp_path / "first") + second = run_once(tmp_path / "second") + + assert first.best_prompt == second.best_prompt == "same prompt" + assert first.best_seed == second.best_seed == 12345 + assert first.best_export_record.sha256 == second.best_export_record.sha256 + assert first.best_score_record.sha256 == second.best_score_record.sha256 + assert first.best_score == second.best_score + + +def test_same_candidate_hash_is_independent_of_evaluation_order(tmp_path): + base = PromptSearchSpace(prompt_banks={"subject": ["same prompt", "filler prompt"]}) + space = SeededSearchSpace(base, [12345, 54321]) + scorer = AssetScorer(MeanScorer(), target=None, preprocess_spec=EncoderPreprocessSpec(size=8, num_frames=1)) + spec = ImageAssetSpec(width=8, height=8) + + def run_order(inds, out_dir): + search = GeneticSearch( + max_evals=len(inds), + population_size=len(inds), + n_init=len(inds), + mutation_rate=0.0, + image_kwargs={}, + video_kwargs={}, + score_size=8, + num_frames=1, + asset_scorer=scorer, + asset_dir=out_dir, + asset_type="image", + image_asset_spec=spec, + ) + # Exercise the same evaluation path GeneticSearch uses, but with fixed order. + from stimulus_synthesis.search.evaluation import EvaluationManager, candidate_key + evaluator = EvaluationManager() + search._evaluate(evaluator, space, SeedColorTextToImage(), StaticImageToVideo(), MeanScorer(), None, inds, seed=999) + key = candidate_key("same prompt", 12345) + return evaluator.asset_export_cache[key].sha256, evaluator.asset_score_cache[key].score + + first_hash, first_score = run_order([[0, 0], [1, 1]], tmp_path / "first_order") + second_hash, second_score = run_order([[1, 1], [0, 0]], tmp_path / "second_order") + + assert first_hash == second_hash + assert first_score == second_score diff --git a/tests/test_smoke_pipeline.py b/tests/test_smoke_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..f9cd38c8589adc5ab4ade0e100ffcb8bff00194b --- /dev/null +++ b/tests/test_smoke_pipeline.py @@ -0,0 +1,61 @@ +from pathlib import Path + +import torch +from PIL import Image +from diffusers import DiffusionPipeline + +from stimulus_synthesis.pipeline import NevoPipeline +from stimulus_synthesis.spaces import PromptSearchSpace + + +class MockTextToImage: + def generate(self, prompts, *, generator=None, **kwargs): + return [Image.new("RGB", (8, 8), color=(len(prompt) % 255, 0, 0)) for prompt in prompts] + + +class MockImageToVideo: + def generate(self, image, prompt, *, generator=None, **kwargs): + value = len(prompt) / 100.0 + return torch.full((2, 3, 8, 8), value) + + def generate_batch(self, images, prompts, *, generators=None, **kwargs): + return [self.generate(image, prompt, **kwargs) for image, prompt in zip(images, prompts)] + + +class MockScorer: + def score(self, videos, target, **kwargs): + return videos.mean(dim=(1, 2, 3, 4)).tolist() + + +def test_pipeline_smoke_with_mock_components(): + space = PromptSearchSpace( + prompt_banks={ + "subject": ["a", "long subject"], + "action": ["runs", "walks"], + } + ) + pipe = NevoPipeline( + text_to_image=MockTextToImage(), + image_to_video=MockImageToVideo(), + scorer=MockScorer(), + search_space=space, + ) + + out = pipe(target=[0], image_max_evals=2, population_size=2, seed=123, score_size=None) + + assert out.best_prompt + assert isinstance(out.best_score, float) + assert out.best.image is not None + assert out.best.video is not None + assert out.metadata["max_evals"] == 2 + + +def test_diffusers_custom_pipeline_loads_local_repo(): + repo = str(Path(__file__).resolve().parents[1]) + pipe = DiffusionPipeline.from_pretrained( + repo, + custom_pipeline=repo, + ) + + assert isinstance(pipe, NevoPipeline) + assert pipe.synthesis_config.encoder_call == "predict_fmri" diff --git a/tests/test_video_normalize.py b/tests/test_video_normalize.py new file mode 100644 index 0000000000000000000000000000000000000000..0225dfa855566af205d727b33fedee12a2bf2500 --- /dev/null +++ b/tests/test_video_normalize.py @@ -0,0 +1,34 @@ +import numpy as np +import torch +from PIL import Image + +from stimulus_synthesis.media.normalize import videos_to_b_t_c_h_w + + +def test_videos_to_b_t_c_h_w_from_pil_frames(): + frames = [ + Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)), + Image.fromarray(np.full((8, 8, 3), 255, dtype=np.uint8)), + ] + + video = videos_to_b_t_c_h_w([frames], size=4, num_frames=3) + + assert video.shape == (1, 3, 3, 4, 4) + assert torch.all(video >= 0) + assert torch.all(video <= 1) + + +def test_videos_to_b_t_c_h_w_from_tensor_thwc(): + tensor = torch.zeros(2, 8, 8, 3) + + video = videos_to_b_t_c_h_w([tensor], size=4) + + assert video.shape == (1, 2, 3, 4, 4) + + +def test_videos_to_b_t_c_h_w_from_numpy_bthwc(): + array = np.zeros((1, 2, 8, 8, 3), dtype=np.float32) + + video = videos_to_b_t_c_h_w([array], size=4) + + assert video.shape == (1, 2, 3, 4, 4)