run_roi_samples.py

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