Fix unseeded generation when seed_base is unset; make the critic's request params configurable

agentic_upsampling/clients.py

@@ -477,8 +477,9 @@ class VLMQualityJudge:
         api_token: str,
         endpoint_url: str = DEFAULT_CRITIC_ENDPOINT_URL,
         model: str = DEFAULT_CRITIC_MODEL,
-        max_tokens: int = 8192,
+        max_tokens: int = 16384,
         image_jpeg_quality: int | None = DEFAULT_JPEG_QUALITY,
+        extra_body: dict[str, Any] | None = None,
     ) -> None:
         self.chat_client = OpenAIChatClient(
             ChatClientConfig(
@@ -487,6 +488,7 @@ class VLMQualityJudge:
                 api_token=api_token,
                 max_tokens=max_tokens,
                 max_retries=3,
+                extra_body=extra_body,
             )
         )
         self.image_jpeg_quality = image_jpeg_quality

agentic_upsampling/runner.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import json
+import random
 import traceback
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass
@@ -112,6 +113,12 @@ class AgenticUpsamplerRunner:
         (item_dir / "failure.json").unlink(missing_ok=True)
         (item_dir / "incomplete.json").unlink(missing_ok=True)
         self._log(f"[prompt {item.prompt_id}] start")
+        # One base per item, reused by every iteration (sample slot j gets base + j), so a
+        # rewrite is compared against the same noise. A resumed run draws a new base.
+        seed_base = self.config.seed_base
+        if seed_base is None:
+            seed_base = random.randint(0, 2**31 - 1 - self.config.samples_per_iteration)
+            self._log(f"[prompt {item.prompt_id}] seed_base drawn: {seed_base}")
         candidates: list[dict[str, Any]] = []
         previous_prompt: dict[str, Any] | None = None
         previous_analysis: dict[str, Any] | None = None
@@ -131,6 +138,7 @@ class AgenticUpsamplerRunner:
                         previous_analysis,
                         previous_negative_prompt,
                         candidates,
+                        seed_base,
                     )
                 except Exception as exc:
                     if not candidates:
@@ -152,7 +160,9 @@ class AgenticUpsamplerRunner:
                 self._log(f"[prompt {item.prompt_id}] early stop at iter={iteration}")
                 break
 
-        return self.finalize_item(item, candidates, incomplete_error=incomplete_error)
+        result = self.finalize_item(item, candidates, incomplete_error=incomplete_error)
+        result["seed_base"] = seed_base
+        return result
 
     def run_item_safely(self, item: PromptItem) -> dict[str, Any]:
         """Run one item and convert failures into structured records."""
@@ -179,6 +189,7 @@ class AgenticUpsamplerRunner:
         previous_analysis: dict[str, Any] | None,
         previous_negative_prompt: str,
         candidates: list[dict[str, Any]],
+        seed_base: int,
     ) -> dict[str, Any]:
         prepared = self.prepare_iteration_prompt(
             item,
@@ -195,6 +206,7 @@ class AgenticUpsamplerRunner:
             iteration,
             prepared.prompt_json,
             prepared.negative_prompt,
+            seed_base,
         )
         return self.finalize_iteration(item, iteration_dir, iteration, sample_candidates, sample_errors)
 
@@ -205,6 +217,7 @@ class AgenticUpsamplerRunner:
         iteration: int,
         prompt_json: dict[str, Any],
         negative_prompt: str,
+        seed_base: int,
     ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
         """Generate seed samples concurrently, then judge successful images in sample order."""
         generation_outputs: dict[int, GenerationOutput] = {}
@@ -218,6 +231,7 @@ class AgenticUpsamplerRunner:
                     sample_index,
                     prompt_json,
                     negative_prompt,
+                    seed_base,
                 ): sample_index
                 for sample_index in range(self.config.samples_per_iteration)
             }
@@ -300,8 +314,11 @@ class AgenticUpsamplerRunner:
         sample_index: int,
         prompt_json: dict[str, Any],
         negative_prompt: str,
+        seed_base: int,
     ) -> dict[str, Any]:
