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