Split into image_num_generations (50) / video_num_generations (25)

README.md

@@ -82,7 +82,7 @@ from stimulus_synthesis.media import save_video, video_to_t_c_h_w   # importable
 save_video(video_to_t_c_h_w(out.best.video), "best_stimulus.mp4")   # save the synthesized video
 ```
 
-This runs the two-stage search with the defaults — a population of 20 evolved for 10 generations per stage, at full-resolution LTX video — which is thorough but heavy (a single run can take a while and a lot of GPU memory).
+This runs the two-stage search with the defaults — a population of 20 evolved for 50 generations (image) then 25 (video), at full-resolution LTX video — which is thorough but heavy (a single run can take a while and a lot of GPU memory).
 
 ### Faster run
 
@@ -92,7 +92,8 @@ For a quicker first result, shrink the search and the video:
 out = pipe(
     roi="FFA",
     progress=True,
-    num_generations=4,        # GA generations per stage (default: 10)
+    image_num_generations=4,  # stage-1 (image) generations (default: 50)
+    video_num_generations=4,  # stage-2 (video) generations (default: 25)
     population_size=8,        # GA population per generation (default: 20)
     seed=0,                   # RNG seed, for reproducibility
     video_kwargs={            # passed straight to LTX-Video; omit to use the model's own (heavier) defaults
@@ -144,14 +145,15 @@ Set in `stimulus_synthesis_config.json`:
 
 | Param | Default | Notes |
 |-------|---------|-------|
-| `default_num_generations` | 10 | number of GA generations per stage (generation 1 = the random init) |
+| `default_image_generations` | 50 | stage-1 (image) GA generations (generation 1 = the random init) |
+| `default_video_generations` | 25 | stage-2 (video) GA generations |
 | `default_population_size` | 20 | GA population per generation (= `n_init`) |
 | `default_mutation_rate` | 0.25 | |
 | `default_elite_frac` | 0.35 | |
 | `default_objective` | `indices_mean` | |
 | `default_score_transform` | disabled | robust augmentation off by default (clean single pass) |
 
-Each stage evaluates a population of `population_size` (default 20) for `num_generations` (default 10) generations — so total generate→score work per stage is about `population_size × num_generations`. Video generation dominates runtime — pass `video_kwargs={"num_inference_steps": 8, "num_frames": 25, ...}` to keep it fast (the distilled LTX model runs in few steps); otherwise LTX's own defaults apply.
+Each stage evaluates a population of `population_size` (default 20) for its own number of generations — `image_num_generations` (default 50) and `video_num_generations` (default 25) — so total generate→score work per stage is about `population_size × generations`. Video generation dominates runtime — pass `video_kwargs={"num_inference_steps": 8, "num_frames": 25, ...}` to keep it fast (the distilled LTX model runs in few steps); otherwise LTX's own defaults apply.
 
 ### Robust scoring
 

stimulus_synthesis/config.py

@@ -15,7 +15,8 @@ class StimulusSynthesisConfig:
     default_objective: str = "indices_mean"
     default_device: str = "cuda"
     default_population_size: int = 20
-    default_num_generations: int = 10
+    default_image_generations: int = 50
+    default_video_generations: int = 25
     default_score_frames: int = 24
     default_image_batch_size: int = 16
     default_video_batch_size: int = 8
@@ -52,7 +53,8 @@ class StimulusSynthesisConfig:
             "default_objective": self.default_objective,
             "default_device": self.default_device,
             "default_population_size": self.default_population_size,
-            "default_num_generations": self.default_num_generations,
+            "default_image_generations": self.default_image_generations,
+            "default_video_generations": self.default_video_generations,
             "default_score_frames": self.default_score_frames,
             "default_image_batch_size": self.default_image_batch_size,
             "default_video_batch_size": self.default_video_batch_size,

stimulus_synthesis/pipeline.py

@@ -133,7 +133,8 @@ class NevoPipeline(DiffusionPipeline):
         roi,
         enforce_general_search_space,
         progress,
-        num_generations,
+        image_num_generations,
+        video_num_generations,
         image_batch_size,
         video_batch_size,
         population_size,
@@ -155,12 +156,13 @@ class NevoPipeline(DiffusionPipeline):
         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))
-        gens = max(1, int(num_generations or cfg.default_num_generations))
+        image_gens = max(1, int(image_num_generations or cfg.default_image_generations))
+        video_gens = max(1, int(video_num_generations or cfg.default_video_generations))
         score_frames = num_frames if num_frames is not None else cfg.default_score_frames
 
         # ---- Stage 1: image prompt search ----
         image_search = GeneticSearch(
-            num_generations=gens,
+            num_generations=image_gens,
             population_size=pop,
             n_init=pop,
             mutation_rate=cfg.default_mutation_rate,
@@ -189,7 +191,7 @@ class NevoPipeline(DiffusionPipeline):
 
         # ---- Stage 2: motion prompt search on the fixed best image ----
         video_search = GeneticSearch(
-            num_generations=gens,
+            num_generations=video_gens,
             population_size=pop,
             n_init=pop,
             mutation_rate=cfg.default_mutation_rate,
@@ -240,7 +242,8 @@ class NevoPipeline(DiffusionPipeline):
                 "image_to_video_model_id": cfg.image_to_video_model_id,
                 "objective": cfg.default_objective,
                 "two_stage": True,
-                "num_generations": gens,
+                "image_generations": image_gens,
+                "video_generations": video_gens,
                 "population_size": pop,
             },
         )
@@ -255,7 +258,8 @@ class NevoPipeline(DiffusionPipeline):
         progress: bool = False,
         image_batch_size: int | None = None,
         video_batch_size: int | None = None,
-        num_generations: int | None = None,
+        image_num_generations: int | None = None,
+        video_num_generations: int | None = None,
         population_size: int | None = None,
         seed: int = 0,
         prompt_banks: dict[str, list[str]] | None = None,
@@ -278,7 +282,8 @@ class NevoPipeline(DiffusionPipeline):
         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)
-        gens = max(1, num_generations if num_generations is not None else cfg.default_num_generations)
+        image_gens = max(1, image_num_generations if image_num_generations is not None else cfg.default_image_generations)
+        video_gens = max(1, video_num_generations if video_num_generations is not None else cfg.default_video_generations)
 
         if roi is not None and target is None:
             from .neuro import resolve_driving_voxels
@@ -296,7 +301,8 @@ class NevoPipeline(DiffusionPipeline):
                 roi=roi,
                 enforce_general_search_space=enforce_general_search_space,
                 progress=progress,
-                num_generations=gens,
+                image_num_generations=image_gens,
+                video_num_generations=video_gens,
                 image_batch_size=img_bs,
                 video_batch_size=vid_bs,
                 population_size=pop,
@@ -320,7 +326,7 @@ class NevoPipeline(DiffusionPipeline):
             seed_prompts=seed_prompts,
         )
         search = GeneticSearch(
-            num_generations=gens,
+            num_generations=image_gens,
             population_size=pop,
             n_init=pop,
             mutation_rate=cfg.default_mutation_rate,
@@ -358,7 +364,7 @@ class NevoPipeline(DiffusionPipeline):
                 "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,
-                "num_generations": gens,
+                "num_generations": image_gens,
             },
         )
 

stimulus_synthesis_config.json

@@ -19,5 +19,6 @@
   "default_image_batch_size": 16,
   "default_video_batch_size": 8,
   "default_population_size": 20,
-  "default_num_generations": 10
+  "default_image_generations": 50,
+  "default_video_generations": 25
 }

tests/test_roi_search_space.py

@@ -43,7 +43,7 @@ def test_no_roi_uses_general_prompt_space():
 
 def test_pipeline_runs_with_roi_and_general_override():
     pipe = _pipe()
-    out = pipe(roi="FFA", num_generations=2, population_size=2, seed=0, video_size=None)
+    out = pipe(roi="FFA", image_num_generations=2, video_num_generations=2, population_size=2, seed=0, video_size=None)
     assert out.best_prompt and isinstance(out.best_score, float)
-    out2 = pipe(roi="FFA", enforce_general_search_space=True, num_generations=2, population_size=2, seed=0, video_size=None)
+    out2 = pipe(roi="FFA", enforce_general_search_space=True, image_num_generations=2, video_num_generations=2, population_size=2, seed=0, video_size=None)
     assert out2.best_prompt is not None

tests/test_smoke_pipeline.py

@@ -41,7 +41,7 @@ def test_pipeline_smoke_with_mock_components():
         search_space=space,
     )
 
-    out = pipe(target=[0], num_generations=2, population_size=2, seed=123, video_size=None)
+    out = pipe(target=[0], image_num_generations=2, population_size=2, seed=123, video_size=None)
 
     assert out.best_prompt
     assert isinstance(out.best_score, float)