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