[feat] integrate InterleaveThinker model backends

.agents/exploration/interleavethinker-fastvideo-integration.md

@@ -295,6 +295,146 @@ skill for adding agentic app layers over FastVideo generators:
 - fake-provider unit tests,
 - Modal smoke-test pattern.
 
+## Model / API Backend Integration Extension
+
+Status: in progress as of the user request to integrate the remaining
+InterleaveThinker models and wrap closed-source Nano Banana / Gemini APIs.
+
+User constraints still active:
+
+- Work only in `/tmp/fastvideo-interleavethinker`.
+- Do not edit or overwrite `/home/toolbox/FastVideo`; that checkout belongs to
+  another agent.
+- Commit useful checkpoints and push each commit immediately.
+- Run validation on Modal through `fastvideo/tests/modal/launch_l40s_job.py`.
+- Keep this handoff file current before interruptions or context compaction.
+
+Scope for this slice:
+
+1. Replace the placeholder InterleaveThinker critic actor shell with a concrete
+   Hugging Face Qwen3-VL-compatible `ModelBase` wrapper.
+   - Load `InterleaveThinker/Critic-SFT-8B`,
+     `InterleaveThinker/InterleaveThinker-Critic-8B`, or a local
+     `ckpt/critic_sft` / `ckpt/critic_rl` checkpoint via Transformers.
+   - Preserve FastVideo's `ModelBase` role-module/checkpoint visibility through
+     `self.transformer`.
+   - Implement `generate_interleave_responses(...)` for grouped RL rollouts.
+   - Implement `train_interleave_rollouts(...)` as a lightweight
+     advantage-weighted policy-gradient update over generated critic responses.
+     This is not a full EasyR1/vLLM/FSDP parity port, but it makes the FastVideo
+     RL loop executable with a real Qwen3-VL actor backend.
+   - Add a small JSON/JSONL dataloader for InterleaveThinker critic RL records
+     so the standard FastVideo trainer can receive batches.
+2. Add closed-source API wrappers:
+   - Nano Banana image generation/editing via the official `google-genai`
+     `models.generate_content` API, with model aliases for Nano Banana,
+     Nano Banana Pro, and Nano Banana 2.
+   - Gemini VLM scoring via the same SDK, returning InterleaveThinker
+     semantic/quality scores.
+   - A composite `GeminiNanoBananaEditScorer` callable that the RL reward
+     scorer can instantiate from YAML.
+3. Wire the API scorer into `InterleaveThinkerRLMethod` through an optional
+   `_target_` config block while keeping existing offline/fake-score tests
+   credential-free.
+4. Update examples and focused tests.
+
+Official API facts verified against Google AI docs on 2026-06-19:
+
+- Nano Banana image generation/editing is exposed through Gemini API native
+  image models. Current documented model IDs include:
+  `gemini-3.1-flash-image` (Nano Banana 2), `gemini-3-pro-image`
+  (Nano Banana Pro), and `gemini-2.5-flash-image` (Nano Banana).
+- Python SDK pattern:
+  `from google import genai`; `client = genai.Client(...)`;
+  `client.models.generate_content(model=..., contents=[prompt, image])`.
+- Image responses are returned as response parts with inline data / `as_image`.
+- Structured Gemini outputs use `response_mime_type="application/json"` and a
+  response schema in `GenerateContentConfig`.
+
+Implementation boundary:
+
+- The Qwen3-VL critic is a local/open model wrapper because it must train in the
+  RL loop.
+- Nano Banana and Gemini remain API models; FastVideo should not pretend they
+  are local trainable modules.
+- The API wrappers must import `google-genai` lazily and raise clear dependency
+  / API-key errors only when used.
+- Tests must stub the SDK; no real API calls in CI/Modal validation.
+
+Implementation completed for this slice:
+
+- Replaced the placeholder
+  `fastvideo/train/models/interleave_thinker/critic.py` adapter with a
+  Qwen3-VL-compatible `InterleaveThinkerCriticModel`.
+  - Loads a processor and model through Transformers when `load_backend=true`.
+  - Keeps `load_backend=false` for import/config/unit-test dry runs.
+  - Exposes `self.transformer` for FastVideo role-module/checkpoint visibility.
+  - Freezes visual tower and multimodal projector by default, matching upstream
+    critic SFT.
+  - Implements `build_messages(...)` with the InterleaveThinker two-image
+    critic prompt.
+  - Implements `generate_interleave_responses(...)` for grouped critic
+    rollouts.
+  - Implements `train_interleave_rollouts(...)` as an advantage-weighted
+    policy-gradient update over response token NLL.
+  - Adds a small JSON/JSONL dataloader for critic RL records.
+- Added `NanoBananaImageGeneratorBackend` to
+  `fastvideo/entrypoints/interleave/generator.py`.
+  - Supports model aliases: `nano-banana`, `nano-banana-pro`,
+    `nano-banana-2`.
+  - Uses the official `google-genai` `models.generate_content` API lazily.
+  - Implements the existing `ImageGeneratorBackend` protocol.
+- Added `fastvideo/train/methods/rl/rewards/interleave_api.py`.
+  - `GeminiInterleaveImageScorer` wraps Gemini VLM JSON scoring.
+  - `GeminiNanoBananaEditScorer` generates edits with Nano Banana and scores
+    them with Gemini, matching `EditScoreProvider`.
+  - `ConstantInterleaveEditScorer` gives tests/offline configs a deterministic
+    scorer.
+- Wired `InterleaveThinkerRLMethod` to accept an optional `method.edit_scorer`
+  `_target_` block.
+- Updated `examples/train/configs/rl/interleave_thinker/critic_grpo.yaml` to
+  point at the Qwen3-VL critic wrapper and the Gemini/Nano Banana scorer.
+- Added optional dependency extra: `fastvideo[interleave-api] = google-genai`.
+- Documented the API-backed path in `examples/interleave/README.md`.
+- Added tests:
+  - `tests/local_tests/test_interleave_thinker_api_models.py`
+  - `tests/local_tests/test_interleave_thinker_critic_model.py`
+  - extended `tests/local_tests/test_interleave_thinker_method.py`
+
+Validation for this slice:
+
+- Initial Modal pytest run:
+  - App: `https://modal.com/apps/hao-ai-lab/main/ap-rGgMjcQcI2bNAl3Viy0eES`
+  - Result: one test assertion failure in
+    `test_interleave_thinker_critic_builds_qwen_vl_messages_without_loading_backend`.
+  - Cause: test assumed fixed content indices, but adjacent `<image><image>`
+    placeholders produce adjacent image parts.
+  - Fix: assert by part type instead of index.
+- Second Modal pytest run:
+  - App: `https://modal.com/apps/hao-ai-lab/main/ap-SIztyW6PC3yRAKdybFxEEE`
+  - Result: `27 passed, 14 warnings in 19.71s`.
+- First Modal pre-commit run:
+  - App: `https://modal.com/apps/hao-ai-lab/main/ap-E0mnD27dHXDDnOquRHU5QD`
+  - Result: yapf and ruff modified files; mypy passed on Modal.
+- Local formatter convergence:
+  - Installed `pre-commit==4.0.1` locally with hook cache under `/tmp`.
+  - Local hook run converged for yapf/ruff/codespell.
+  - Local mypy cannot run from `/tmp/fastvideo-interleavethinker` because the
+    hyphenated directory name is interpreted as an invalid package name; Modal
+    mypy remains the authoritative result.
+- Post-format Modal pytest:
+  - App: `https://modal.com/apps/hao-ai-lab/main/ap-ykgfYTZyJgO4EgRlHv6l81`
+  - Result: `27 passed, 14 warnings in 17.04s`.
+- Final Modal pre-commit:
+  - App: `https://modal.com/apps/hao-ai-lab/main/ap-QOKlzapm5bSAo3c21lprwv`
+  - Result: yapf, ruff, codespell, mypy, filename check, and suggestion hooks
+    passed; PyMarkdown/actionlint skipped with no files to check.
+
+Pending:
+
+- Commit this model/API integration slice.
+- Push the commit immediately after committing, per user instruction.
+
 ## Training Loop Extension
 
 Status: in progress as of the user request to "Add a training loop" for

