run_regional_asset_pilot.py

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