Duplicate from epfl-neuroai/NEvo

Co-authored-by: david <dawidtang@users.noreply.huggingface.co>

.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. <repo>/cache/  (only if no default is resolvable)
+# NEVO_CACHE_DIR=/absolute/path/to/nevo_cache

.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

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.DS_Store
+cache/
+.env

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=<int> 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-<n>` (both hemispheres), `SL-<n>_lh`, `SL-<n>_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).

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)

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)

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)

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)

model_index.json

@@ -0,0 +1,5 @@
+{
+  "_class_name": "NevoPipeline",
+  "_diffusers_version": "0.30.0",
+  "_module": "pipeline"
+}

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

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

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

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

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

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

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

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,
+        }

stimulus_synthesis/data/roi_masks.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3be584b579472777c7f41e3589e47b6cee5b608636c633bf605e1dd266f76914
+size 5758

stimulus_synthesis/data/searchlight_both.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2c3e3fec1a88d56e182ccb1362f303f50fa24bae1d8796fee99f80b30e2bea76
+size 2585

stimulus_synthesis/data/searchlight_lh.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a9ba6bc47b7d7075c0faf0cf6edd8acebaf9f2fe1f588899d6e7323b00f89979
+size 1379

stimulus_synthesis/data/searchlight_rh.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:838c8ef226320a484c1776092d653d73d77721881b2c64113a88d13b83190b3a
+size 1478

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

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)
+        ]

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

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

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",
+]

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)

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

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

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)

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

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

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-<n>``         -> both-hemisphere searchlight region (1-indexed)
+- ``SL-<n>_lh`` / ``SL-<n>_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<name>.+?)_(?P<hemi>lh|rh)$", flags=re.IGNORECASE)
+_SL_PATTERN = re.compile(r"^SL-(?P<idx>\d+)(?:_(?P<hemi>lh|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

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]

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 ``<repo>/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
+       ``<repo>/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")

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}

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",
+]

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)

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]:
+        ...

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

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

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

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)