examples/interleave/README.md

@@ -51,6 +51,29 @@ export EDIT_MODEL_NAME=klein
 
 The same app also exposes `/generate` and `/v1/interleave/edit` aliases.
 
+## Nano Banana / Gemini API Backends
+
+The InterleaveThinker RL config can use closed-source Google models through
+lazy wrappers:
+
+- `fastvideo.train.methods.rl.rewards.GeminiNanoBananaEditScorer`
+  generates edits with Nano Banana and scores them with Gemini.
+- `fastvideo.entrypoints.interleave.generator.NanoBananaImageGeneratorBackend`
+  implements the same image backend protocol as the local FastVideo generator.
+
+Install the optional SDK and provide a key only when using these API backends:
+
+```bash
+uv pip install -e ".[interleave-api]"
+export GEMINI_API_KEY=...
+```
+
+Supported Nano Banana aliases are:
+
+- `nano-banana` -> `gemini-2.5-flash-image`
+- `nano-banana-pro` -> `gemini-3-pro-image`
+- `nano-banana-2` -> `gemini-3.1-flash-image`
+
 ## Native Single-Prompt Trace
 
 Run:

examples/train/configs/rl/interleave_thinker/critic_grpo.yaml

@@ -2,16 +2,22 @@
 #
 # Upstream InterleaveThinker trains a Qwen3-VL critic checkpoint with EasyR1
 # GRPO. This config expresses the same outer loop through FastVideo's modular
-# trainer. The default model target is an adapter shell; replace or extend it
-# with a concrete Qwen/VLM backend that implements:
-#
-#   generate_interleave_responses(...)
-#   train_interleave_rollouts(...)
+# trainer with a local Qwen3-VL actor wrapper plus API-backed Gemini/Nano
+# Banana rewards.
 
 models:
   student:
     _target_: fastvideo.train.models.interleave_thinker.InterleaveThinkerCriticModel
-    init_from: ckpt/critic_sft
+    init_from: InterleaveThinker/Critic-SFT-8B
+    processor_from: Qwen/Qwen3-VL-8B-Instruct
+    load_backend: true
+    torch_dtype: auto
+    attn_implementation: flash_attention_2
+    freeze_vision_tower: true
+    freeze_multi_modal_projector: true
+    enable_gradient_checkpointing: true
+    max_prompt_length: 16384
+    max_response_length: 4096
     trainable: true
 
 method:
@@ -26,6 +32,15 @@ method:
   semantic_weight: 0.6
   quality_weight: 0.2
   fallback_edit_reward: 0.5
+  edit_scorer:
+    _target_: fastvideo.train.methods.rl.rewards.GeminiNanoBananaEditScorer
+    image_model: gemini-3.1-flash-image
+    judge_model: gemini-2.5-pro
+    output_dir: outputs/interleave_thinker_critic_grpo/reward_api
+    width: 1024
+    height: 1024
+    num_inference_steps: 4
+    guidance_scale: 1.0
   advantage_eps: 0.0001
   advantage_clip: 5.0
   max_grad_norm: 1.0

fastvideo/entrypoints/interleave/generator.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import base64
+import io
 import os
 import time
 import uuid
@@ -23,6 +24,12 @@ from fastvideo.entrypoints.interleave.schema import (
     InterleaveEditRequest,
 )
 
+_NANO_BANANA_MODEL_ALIASES = {
+    "nano-banana": "gemini-2.5-flash-image",
+    "nano-banana-pro": "gemini-3-pro-image",
+    "nano-banana-2": "gemini-3.1-flash-image",
+}
+
 
 class ImageGeneratorBackend(Protocol):
     """Minimal image-generation backend used by the Interleave app layer."""
@@ -138,6 +145,107 @@ class FastVideoImageGeneratorBackend:
         )
 
 
+class NanoBananaImageGeneratorBackend:
+    """Google Gemini API image backend for Nano Banana models.
+
+    This wraps the closed-source Gemini native-image API behind the same
+    ``ImageGeneratorBackend`` protocol used by the FastVideo compatibility
+    service. The SDK import and API-key validation are intentionally lazy so
+    installing FastVideo does not require ``google-genai`` unless this backend is
+    configured.
+    """
+
+    def __init__(
+        self,
+        *,
+        model: str = "gemini-3.1-flash-image",
+        api_key: str | None = None,
+        base_url: str | None = None,
+        output_dir: str = "outputs/nano_banana",
+        aspect_ratio: str | None = None,
+        image_size: str | None = None,
+        max_attempts: int = 3,
+        retry_delay_s: float = 2.0,
+    ) -> None:
+        self.model = _NANO_BANANA_MODEL_ALIASES.get(model, model)
+        self.api_key = api_key
+        self.base_url = base_url
+        self.output_dir = output_dir
+        self.aspect_ratio = aspect_ratio
+        self.image_size = image_size
+        self.max_attempts = max(1, int(max_attempts))
+        self.retry_delay_s = float(retry_delay_s)
+        self._client: Any | None = None
+
+    def generate(
+        self,
+        request: InterleaveEditRequest,
+        *,
+        request_id: str | None = None,
+    ) -> GeneratedImage:
+        request_id = request_id or uuid.uuid4().hex
+        output_format = (request.output_format or "png").lower()
+        if output_format == "jpg":
+            output_format = "jpeg"
+        output_path = Path(self.output_dir) / "interleave" / f"{request_id}.{output_format}"
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        contents: list[Any] = [request.prompt]
+        if request.image:
+            contents.append(_decode_base64_to_pil(request.image))
+
+        last_exc: Exception | None = None
+        start = time.perf_counter()
+        for attempt in range(self.max_attempts):
+            try:
+                response = self._client_instance().models.generate_content(
+                    model=self.model,
+                    contents=contents,
+                    config=self._make_generate_config(),
+                )
+                image = _extract_first_response_image(response)
+                image.save(output_path)
+                return GeneratedImage(
+                    prompt=request.prompt,
+                    image_base64=encode_file_to_base64(output_path),
+                    file_path=str(output_path.resolve()),
+                    inference_time_s=time.perf_counter() - start,
+                    metadata={
+                        "request_id": request_id,
+                        "model": self.model,
+                        "attempt": attempt + 1,
+                    },
+                )
+            except Exception as exc:  # noqa: BLE001 - remote API errors vary by SDK version
+                last_exc = exc
+                if attempt + 1 < self.max_attempts:
+                    time.sleep(self.retry_delay_s)
+        raise RuntimeError(
+            f"Nano Banana generation failed after {self.max_attempts} attempts: {last_exc}") from last_exc
+
+    def _client_instance(self) -> Any:
+        if self._client is not None:
+            return self._client
+        genai, _ = _import_google_genai()
+        kwargs: dict[str, Any] = {"api_key": _resolve_google_api_key(self.api_key)}
+        if self.base_url:
+            kwargs["http_options"] = {"base_url": self.base_url}
+        self._client = genai.Client(**kwargs)
+        return self._client
+
+    def _make_generate_config(self) -> Any:
+        _, types = _import_google_genai()
+        kwargs: dict[str, Any] = {"response_modalities": ["TEXT", "IMAGE"]}
+        if self.aspect_ratio or self.image_size:
+            image_kwargs: dict[str, Any] = {}
+            if self.aspect_ratio:
+                image_kwargs["aspect_ratio"] = self.aspect_ratio
+            if self.image_size:
+                image_kwargs["image_size"] = self.image_size
+            kwargs["image_config"] = types.ImageConfig(**image_kwargs)
+        return types.GenerateContentConfig(**kwargs)
+
+
 def _safe_explicit_request_updates(request: GenerationRequest) -> dict[str, Any]:
     try:
         return explicit_request_updates(request)