-        generation = self.run_generation_sample(item, iteration_dir, sample_index, prompt_json, negative_prompt)
+        generation = self.run_generation_sample(
+            item, iteration_dir, sample_index, prompt_json, negative_prompt, seed_base
+        )
         return self.judge_iteration_sample(
             item,
             iteration_dir,
@@ -319,6 +336,7 @@ class AgenticUpsamplerRunner:
         sample_index: int,
         prompt_json: dict[str, Any],
         negative_prompt: str,
+        seed_base: int,
     ) -> GenerationOutput:
         """Generate one sample image for an iteration."""
         sample_dir = self._sample_dir(iteration_dir, sample_index)
@@ -328,7 +346,7 @@ class AgenticUpsamplerRunner:
             prompt_json=prompt_json,
             prompt_id=item.prompt_id,
             output_dir=sample_dir,
-            seed=self._sample_seed(sample_index),
+            seed=seed_base + sample_index,
             negative_prompt=negative_prompt,
             jpeg_quality=self.config.jpeg_quality,
         )
@@ -422,11 +440,6 @@ class AgenticUpsamplerRunner:
         if self.config.verbose:
             print(message, flush=True)
 
-    def _sample_seed(self, sample_index: int) -> int | None:
-        if self.config.seed_base is None:
-            return None
-        return self.config.seed_base + sample_index
-
     def _sample_dir(self, iteration_dir: Path, sample_index: int) -> Path:
         if self.config.samples_per_iteration == 1:
             return iteration_dir

tests/test_agentic_upsampling.py

@@ -10,7 +10,7 @@ from typing import Any
 
 from PIL import Image
 
-from agentic_upsampling.clients import ImageGenerationClient, PromptRewriterClient
+from agentic_upsampling.clients import ImageGenerationClient, PromptRewriterClient, VLMQualityJudge
 from agentic_upsampling.constants import (
     DEFAULT_CRITIC_ENDPOINT_URL,
     DEFAULT_CRITIC_MODEL,
@@ -413,7 +413,32 @@ def test_runner_early_stops_by_default(tmp_path: Path) -> None:
     assert result["best_iteration"] == 0
     assert rewriter.initial_calls == 1
     assert rewriter.joint_rewrite_calls == 0
-    assert generator.seeds == [None]
+    assert generator.seeds == [result["seed_base"]]
+
+
+def test_runner_draws_one_seed_base_when_unset(tmp_path: Path) -> None:
+    """An unset seed_base is drawn once and reused, so slot j gets base + j every iteration."""
+    rewriter = FakeRewriter()
+    generator = FakeGenerator()
+    runner = AgenticUpsamplerRunner(
+        rewriter=rewriter,
+        generator=generator,  # type: ignore[arg-type]
+        judge=FakeJudge([5.0, 6.0, 5.5, 6.5]),
+        config=RunnerConfig(
+            output_dir=tmp_path,
+            max_iterations=2,
+            samples_per_iteration=2,
+            early_stop=False,
+        ),
+    )
+
+    result = runner.run_item(_item())
+
+    seed_base = result["seed_base"]
+    assert seed_base is not None
+    # Samples are generated concurrently, so compare each iteration's seeds as a set.
+    assert sorted(generator.seeds[:2]) == [seed_base, seed_base + 1]
+    assert sorted(generator.seeds[2:]) == [seed_base, seed_base + 1]
 
 
 def test_runner_can_disable_early_stop_and_select_best_sample(tmp_path: Path) -> None:
@@ -496,3 +521,14 @@ def test_extract_best_images_copies_images_and_writes_manifests(tmp_path: Path)
     assert (tmp_path / "export" / "best_generations.jsonl").exists()
     assert (tmp_path / "export" / "best_generations.csv").exists()
 
+
+def test_vlm_quality_judge_sends_no_extra_body_by_default() -> None:
+    judge = VLMQualityJudge(api_token="token")
+
+    assert judge.chat_client.config.extra_body is None
+
+
+def test_vlm_quality_judge_forwards_extra_body() -> None:
+    judge = VLMQualityJudge(api_token="token", extra_body={"reasoning_effort": "high"})
+
+    assert judge.chat_client.config.extra_body == {"reasoning_effort": "high"}