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}