@@ -170,3 +278,59 @@ def decode_base64_image_to_path(
     path.parent.mkdir(parents=True, exist_ok=True)
     path.write_bytes(data)
     return str(path)
+
+
+def _resolve_google_api_key(explicit: str | None = None) -> str:
+    if explicit:
+        return explicit.strip()
+    for env_name in ("GEMINI_API_KEY", "GOOGLE_API_KEY"):
+        value = os.environ.get(env_name)
+        if value:
+            return value.strip()
+    token_path = Path("~/.gemini_token").expanduser()
+    if token_path.is_file():
+        return token_path.read_text().strip()
+    raise ValueError("Google Gemini API access requires GEMINI_API_KEY, GOOGLE_API_KEY, "
+                     "an explicit api_key, or ~/.gemini_token.")
+
+
+def _import_google_genai() -> tuple[Any, Any]:
+    try:
+        from google import genai
+        from google.genai import types
+    except ImportError as exc:
+        raise RuntimeError("Nano Banana API backend requires google-genai. "
+                           "Install with `uv pip install -e '.[interleave-api]'`.") from exc
+    return genai, types
+
+
+def _decode_base64_to_pil(image_base64: str) -> Any:
+    from PIL import Image
+
+    payload = image_base64.strip()
+    if payload.startswith("data:") and "," in payload:
+        payload = payload.split(",", 1)[1]
+    return Image.open(io.BytesIO(base64.b64decode(payload))).convert("RGB")
+
+
+def _extract_first_response_image(response: Any) -> Any:
+    from PIL import Image
+
+    parts = getattr(response, "parts", None)
+    if parts is None:
+        candidates = getattr(response, "candidates", None) or []
+        if candidates:
+            parts = getattr(getattr(candidates[0], "content", None), "parts", None)
+    for part in parts or []:
+        as_image = getattr(part, "as_image", None)
+        if callable(as_image):
+            image = as_image()
+            if isinstance(image, Image.Image):
+                return image
+        inline_data = getattr(part, "inline_data", None) or getattr(part, "inlineData", None)
+        data = getattr(inline_data, "data", None)
+        if data:
+            if isinstance(data, str):
+                data = base64.b64decode(data)
+            return Image.open(io.BytesIO(data)).convert("RGB")
+    raise RuntimeError("Gemini image response did not include an image part")

fastvideo/train/methods/rl/interleave_thinker.py

@@ -13,17 +13,17 @@ import torch.distributed as dist
 from fastvideo.logger import init_logger
 from fastvideo.train.methods.base import LogScalar, TrainingMethod
 from fastvideo.train.methods.rl.rewards import (
-    InterleaveThinkerRewardScorer,
-)
+    InterleaveThinkerRewardScorer, )
 from fastvideo.train.models.base import ModelBase
 from fastvideo.train.utils.config import (
     get_optional_float,
     get_optional_int,
     parse_betas,
 )
+from fastvideo.train.utils.instantiate import (
+    instantiate, )
 from fastvideo.train.utils.optimizer import (
-    build_optimizer_and_scheduler,
-)
+    build_optimizer_and_scheduler, )
 
 logger = init_logger(__name__)
 
@@ -76,6 +76,7 @@ class InterleaveThinkerRLMethod(TrainingMethod):
             semantic_weight=self._read_float("semantic_weight", 0.6),
             quality_weight=self._read_float("quality_weight", 0.2),
             fallback_edit_reward=self._read_float("fallback_edit_reward", 0.5),
+            edit_scorer=self._build_edit_scorer(self.method_config.get("edit_scorer")),
         )
         self._student_optimizer: torch.optim.Optimizer | None = None
         self._student_lr_scheduler: Any | None = None
@@ -186,6 +187,18 @@ class InterleaveThinkerRLMethod(TrainingMethod):
             scheduler_name=str(self.training_config.optimizer.lr_scheduler),
         )
 
