diff --git a/.agents/exploration/interleavethinker-fastvideo-integration.md b/.agents/exploration/interleavethinker-fastvideo-integration.md index b74704fc4..01ff9c88d 100644 --- a/.agents/exploration/interleavethinker-fastvideo-integration.md +++ b/.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 `` + 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 diff --git a/examples/interleave/README.md b/examples/interleave/README.md index cb36c031d..0e83878a8 100644 --- a/examples/interleave/README.md +++ b/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: diff --git a/examples/train/configs/rl/interleave_thinker/critic_grpo.yaml b/examples/train/configs/rl/interleave_thinker/critic_grpo.yaml index bc67d9aef..5c924c252 100644 --- a/examples/train/configs/rl/interleave_thinker/critic_grpo.yaml +++ b/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 diff --git a/fastvideo/entrypoints/interleave/generator.py b/fastvideo/entrypoints/interleave/generator.py index 07c21c800..28c16cddb 100644 --- a/fastvideo/entrypoints/interleave/generator.py +++ b/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") diff --git a/fastvideo/train/methods/rl/interleave_thinker.py b/fastvideo/train/methods/rl/interleave_thinker.py index 213069b17..253a09880 100644 --- a/fastvideo/train/methods/rl/interleave_thinker.py +++ b/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), diff --git a/fastvideo/train/methods/rl/rewards/__init__.py b/fastvideo/train/methods/rl/rewards/__init__.py index 5c16cb02c..ad69c60f6 100644 --- a/fastvideo/train/methods/rl/rewards/__init__.py +++ b/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", diff --git a/fastvideo/train/methods/rl/rewards/interleave_api.py b/fastvideo/train/methods/rl/rewards/interleave_api.py new file mode 100644 index 000000000..e8aa0433e --- /dev/null +++ b/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", +] diff --git a/fastvideo/train/models/interleave_thinker/__init__.py b/fastvideo/train/models/interleave_thinker/__init__.py index ab3a94aaf..4ceee762e 100644 --- a/fastvideo/train/models/interleave_thinker/__init__.py +++ b/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"] diff --git a/fastvideo/train/models/interleave_thinker/critic.py b/fastvideo/train/models/interleave_thinker/critic.py index e67830ef8..1356dff36 100644 --- a/fastvideo/train/models/interleave_thinker/critic.py +++ b/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 = """ + +# 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: + +Detailed evaluation and refinement reasoning. + + +{{"previous_step_success": true_or_false, "refine_prompt": "prompt text"}} + +""" + + +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)", 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"] diff --git a/pyproject.toml b/pyproject.toml index 35d3033d2..d766d6b02 100644 --- a/pyproject.toml +++ b/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]", ] diff --git a/tests/local_tests/test_interleave_thinker_api_models.py b/tests/local_tests/test_interleave_thinker_api_models.py new file mode 100644 index 000000000..ba75b72a4 --- /dev/null +++ b/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 diff --git a/tests/local_tests/test_interleave_thinker_critic_model.py b/tests/local_tests/test_interleave_thinker_critic_model.py new file mode 100644 index 000000000..bec6a091b --- /dev/null +++ b/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 diff --git a/tests/local_tests/test_interleave_thinker_method.py b/tests/local_tests/test_interleave_thinker_method.py index 692ac209e..c46b24c1c 100644 --- a/tests/local_tests/test_interleave_thinker_method.py +++ b/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