+    def _build_edit_scorer(
+        self,
+        raw: Any,
+    ) -> Any:
+        if raw is None:
+            return None
+        if callable(raw):
+            return raw
+        if isinstance(raw, Mapping):
+            return instantiate(dict(raw))
+        raise TypeError("method.edit_scorer must be a callable or a mapping with _target_")
+
     def _generate_rollouts(
         self,
         batch: dict[str, Any],
@@ -353,9 +366,7 @@ class InterleaveThinkerRLMethod(TrainingMethod):
         raise TypeError(f"Expected scalar metric, got {type(value).__name__}")
 
     @staticmethod
-    def _batch_to_items(
-        batch: Mapping[str, Any],
-    ) -> list[dict[str, Any]]:
+    def _batch_to_items(batch: Mapping[str, Any], ) -> list[dict[str, Any]]:
         if "items" in batch and isinstance(batch["items"], Sequence):
             return [dict(item) for item in batch["items"]]
 
@@ -372,9 +383,8 @@ class InterleaveThinkerRLMethod(TrainingMethod):
         for idx in range(batch_size):
             item: dict[str, Any] = {}
             for key, value in batch.items():
-                if isinstance(value, list | tuple) and len(value) == batch_size:
-                    item[key] = value[idx]
-                elif torch.is_tensor(value) and value.ndim > 0 and int(value.shape[0]) == batch_size:
+                if isinstance(value, list | tuple) and len(value) == batch_size or torch.is_tensor(
+                        value) and value.ndim > 0 and int(value.shape[0]) == batch_size:
                     item[key] = value[idx]
                 else:
                     item[key] = value
@@ -435,9 +445,7 @@ class InterleaveThinkerRLMethod(TrainingMethod):
         group_sizes: dict[str, int] = defaultdict(int)
         for key in group_keys:
             group_sizes[key] += 1
-        group_size_values = torch.tensor(list(group_sizes.values()),
-                                         device=advantages.device,
-                                         dtype=torch.float32)
+        group_size_values = torch.tensor(list(group_sizes.values()), device=advantages.device, dtype=torch.float32)
         return {
             "interleave/advantage_mean": advantages.detach().float().mean(),
             "interleave/advantage_std": advantages.detach().float().std(unbiased=False),

fastvideo/train/methods/rl/rewards/__init__.py

@@ -5,6 +5,11 @@ from fastvideo.train.methods.rl.rewards.frame_rewards import (
     ClipScoreScorer,
     PickScoreScorer,
 )
+from fastvideo.train.methods.rl.rewards.interleave_api import (
+    ConstantInterleaveEditScorer,
+    GeminiInterleaveImageScorer,
+    GeminiNanoBananaEditScorer,
+)
 from fastvideo.train.methods.rl.rewards.interleave_thinker import (
     EditScoreProvider,
     InterleaveThinkerAnswer,
@@ -42,7 +47,10 @@ def build_multi_reward_scorer(
 
 __all__ = [
     "ClipScoreScorer",
+    "ConstantInterleaveEditScorer",
     "EditScoreProvider",
+    "GeminiInterleaveImageScorer",
+    "GeminiNanoBananaEditScorer",
     "InterleaveThinkerAnswer",
     "InterleaveThinkerEditRequest",
     "InterleaveThinkerEditScore",

fastvideo/train/methods/rl/rewards/interleave_api.py

@@ -0,0 +1,324 @@
+# SPDX-License-Identifier: Apache-2.0
+"""API-backed InterleaveThinker reward services.
+
+These classes wrap the closed-source models used by InterleaveThinker without
+making them mandatory FastVideo dependencies:
+
+- Nano Banana / Gemini native-image models generate or edit images.
+- Gemini VLM models score semantic alignment and perceptual quality.
+
+The reusable reward parser in ``interleave_thinker.py`` remains pure; this file
+contains the optional networked scorer that can be configured from YAML.
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+import json
+import os
+from pathlib import Path
+import time
+import uuid
+from typing import Any
+
+from fastvideo.entrypoints.interleave.generator import (
+    NanoBananaImageGeneratorBackend,
+    encode_file_to_base64,
+)
+from fastvideo.entrypoints.interleave.schema import (
+    InterleaveEditRequest as InterleaveAPIEditRequest, )
+from fastvideo.train.methods.rl.rewards.interleave_thinker import (
+    InterleaveThinkerEditRequest,
+    InterleaveThinkerEditScore,
+)
+
+_SCORE_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "semantic_score": {
+            "type": "number"
+        },
+        "quality_score": {
+            "type": "number"
+        },
+        "semantic_analysis": {
+            "type": "string"
+        },
+        "quality_analysis": {
+            "type": "string"
+        },
+    },
+    "required": ["semantic_score", "quality_score"],
+}
+
+_SCORE_SYSTEM_PROMPT = ("You are a strict image editing evaluator. Return only JSON. Scores must be "
+                        "numbers from 0 to 10, where 10 is best.")
+
+_SCORE_USER_PROMPT = """\
+Evaluate the edited image for an iterative image generation/editing step.
+
+Instruction:
+{instruction}
+
+Return:
+- semantic_score: how well the edited image satisfies the instruction and preserves required content.
+- quality_score: visual quality, realism/coherence, artifact absence, and logical consistency.
+
+Use the full 0-10 range. Do not reward unrelated changes.
+"""
+
+
+class GeminiInterleaveImageScorer:
+    """Gemini VLM scorer returning InterleaveThinker semantic/quality scores."""
+
+    def __init__(
+        self,
+        *,
+        model: str = "gemini-2.5-pro",
+        api_key: str | None = None,
+        base_url: str | None = None,
+        max_attempts: int = 6,
+        retry_delay_s: float = 2.0,
+        temperature: float = 0.2,
+    ) -> None:
+        self.model = model
+        self.api_key = api_key
+        self.base_url = base_url
+        self.max_attempts = max(1, int(max_attempts))
+        self.retry_delay_s = float(retry_delay_s)
+        self.temperature = float(temperature)
+        self._client: Any | None = None
+
+    def score_images(
+        self,
+        *,
+        edited_image_path: str,
+        instruction: str,
+        reference_image_path: str | None = None,
+    ) -> InterleaveThinkerEditScore:
+        prompt = _SCORE_USER_PROMPT.format(instruction=(instruction or "").strip())
+        contents: list[Any] = [prompt]
+        if reference_image_path:
+            contents.extend(["Reference image before this step:", self._image_part(reference_image_path)])
+        contents.extend(["Edited/generated image to score:", self._image_part(edited_image_path)])
+
+        last_exc: Exception | None = None
+        for attempt in range(self.max_attempts):
+            try:
+                response = self._client_instance().models.generate_content(
+                    model=self.model,
+                    contents=contents,
+                    config=self._make_score_config(),
+                )
+                payload = _load_jsonish(getattr(response, "text", "") or "")
+                return InterleaveThinkerEditScore(
+                    semantic_score=float(payload["semantic_score"]),
+                    quality_score=float(payload["quality_score"]),
+                )
+            except Exception as exc:  # noqa: BLE001 - remote API errors vary by SDK version
+                last_exc = exc
+                if attempt + 1 < self.max_attempts:
+                    time.sleep(self.retry_delay_s)
+        raise RuntimeError(f"Gemini image scoring failed after {self.max_attempts} attempts: {last_exc}") from last_exc
+
+    def _client_instance(self) -> Any:
+        if self._client is not None:
+            return self._client
+        genai, _ = _import_google_genai()
+        kwargs: dict[str, Any] = {"api_key": _resolve_google_api_key(self.api_key)}
+        if self.base_url:
+            kwargs["http_options"] = {"base_url": self.base_url}
+        self._client = genai.Client(**kwargs)
+        return self._client
+
+    def _make_score_config(self) -> Any:
+        _, types = _import_google_genai()
+        return types.GenerateContentConfig(
+            system_instruction=_SCORE_SYSTEM_PROMPT,
+            response_mime_type="application/json",
+            response_schema=_SCORE_SCHEMA,
+            temperature=self.temperature,
+        )
+
+    def _image_part(self, path: str) -> Any:
+        _, types = _import_google_genai()
+        image_path = Path(path)
+        mime_type = "image/png" if image_path.suffix.lower() == ".png" else "image/jpeg"
+        return types.Part.from_bytes(
+            data=image_path.read_bytes(),
+            mime_type=mime_type,
+        )
+
+
+class GeminiNanoBananaEditScorer:
+    """Generate an edit with Nano Banana and score it with Gemini.
+
+    This callable matches ``EditScoreProvider`` from
+    ``interleave_thinker.py`` and can be passed to
+    ``InterleaveThinkerRewardScorer``.
+    """
+
+    def __init__(
+        self,
+        *,
+        image_model: str = "gemini-3.1-flash-image",
+        judge_model: str = "gemini-2.5-pro",
+        api_key: str | None = None,
+        base_url: str | None = None,
+        output_dir: str = "outputs/interleave_thinker_reward_api",
+        width: int = 1024,
+        height: int = 1024,
+        num_inference_steps: int = 4,
+        guidance_scale: float = 1.0,
+        aspect_ratio: str | None = "1:1",
+        image_size: str | None = None,
+        max_attempts: int = 3,
+        retry_delay_s: float = 2.0,
+        treat_white_canvas_as_text_to_image: bool = True,
+    ) -> None:
+        self.output_dir = output_dir
+        self.width = int(width)
+        self.height = int(height)
+        self.num_inference_steps = int(num_inference_steps)
+        self.guidance_scale = float(guidance_scale)
+        self.treat_white_canvas_as_text_to_image = bool(treat_white_canvas_as_text_to_image)
+        self.editor = NanoBananaImageGeneratorBackend(
+            model=image_model,
+            api_key=api_key,
+            base_url=base_url,
+            output_dir=output_dir,
+            aspect_ratio=aspect_ratio,
+            image_size=image_size,
+            max_attempts=max_attempts,
+            retry_delay_s=retry_delay_s,
+        )
+        self.scorer = GeminiInterleaveImageScorer(
+            model=judge_model,
+            api_key=api_key,
+            base_url=base_url,
+            max_attempts=max_attempts,
+            retry_delay_s=retry_delay_s,
+        )
+
+    def __call__(
+        self,
+        request: InterleaveThinkerEditRequest,
+    ) -> InterleaveThinkerEditScore | None:
+        input_image_path = request.previous_image_path or request.origin_image_path
+        image_base64 = None
+        if input_image_path and not self._is_text_to_image_canvas(input_image_path):
+            image_base64 = encode_file_to_base64(input_image_path)
+
+        edit_request = InterleaveAPIEditRequest(
+            prompt=request.refine_prompt,
+            image=image_base64,
+            width=self.width,
+            height=self.height,
+            num_inference_steps=self.num_inference_steps,
+            guidance_scale=self.guidance_scale,
+            output_format="png",
+            enhance_prompt=False,
+        )
+        generated = self.editor.generate(
+            edit_request,
+            request_id=f"reward_{request.index}_{uuid.uuid4().hex}",
+        )
+        if not generated.file_path:
+            return None
+        return self.scorer.score_images(
+            reference_image_path=input_image_path if input_image_path else None,
+            edited_image_path=generated.file_path,
+            instruction=request.origin_prompt or request.refine_prompt,
+        )
+
+    def _is_text_to_image_canvas(self, path: str) -> bool:
+        if not self.treat_white_canvas_as_text_to_image:
+            return False
+        normalized = path.replace(os.sep, "/")
+        return normalized.endswith("data/interleave/white.png") or Path(path).name == "white.png"
+
+
+class ConstantInterleaveEditScorer:
+    """Small deterministic scorer for smoke tests and offline debugging."""
+
+    def __init__(
+        self,
+        *,
+        semantic_reward: float = 0.5,
+        quality_reward: float = 0.5,
+    ) -> None:
+        self.semantic_reward = float(semantic_reward)
+        self.quality_reward = float(quality_reward)
+
+    def __call__(
+        self,
+        request: InterleaveThinkerEditRequest,
+    ) -> InterleaveThinkerEditScore:
+        del request
+        return InterleaveThinkerEditScore(
+            semantic_reward=self.semantic_reward,
+            quality_reward=self.quality_reward,
+        )
+
+
+def _resolve_google_api_key(explicit: str | None = None) -> str:
+    if explicit:
+        return explicit.strip()
+    for env_name in ("GEMINI_API_KEY", "GOOGLE_API_KEY"):
+        value = os.environ.get(env_name)
+        if value:
+            return value.strip()
+    token_path = Path("~/.gemini_token").expanduser()
+    if token_path.is_file():
+        return token_path.read_text().strip()
+    raise ValueError("Gemini API wrappers require GEMINI_API_KEY, GOOGLE_API_KEY, "
+                     "an explicit api_key, or ~/.gemini_token.")
+
+
+def _import_google_genai() -> tuple[Any, Any]:
+    try:
+        from google import genai
+        from google.genai import types
+    except ImportError as exc:
+        raise RuntimeError("Gemini / Nano Banana API wrappers require google-genai. "
+                           "Install with `uv pip install -e '.[interleave-api]'`.") from exc
+    return genai, types
+
+
+def _load_jsonish(raw: str) -> dict[str, Any]:
+    text = (raw or "").strip()
+    if text.startswith("```"):
+        text = text.strip("`")
+        if text.startswith("json"):
+            text = text[4:].strip()
+    try:
+        payload = json.loads(text)
+    except json.JSONDecodeError:
+        start = text.find("{")
+        end = text.rfind("}")
+        if start < 0 or end <= start:
+            raise
+        payload = json.loads(text[start:end + 1])
+    if not isinstance(payload, dict):
+        raise ValueError(f"Gemini scorer returned non-object JSON: {type(payload).__name__}")
+    return payload
+
+
+def image_bytes_to_base64(data: bytes) -> str:
+    return base64.b64encode(data).decode("utf-8")
+
+
+def image_to_png_base64(image: Any) -> str:
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    return image_bytes_to_base64(buffer.getvalue())
+
+
+__all__ = [
+    "ConstantInterleaveEditScorer",
+    "GeminiInterleaveImageScorer",
+    "GeminiNanoBananaEditScorer",
+    "image_bytes_to_base64",
+    "image_to_png_base64",
+]

fastvideo/train/models/interleave_thinker/__init__.py

@@ -2,7 +2,8 @@
 """InterleaveThinker training model adapters."""
 
 from fastvideo.train.models.interleave_thinker.critic import (
+    INTERLEAVE_CRITIC_PROMPT,
     InterleaveThinkerCriticModel,
 )
 
-__all__ = ["InterleaveThinkerCriticModel"]
+__all__ = ["INTERLEAVE_CRITIC_PROMPT", "InterleaveThinkerCriticModel"]

fastvideo/train/models/interleave_thinker/critic.py

@@ -1,65 +1,407 @@
 # SPDX-License-Identifier: Apache-2.0
-"""InterleaveThinker critic actor adapter shell."""
+"""InterleaveThinker Qwen3-VL critic actor adapter."""
 
 from __future__ import annotations
 
+from collections.abc import Mapping, Sequence
+import json
+import os
+from pathlib import Path
+import re
 from types import SimpleNamespace
 from typing import Any, Literal, TYPE_CHECKING
 
 import torch
+from torch.utils.data import DataLoader, Dataset
 
 from fastvideo.pipelines import TrainingBatch
 from fastvideo.train.models.base import ModelBase
 
 if TYPE_CHECKING:
+    from fastvideo.train.utils.lora import LoraConfig
     from fastvideo.train.utils.training_config import TrainingConfig
 
+INTERLEAVE_CRITIC_PROMPT = """
+<image><image>
+# Generation/Edit Evaluation and Prompt Refinement System
 
-class _InterleaveThinkerActorModule(torch.nn.Module):
-    """Checkpoint-visible placeholder for future actor backends."""
+You are an expert image editing evaluator and prompt engineer. Your task is to:
+1. Evaluate the edited image and output the result in boolean format (True/False).
+2. If you think the edited image is not good enough (False), generate an optimized rewritten prompt that addresses the original shortcomings; if you think it is good enough (True), output the original rewritten prompt.
+
+## Input Information
+You have been presented with two images in sequence:
+- Original Image: The input image before editing. For the initial generation step, this may be a pure white or blank canvas.
+- Generated/Edited Image: The resulting image after applying the instruction or prompt.
+
+Original User Instruction: "{original_instruction}"
+Rewritten Prompt: "{rewritten_prompt}"
+
+## Evaluation Instructions
+Check both intent matching and anomaly / logic errors. The step is successful only if the generated or edited image satisfies the instruction and has no fatal artifacts, collateral damage, impossible anatomy, identity loss, or unrelated changes.
+
+## Prompt Refinement Strategy
+If the previous step is not good enough, explain what failed and produce a clearer rewritten prompt that fixes the issue while preserving what worked. If the previous step is good enough, return the original rewritten prompt.
+
+## Output
+Return exactly:
+<think>
+Detailed evaluation and refinement reasoning.
+</think>
+<answer>
+{{"previous_step_success": true_or_false, "refine_prompt": "prompt text"}}
+</answer>
+"""
+
+
+class _InterleaveJSONLDataset(Dataset):
+    """Tiny JSON/JSONL dataset for InterleaveThinker critic RL records."""
+
+    def __init__(
+        self,
+        data_path: str,
+    ) -> None:
+        self.records = list(_load_interleave_records(data_path))
+        if not self.records:
+            raise ValueError(f"No InterleaveThinker records found at {data_path!r}")
+
+    def __len__(self) -> int:
+        return len(self.records)
+
+    def __getitem__(
+        self,
+        index: int,
+    ) -> dict[str, Any]:
+        return dict(self.records[index])
+
+
+class _PlaceholderActorModule(torch.nn.Module):
+    """Checkpoint-visible placeholder used when backend loading is disabled."""
 
 
 class InterleaveThinkerCriticModel(ModelBase):
-    """ModelBase adapter slot for InterleaveThinker critic RL.
+    """Qwen3-VL actor wrapper for InterleaveThinker critic RL.
 
-    Upstream InterleaveThinker trains a Qwen3-VL critic with EasyR1. FastVideo
-    does not yet ship a native Qwen3-VL actor wrapper, so this class documents
-    the hooks consumed by :class:`InterleaveThinkerRLMethod` and gives YAML
-    configs an importable model target. A concrete backend should override
-    ``generate_interleave_responses`` and ``train_interleave_rollouts``.
+    The wrapper loads an InterleaveThinker critic checkpoint through
+    Transformers and exposes the two hooks consumed by
+    :class:`InterleaveThinkerRLMethod`:
+
+    - ``generate_interleave_responses`` for grouped critic rollouts.
+    - ``train_interleave_rollouts`` for an advantage-weighted policy update.
+
+    Set ``load_backend=false`` in tests or dry-run configs to skip heavyweight
+    checkpoint loading while keeping the YAML target importable.
     """
 
     def __init__(
         self,
         *,
-        init_from: str = "",
+        init_from: str = "InterleaveThinker/Critic-SFT-8B",
+        processor_from: str | None = None,
         training_config: TrainingConfig | None = None,
         trainable: bool = True,
+        load_backend: bool = True,
+        torch_dtype: str = "auto",
+        device_map: str | dict[str, Any] | None = None,
+        attn_implementation: str | None = None,
+        trust_remote_code: bool = False,
+        use_cache: bool = False,
+        freeze_vision_tower: bool = True,
+        freeze_multi_modal_projector: bool = True,
+        enable_gradient_checkpointing: bool = True,
+        max_prompt_length: int = 16384,
+        max_response_length: int = 4096,
+        prompt_template: str = INTERLEAVE_CRITIC_PROMPT,
+        lora: LoraConfig | dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         del kwargs
-        super().__init__(trainable=trainable)
-        self.init_from = str(init_from or "")
+        super().__init__(trainable=trainable, lora=lora)
+        self.init_from = str(init_from)
+        self.processor_from = str(processor_from or init_from)
         self.training_config = training_config
-        self.transformer = _InterleaveThinkerActorModule()
+        self.max_prompt_length = int(max_prompt_length)
+        self.max_response_length = int(max_response_length)
+        self.prompt_template = prompt_template
+        self.processor: Any | None = None
+        self.dataloader: Any = None
+        self.start_step = 0
         self.noise_scheduler = SimpleNamespace(num_train_timesteps=0)
 
+        if not load_backend:
+            self.transformer = _PlaceholderActorModule()
+            return
+
+        self.processor, self.transformer = self._load_backend(
+            init_from=self.init_from,
+            processor_from=self.processor_from,
+            torch_dtype=torch_dtype,
+            device_map=device_map,
+            attn_implementation=attn_implementation,
+            trust_remote_code=trust_remote_code,
+            use_cache=use_cache,
+        )
+        self._apply_trainable_policy(
+            trainable=trainable,
+            freeze_vision_tower=freeze_vision_tower,
+            freeze_multi_modal_projector=freeze_multi_modal_projector,
+        )
+        if trainable and enable_gradient_checkpointing and hasattr(self.transformer, "gradient_checkpointing_enable"):
+            self.transformer.gradient_checkpointing_enable()
+        if self._enable_lora_if_configured(self.transformer):
+            return
+
+    def _load_backend(
+        self,
+        *,
+        init_from: str,
+        processor_from: str,
+        torch_dtype: str,
+        device_map: str | dict[str, Any] | None,
+        attn_implementation: str | None,
+        trust_remote_code: bool,
+        use_cache: bool,
+    ) -> tuple[Any, torch.nn.Module]:
+        from transformers import AutoProcessor
+        import transformers
+
+        processor = AutoProcessor.from_pretrained(
+            processor_from,
+            trust_remote_code=trust_remote_code,
+        )
+        model_cls = getattr(transformers, "AutoModelForImageTextToText", None)
+        if model_cls is None:
+            model_cls = transformers.AutoModelForCausalLM
+        load_kwargs: dict[str, Any] = {
+            "trust_remote_code": trust_remote_code,
+            "use_cache": use_cache,
+        }
+        dtype_arg = _resolve_transformers_dtype(torch_dtype)
+        if dtype_arg is not None:
+            load_kwargs["dtype"] = dtype_arg
+        if device_map is not None:
+            load_kwargs["device_map"] = device_map
+        if attn_implementation:
+            load_kwargs["attn_implementation"] = attn_implementation
+        model = model_cls.from_pretrained(init_from, **load_kwargs)
+        if device_map is None:
+            model = model.to(self.device)
+        return processor, model
+
+    def _apply_trainable_policy(
+        self,
+        *,
+        trainable: bool,
+        freeze_vision_tower: bool,
+        freeze_multi_modal_projector: bool,
+    ) -> None:
+        for name, param in self.transformer.named_parameters():
+            requires_grad = bool(trainable)
+            lower_name = name.lower()
+            if freeze_vision_tower and _is_vision_parameter(lower_name):
+                requires_grad = False
+            if freeze_multi_modal_projector and _is_mm_projector_parameter(lower_name):
+                requires_grad = False
+            param.requires_grad_(requires_grad)
+
+    def init_preprocessors(
+        self,
+        training_config: TrainingConfig,
+    ) -> None:
+        self.training_config = training_config
+        data_path = str(getattr(training_config.data, "data_path", "") or "")
+        if not data_path:
+            return
+        if not os.path.exists(os.path.expanduser(data_path)):
+            raise FileNotFoundError(f"InterleaveThinker data_path not found: {data_path}")
+        dataset = _InterleaveJSONLDataset(data_path)
+        self.dataloader = DataLoader(
+            dataset,
+            batch_size=max(1, int(training_config.data.train_batch_size or 1)),
+            shuffle=True,
+            num_workers=max(0, int(training_config.data.dataloader_num_workers or 0)),
+            collate_fn=_collate_interleave_records,
+        )
+
+    @torch.no_grad()
     def generate_interleave_responses(
         self,
         batch: dict[str, Any],
         **kwargs: Any,
     ) -> list[dict[str, Any]]:
-        del batch, kwargs
-        raise NotImplementedError("InterleaveThinkerCriticModel requires a concrete Qwen/VLM actor backend "
-                                  "that implements generate_interleave_responses().")
+        self._require_backend()
+        self.transformer.eval()
+        num_generations = max(1, int(kwargs.get("num_generations", 1) or 1))
+        temperature = float(kwargs.get("temperature", 1.0) or 1.0)
+        top_p = float(kwargs.get("top_p", 1.0) or 1.0)
+        max_new_tokens = int(kwargs.get("max_new_tokens") or self.max_response_length)
+
+        rollouts: list[dict[str, Any]] = []
+        for item_idx, item in enumerate(_batch_to_items(batch)):
+            messages = self.build_messages(item)
+            inputs = self._apply_chat_template(
+                messages,
+                add_generation_prompt=True,
+            )
+            input_ids = inputs["input_ids"]
+            generate_kwargs: dict[str, Any] = {
+                "max_new_tokens": max_new_tokens,
+                "num_return_sequences": num_generations,
+                "do_sample": num_generations > 1 or temperature > 0.0,
+                "temperature": temperature,
+                "top_p": top_p,
+            }
+            pad_token_id = getattr(getattr(self.processor, "tokenizer", None), "pad_token_id", None)
+            eos_token_id = getattr(getattr(self.processor, "tokenizer", None), "eos_token_id", None)
+            if pad_token_id is not None:
+                generate_kwargs["pad_token_id"] = pad_token_id
+            if eos_token_id is not None:
+                generate_kwargs["eos_token_id"] = eos_token_id
+            generated_ids = self.transformer.generate(**inputs, **generate_kwargs)
+            prompt_len = int(input_ids.shape[-1])
+            decoded = self.processor.batch_decode(
+                [output_ids[prompt_len:] for output_ids in generated_ids],
+                skip_special_tokens=True,
+                clean_up_tokenization_spaces=False,
+            )
+            for generation_idx, response in enumerate(decoded):
+                rollout = dict(item)
+                rollout["response"] = str(response).strip()
+                rollout.setdefault("sample_index", item_idx)
+                rollout.setdefault("generation_index", generation_idx)
+                rollout.setdefault("group_key", _rollout_group_key(rollout, item_idx))
+                rollouts.append(rollout)
+        if self._trainable:
+            self.transformer.train()
+        return rollouts
 
     def train_interleave_rollouts(
         self,
         **kwargs: Any,
     ) -> tuple[dict[str, torch.Tensor], dict[str, float]]:
-        del kwargs
-        raise NotImplementedError("InterleaveThinkerCriticModel requires a concrete Qwen/VLM actor backend "
-                                  "that implements train_interleave_rollouts().")
+        self._require_backend()
+        rollouts = kwargs["rollouts"]
+        advantages = kwargs["advantages"].detach().to(self.device).float()
+        optimizer = kwargs.get("optimizer")
+        lr_scheduler = kwargs.get("lr_scheduler")
+        max_grad_norm = float(kwargs.get("max_grad_norm", 0.0) or 0.0)
+        grad_accum = max(1, int(kwargs.get("gradient_accumulation_steps", 1) or 1))
+
+        if optimizer is None:
+            raise RuntimeError("InterleaveThinkerCriticModel.train_interleave_rollouts() requires an optimizer")
+        optimizer.zero_grad(set_to_none=True)
+        self.transformer.train()
+
+        weighted_losses: list[torch.Tensor] = []
+        nll_values: list[torch.Tensor] = []
+        token_counts: list[int] = []
+        for idx, rollout in enumerate(rollouts):
+            nll, token_count = self._response_nll(rollout)
+            weighted = nll * advantages[idx]
+            (weighted / grad_accum).backward()
+            weighted_losses.append(weighted.detach())
+            nll_values.append(nll.detach())
+            token_counts.append(token_count)
+
+        if max_grad_norm > 0.0:
+            torch.nn.utils.clip_grad_norm_(
+                [p for p in self.transformer.parameters() if p.requires_grad],
+                max_grad_norm,
+            )
+        optimizer.step()
+        if lr_scheduler is not None:
+            lr_scheduler.step()
+
+        total_loss = torch.stack(weighted_losses).mean() if weighted_losses else torch.zeros((), device=self.device)
+        mean_nll = torch.stack(nll_values).mean() if nll_values else torch.zeros((), device=self.device)
+        return (
+            {
+                "total_loss": total_loss.detach()
+            },
+            {
+                "actor/policy_loss": float(total_loss.detach().cpu()),
+                "actor/response_nll": float(mean_nll.detach().cpu()),
+                "actor/mean_advantage": float(advantages.mean().detach().cpu()),
+                "actor/response_tokens": float(sum(token_counts) / max(1, len(token_counts))),
+            },
+        )
+
+    def build_messages(
+        self,
+        item: Mapping[str, Any],
+    ) -> list[dict[str, Any]]:
+        prompt = self.prompt_template.format(
+            original_instruction=str(item.get("origin_prompt", item.get("prompt", "")) or ""),
+            rewritten_prompt=str(item.get("previous_prompt", item.get("rewritten_prompt", "")) or ""),
+        )
+        images = _item_image_paths(item)
+        content: list[dict[str, Any]] = []
+        parts = re.split(r"(?i)<image>", prompt)
+        for idx, text in enumerate(parts):
+            if text:
+                content.append({
+                    "type": "text",
+                    "text": text,
+                })
+            if idx < len(images):
+                content.append({
+                    "type": "image",
+                    "image": images[idx],
+                })
+        return [{
+            "role": "user",
+            "content": content,
+        }]
+
+    def _response_nll(
+        self,
+        rollout: Mapping[str, Any],
+    ) -> tuple[torch.Tensor, int]:
+        response = str(rollout.get("response", "") or "")
+        if not response:
+            raise ValueError("InterleaveThinker rollout response is empty")
+        prompt_messages = self.build_messages(rollout)
+        full_messages = [*prompt_messages, {
+            "role": "assistant",
+            "content": response,
+        }]
+        prompt_inputs = self._apply_chat_template(
+            prompt_messages,
+            add_generation_prompt=True,
+        )
+        full_inputs = self._apply_chat_template(
+            full_messages,
+            add_generation_prompt=False,
+        )
+        labels = full_inputs["input_ids"].clone()
+        prompt_len = int(prompt_inputs["input_ids"].shape[-1])
+        labels[:, :prompt_len] = -100
+        token_count = int((labels != -100).sum().detach().cpu())
+        if token_count == 0:
+            raise ValueError("InterleaveThinker rollout has no trainable response tokens")
+        outputs = self.transformer(**full_inputs, labels=labels)
+        return outputs.loss, token_count
+
+    def _apply_chat_template(
+        self,
+        messages: list[dict[str, Any]],
+        *,
+        add_generation_prompt: bool,
+    ) -> Mapping[str, torch.Tensor]:
+        assert self.processor is not None
+        inputs = self.processor.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=add_generation_prompt,
+            return_dict=True,
+            return_tensors="pt",
+        )
+        return _move_to_device(inputs, self.device)
+
+    def _require_backend(self) -> None:
+        if self.processor is None or isinstance(self.transformer, _PlaceholderActorModule):
+            raise RuntimeError("InterleaveThinkerCriticModel was created with load_backend=false; "
+                               "enable backend loading for generation or RL updates.")
 
     def prepare_batch(
         self,
@@ -102,3 +444,108 @@ class InterleaveThinkerCriticModel(ModelBase):
     ) -> None:
         del loss, ctx, grad_accum_rounds
         raise NotImplementedError("InterleaveThinkerCriticModel uses train_interleave_rollouts().")
+
+
+def _load_interleave_records(data_path: str, ) -> list[dict[str, Any]]:
+    path = Path(os.path.expanduser(data_path))
+    files = sorted(path.glob("*.json*")) if path.is_dir() else [path]
+    records: list[dict[str, Any]] = []
+    for file_path in files:
+        if file_path.suffix == ".jsonl":
+            for line in file_path.read_text().splitlines():
+                if line.strip():
+                    raw = json.loads(line)
+                    if isinstance(raw, Mapping):
+                        records.append(dict(raw))
+        elif file_path.suffix == ".json":
+            raw_json = json.loads(file_path.read_text())
+            if isinstance(raw_json, list):
+                records.extend(dict(item) for item in raw_json if isinstance(item, Mapping))
+            elif isinstance(raw_json, Mapping):
+                records.append(dict(raw_json))
+        else:
+            raise ValueError(f"Unsupported InterleaveThinker data file: {file_path}")
+    return records
+
+
+def _collate_interleave_records(records: Sequence[Mapping[str, Any]], ) -> dict[str, Any]:
+    return {"items": [dict(record) for record in records]}
+
+
+def _batch_to_items(batch: Mapping[str, Any], ) -> list[dict[str, Any]]:
+    if "items" in batch and isinstance(batch["items"], Sequence):
+        return [dict(item) for item in batch["items"]]
+    length = 1
+    for value in batch.values():
+        if isinstance(value, Sequence) and not isinstance(value, str | bytes):
+            length = len(value)
+            break
+    items: list[dict[str, Any]] = []
+    for idx in range(length):
+        item: dict[str, Any] = {}
+        for key, value in batch.items():
+            if isinstance(value, Sequence) and not isinstance(value, str | bytes) and len(value) == length:
+                item[key] = value[idx]
+            else:
+                item[key] = value
+        items.append(item)
+    return items
+
+
+def _item_image_paths(item: Mapping[str, Any], ) -> list[str]:
+    image_paths: list[str] = []
+    for key in ("origin_image_path", "previous_image_path", "edited_image_path"):
+        value = item.get(key)
+        if isinstance(value, str) and value:
+            image_paths.append(value)
+    if len(image_paths) == 1:
+        image_paths.append(image_paths[0])
+    return image_paths[:2]
+
+
+def _move_to_device(
+    value: Any,
+    device: torch.device,
+) -> Any:
+    if torch.is_tensor(value):
+        return value.to(device)
+    if hasattr(value, "to"):
+        return value.to(device)
+    if isinstance(value, Mapping):
+        return {key: _move_to_device(item, device) for key, item in value.items()}
+    return value
+
+
+def _resolve_transformers_dtype(raw: str, ) -> Any:
+    value = str(raw or "auto").lower()
+    if value == "auto":
+        return "auto"
+    if value in {"bf16", "bfloat16"}:
+        return torch.bfloat16
+    if value in {"fp16", "float16"}:
+        return torch.float16
+    if value in {"fp32", "float32"}:
+        return torch.float32
+    raise ValueError(f"Unsupported torch_dtype for InterleaveThinker critic: {raw!r}")
+
+
+def _is_vision_parameter(lower_name: str, ) -> bool:
+    return any(token in lower_name for token in ("visual", "vision_tower", "vision_model"))
+
+
+def _is_mm_projector_parameter(lower_name: str, ) -> bool:
+    return any(token in lower_name for token in ("mm_projector", "multi_modal_projector", "visual_merger"))
+
+
+def _rollout_group_key(
+    rollout: Mapping[str, Any],
+    index: int,
+) -> str:
+    for key in ("group_key", "problem_id", "sample_index", "origin_prompt", "prompt"):
+        value = rollout.get(key)
+        if value is not None:
+            return str(value)
+    return str(index)
+
+
+__all__ = ["INTERLEAVE_CRITIC_PROMPT", "InterleaveThinkerCriticModel"]

pyproject.toml

@@ -156,6 +156,7 @@ eval = [
 ]
 eval-fast-decode = ["decord"]
 eval-full = ["fastvideo[eval]", "fastvideo[eval-audio]", "qwen-omni-utils"]
+interleave-api = ["google-genai"]
 
 dev = [ "fastvideo[lint]", "fastvideo[test]", ]
 

tests/local_tests/test_interleave_thinker_api_models.py

@@ -0,0 +1,153 @@
+import base64
+import io
+import sys
+import types
+
+from PIL import Image
+import pytest
+
+from fastvideo.entrypoints.interleave.generator import (
+    NanoBananaImageGeneratorBackend,
+)
+from fastvideo.entrypoints.interleave.schema import InterleaveEditRequest
+from fastvideo.train.methods.rl.rewards.interleave_api import (
+    GeminiNanoBananaEditScorer,
+)
+from fastvideo.train.methods.rl.rewards.interleave_thinker import (
+    InterleaveThinkerEditRequest,
+)
+
+
+def _png_base64(color="red"):
+    image = Image.new("RGB", (8, 8), color)
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+def _install_fake_google_genai(monkeypatch):
+    captured = {
+        "calls": [],
+        "configs": [],
+        "parts": [],
+    }
+
+    class FakeImagePart:
+
+        def __init__(self, image):
+            self._image = image
+
+        def as_image(self):
+            return self._image
+
+    class FakePartFactory:
+
+        @staticmethod
+        def from_bytes(data, mime_type):
+            captured["parts"].append((data, mime_type))
+            return {
+                "data": data,
+                "mime_type": mime_type,
+            }
+
+    class FakeImageConfig:
+
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+    class FakeGenerateContentConfig:
+
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+            captured["configs"].append(kwargs)
+
+    class FakeTypes:
+        GenerateContentConfig = FakeGenerateContentConfig
+        ImageConfig = FakeImageConfig
+        Part = FakePartFactory
+
+    class FakeModels:
+
+        def generate_content(self, **kwargs):
+            captured["calls"].append(kwargs)
+            model = kwargs["model"]
+            if "image" in model:
+                return types.SimpleNamespace(parts=[FakeImagePart(Image.new("RGB", (8, 8), "blue"))])
+            return types.SimpleNamespace(text='{"semantic_score": 8.0, "quality_score": 7.0}')
+
+    class FakeClient:
+
+        def __init__(self, **kwargs):
+            captured["client_kwargs"] = kwargs
+            self.models = FakeModels()
+
+    google_mod = types.ModuleType("google")
+    genai_mod = types.ModuleType("google.genai")
+    genai_mod.Client = FakeClient
+    genai_mod.types = FakeTypes
+    google_mod.genai = genai_mod
+
+    monkeypatch.setitem(sys.modules, "google", google_mod)
+    monkeypatch.setitem(sys.modules, "google.genai", genai_mod)
+    return captured
+
+
+def test_nano_banana_backend_wraps_google_genai_image_api(monkeypatch, tmp_path):
+    captured = _install_fake_google_genai(monkeypatch)
+    monkeypatch.setenv("GEMINI_API_KEY", "test-key")
+    backend = NanoBananaImageGeneratorBackend(
+        model="nano-banana-2",
+        output_dir=str(tmp_path),
+        aspect_ratio="1:1",
+        image_size="1K",
+    )
+
+    result = backend.generate(
+        InterleaveEditRequest(
+            prompt="make a blue square",
+            image=_png_base64("red"),
+        ),
+        request_id="abc",
+    )
+
+    assert result.file_path is not None
+    assert Image.open(result.file_path).getpixel((0, 0)) == (0, 0, 255)
+    assert captured["client_kwargs"]["api_key"] == "test-key"
+    assert captured["calls"][0]["model"] == "gemini-3.1-flash-image"
+    assert captured["configs"][0]["image_config"].kwargs == {
+        "aspect_ratio": "1:1",
+        "image_size": "1K",
+    }
+
+
+def test_gemini_nano_banana_edit_scorer_generates_and_scores(monkeypatch, tmp_path):
+    captured = _install_fake_google_genai(monkeypatch)
+    monkeypatch.setenv("GEMINI_API_KEY", "test-key")
+    source = tmp_path / "source.png"
+    Image.new("RGB", (8, 8), "white").save(source)
+    scorer = GeminiNanoBananaEditScorer(
+        image_model="nano-banana",
+        judge_model="gemini-2.5-pro",
+        output_dir=str(tmp_path / "reward"),
+        max_attempts=1,
+    )
+
+    score = scorer(
+        InterleaveThinkerEditRequest(
+            index=0,
+            origin_prompt="draw a blue square",
+            previous_prompt="blue square",
+            refine_prompt="make it a clean blue square",
+            origin_image_path=str(source),
+            previous_image_path=str(source),
+            previous_step_success=False,
+            previous_semantic_score=4.0,
+            previous_quality_score=5.0,
+        ))
+
+    assert score is not None
+    assert score.semantic_score == pytest.approx(8.0)
+    assert score.quality_score == pytest.approx(7.0)
+    assert captured["calls"][0]["model"] == "gemini-2.5-flash-image"
+    assert captured["calls"][1]["model"] == "gemini-2.5-pro"
+    assert len(captured["parts"]) == 2

tests/local_tests/test_interleave_thinker_critic_model.py

@@ -0,0 +1,59 @@
+import json
+
+from fastvideo.train.models.interleave_thinker import (
+    INTERLEAVE_CRITIC_PROMPT,
+    InterleaveThinkerCriticModel,
+)
+from fastvideo.train.utils.training_config import (
+    DataConfig,
+    TrainingConfig,
+)
+
+
+def test_interleave_thinker_critic_builds_qwen_vl_messages_without_loading_backend():
+    model = InterleaveThinkerCriticModel(load_backend=False)
+
+    messages = model.build_messages({
+        "origin_prompt": "draw a vase",
+        "previous_prompt": "a vase on a table",
+        "origin_image_path": "before.png",
+        "edited_image_path": "after.png",
+    })
+
+    content = messages[0]["content"]
+    images = [part for part in content if part["type"] == "image"]
+    text = "\n".join(part["text"] for part in content if part["type"] == "text")
+    assert messages[0]["role"] == "user"
+    assert images == [{
+        "type": "image",
+        "image": "before.png",
+    }, {
+        "type": "image",
+        "image": "after.png",
+    }]
+    assert "draw a vase" in text
+    assert "a vase on a table" in text
+
+
+def test_interleave_thinker_critic_initializes_jsonl_dataloader(tmp_path):
+    data_path = tmp_path / "critic_rl.jsonl"
+    data_path.write_text(
+        json.dumps({
+            "origin_prompt": "draw a chair",
+            "rewritten_prompt": "a wooden chair",
+            "origin_image_path": "before.png",
+            "edited_image_path": "after.png",
+            "evaluation": {
+                "success": True,
+                "semantics": 7.0,
+                "quality": 8.0,
+            },
+        }) + "\n")
+    model = InterleaveThinkerCriticModel(load_backend=False)
+
+    model.init_preprocessors(TrainingConfig(data=DataConfig(data_path=str(data_path), train_batch_size=1)))
+    batch = next(iter(model.dataloader))
+
+    assert batch["items"][0]["origin_prompt"] == "draw a chair"
+    assert batch["items"][0]["evaluation"]["success"] is True
+    assert "{original_instruction}" in INTERLEAVE_CRITIC_PROMPT

tests/local_tests/test_interleave_thinker_method.py

@@ -184,12 +184,43 @@ def test_interleave_thinker_method_can_use_offline_response_batches():
     assert len(actor.train_calls) == 1
 
 
+def test_interleave_thinker_method_instantiates_configured_edit_scorer():
+    actor = _FakeInterleaveActor()
+    actor.generate_interleave_responses = None
+    method = InterleaveThinkerRLMethod(
+        cfg=_cfg({
+            "edit_scorer": {
+                "_target_": "fastvideo.train.methods.rl.rewards.ConstantInterleaveEditScorer",
+                "semantic_reward": 0.9,
+                "quality_reward": 0.1,
+            }
+        }),
+        role_models={"student": actor},
+    )
+    batch = {
+        "origin_prompt": ["prompt-a"],
+        "previous_prompt": ["prompt-a previous"],
+        "ground_truth": [{
+            "success": True
+        }],
+        "response": [_response(True)],
+    }
+
+    _, _, metrics = method.managed_train_step(iter([batch]), iteration=2)
+
+    assert torch.isclose(metrics["interleave/reward/edited_image_reward_semantic"], torch.tensor(0.9))
+    assert torch.isclose(metrics["interleave/reward/edited_image_reward_quality"], torch.tensor(0.1))
+
+
 def test_interleave_thinker_config_parses_public_yaml():
     cfg = load_run_config("examples/train/configs/rl/interleave_thinker/critic_grpo.yaml")
 
     assert cfg.models["student"]["_target_"] == (
         "fastvideo.train.models.interleave_thinker.InterleaveThinkerCriticModel")
+    assert cfg.models["student"]["init_from"] == "InterleaveThinker/Critic-SFT-8B"
     assert cfg.method["_target_"] == "fastvideo.train.methods.rl.interleave_thinker.InterleaveThinkerRLMethod"
+    assert cfg.method["edit_scorer"]["_target_"] == (
+        "fastvideo.train.methods.rl.rewards.GeminiNanoBananaEditScorer")
     assert cfg.method["num_generations"] == 8
     assert cfg.training.data.data_path.endswith("critic_rl.jsonl")
     assert cfg.training.optimizer.learning_rate == 2.0e-6