From bb440bf0b7aa0b09bfe896121a9b88a7e779fbd1 Mon Sep 17 00:00:00 2001 From: Mac Lee Date: Mon, 29 Jun 2026 04:36:24 +0000 Subject: [PATCH] [feat]: add InterleaveThinker workflow core --- .../workflow/interleave_thinker/generator.py | 337 ++++++++++++++++++ .../interleave_thinker/orchestrator.py | 190 ++++++++++ .../workflow/interleave_thinker/schema.py | 107 ++++++ .../workflow/interleave_thinker/trace.py | 102 ++++++ 4 files changed, 736 insertions(+) create mode 100644 fastvideo/workflow/interleave_thinker/generator.py create mode 100644 fastvideo/workflow/interleave_thinker/orchestrator.py create mode 100644 fastvideo/workflow/interleave_thinker/schema.py create mode 100644 fastvideo/workflow/interleave_thinker/trace.py diff --git a/fastvideo/workflow/interleave_thinker/generator.py b/fastvideo/workflow/interleave_thinker/generator.py new file mode 100644 index 000000000..0c1905e0d --- /dev/null +++ b/fastvideo/workflow/interleave_thinker/generator.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastVideo generator adapter for InterleaveThinker-style image calls.""" + +from __future__ import annotations + +import base64 +import io +import os +import time +import uuid +from collections.abc import Mapping +from pathlib import Path +from typing import Any, Protocol + +from fastvideo.api.compat import ( + explicit_request_updates, + legacy_generate_call_to_request, + normalize_generation_request, +) +from fastvideo.api.results import GenerationResult +from fastvideo.api.schema import GenerationRequest +from fastvideo.workflow.interleave_thinker.schema import ( + GeneratedImage, + 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.""" + + def generate( + self, + request: InterleaveEditRequest, + *, + request_id: str | None = None, + ) -> GeneratedImage: + ... + + +class FastVideoImageGeneratorBackend: + """Translate InterleaveThinker image requests into ``VideoGenerator`` calls.""" + + def __init__( + self, + generator: Any, + *, + output_dir: str, + default_request: GenerationRequest | Mapping[str, Any] | None = None, + ) -> None: + self.generator = generator + self.output_dir = output_dir + self.default_request = normalize_generation_request(default_request) if default_request is not None else None + + def generate( + self, + request: InterleaveEditRequest, + *, + request_id: str | None = None, + ) -> GeneratedImage: + request_id = request_id or uuid.uuid4().hex + request_output_dir = os.path.join(self.output_dir, "interleave") + upload_dir = os.path.join(self.output_dir, "uploads") + os.makedirs(request_output_dir, exist_ok=True) + + input_path = None + if request.image: + os.makedirs(upload_dir, exist_ok=True) + input_path = decode_base64_image_to_path( + request.image, + os.path.join(upload_dir, f"{request_id}_input.png"), + ) + + output_path = os.path.join(request_output_dir, f"{request_id}.png") + generation_request = self._build_generation_request( + request, + output_path=output_path, + input_image_path=input_path, + ) + + start = time.perf_counter() + result = self.generator.generate(generation_request) + elapsed = time.perf_counter() - start + result = _first_generation_result(result) + file_path = result.video_path or output_path + if not file_path or not os.path.exists(file_path): + raise RuntimeError(f"FastVideo generation did not produce an image at {file_path!r}") + + return GeneratedImage( + prompt=request.prompt, + image_base64=encode_file_to_base64(file_path), + file_path=os.path.abspath(file_path), + inference_time_s=result.generation_time or elapsed, + metadata={ + "request_id": request_id, + "input_image_path": input_path, + "peak_memory_mb": result.peak_memory_mb, + }, + ) + + def _build_generation_request( + self, + request: InterleaveEditRequest, + *, + output_path: str, + input_image_path: str | None, + ) -> GenerationRequest: + kwargs = {} + if self.default_request is not None: + kwargs.update(_safe_explicit_request_updates(self.default_request)) + + kwargs.update({ + "num_frames": 1, + "fps": 1, + "save_video": True, + "return_frames": False, + "output_path": output_path, + }) + if input_image_path is not None: + kwargs["image_path"] = input_image_path + if request.width is not None: + kwargs["width"] = int(request.width) + if request.height is not None: + kwargs["height"] = int(request.height) + if request.seed is not None: + kwargs["seed"] = int(request.seed) + if request.resolved_num_inference_steps() is not None: + kwargs["num_inference_steps"] = int(request.resolved_num_inference_steps()) + if request.guidance_scale is not None: + kwargs["guidance_scale"] = float(request.guidance_scale) + if request.true_cfg_scale is not None: + kwargs["true_cfg_scale"] = float(request.true_cfg_scale) + if request.negative_prompt is not None: + kwargs["negative_prompt"] = request.negative_prompt + + return legacy_generate_call_to_request( + request.prompt, + None, + legacy_kwargs=kwargs, + ) + + +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 Interleave orchestration. 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() + client = self._client_instance() + for attempt in range(self.max_attempts): + try: + response = client.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) + except AssertionError: + return explicit_request_updates(normalize_generation_request(request)) + + +def _first_generation_result(result: GenerationResult | list[GenerationResult]) -> GenerationResult: + if isinstance(result, list): + if not result: + raise RuntimeError("FastVideo generation returned an empty result list") + return result[0] + return result + + +def encode_file_to_base64(path: str | os.PathLike[str]) -> str: + with open(path, "rb") as handle: + return base64.b64encode(handle.read()).decode("utf-8") + + +def decode_base64_image_to_path( + image_base64: str, + output_path: str | os.PathLike[str], +) -> str: + payload = image_base64.strip() + if payload.startswith("data:") and "," in payload: + payload = payload.split(",", 1)[1] + data = base64.b64decode(payload) + path = Path(output_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 google-genai directly or with `uv pip install -e '.[eval-judge]'`.") 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/workflow/interleave_thinker/orchestrator.py b/fastvideo/workflow/interleave_thinker/orchestrator.py new file mode 100644 index 000000000..4b67f1bc7 --- /dev/null +++ b/fastvideo/workflow/interleave_thinker/orchestrator.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Provider-based interleaved generation orchestration.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any, Protocol + +from fastvideo.workflow.interleave_thinker.generator import ( + ImageGeneratorBackend, + encode_file_to_base64, +) +from fastvideo.workflow.interleave_thinker.schema import ( + CriticDecision, + CriticInput, + GeneratedImage, + InterleaveAttempt, + InterleaveEditRequest, + InterleaveTrace, + PlannedInterleaveStep, + PlannerInput, +) + + +class PlannerProvider(Protocol): + """Plans a user instruction into concrete generator calls.""" + + def plan(self, request: PlannerInput) -> Sequence[PlannedInterleaveStep]: + ... + + +class CriticProvider(Protocol): + """Reviews one generated step and optionally proposes a refined prompt.""" + + def review(self, request: CriticInput) -> CriticDecision: + ... + + +class SinglePromptPlanner: + """Fallback planner that runs the instruction as one generator prompt.""" + + def __init__(self, *, max_attempts: int = 1) -> None: + self.max_attempts = max(1, int(max_attempts)) + + def plan(self, request: PlannerInput) -> Sequence[PlannedInterleaveStep]: + return [ + PlannedInterleaveStep( + prompt=request.instruction, + input_image_path=request.initial_image_path, + max_attempts=self.max_attempts, + ) + ] + + +class AcceptAllCritic: + """Fallback critic for smoke tests and simple generation flows.""" + + def review(self, request: CriticInput) -> CriticDecision: + del request + return CriticDecision(success=True) + + +class InterleaveOrchestrator: + """Run planner -> generator -> critic loops for interleaved workflows.""" + + def __init__( + self, + *, + planner: PlannerProvider, + generator: ImageGeneratorBackend, + critic: CriticProvider | None = None, + width: int | None = None, + height: int | None = None, + num_inference_steps: int | None = None, + guidance_scale: float | None = None, + seed: int | None = None, + ) -> None: + self.planner = planner + self.generator = generator + self.critic = critic + self.width = width + self.height = height + self.num_inference_steps = num_inference_steps + self.guidance_scale = guidance_scale + self.seed = seed + + def run( + self, + instruction: str, + *, + initial_image_path: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> InterleaveTrace: + planner_input = PlannerInput( + instruction=instruction, + initial_image_path=initial_image_path, + metadata=dict(metadata or {}), + ) + planned_steps = list(self.planner.plan(planner_input)) + attempts: list[InterleaveAttempt] = [] + previous_image_path = initial_image_path + final_image: GeneratedImage | None = None + + if not planned_steps: + return InterleaveTrace( + instruction=instruction, + attempts=[], + final_image=None, + success=False, + metadata={"error": "planner returned no steps"}, + ) + + for step_index, step in enumerate(planned_steps): + accepted = False + prompt = step.prompt + step_input_path = step.input_image_path or previous_image_path + max_attempts = max(1, int(step.max_attempts)) + + for attempt_index in range(max_attempts): + request = self._build_generation_request( + prompt, + input_image_path=step_input_path, + attempt_index=attempt_index, + ) + generated = self.generator.generate(request) + decision = None + if self.critic is not None: + decision = self.critic.review( + CriticInput( + step=step, + attempt_index=attempt_index, + generated=generated, + previous_image_path=step_input_path, + metadata=dict(step.metadata), + )) + + attempts.append( + InterleaveAttempt( + step_index=step_index, + attempt_index=attempt_index, + prompt=prompt, + generated=generated, + decision=decision, + )) + + if decision is None or decision.success: + accepted = True + final_image = generated + previous_image_path = generated.file_path or previous_image_path + break + + if decision.refine_prompt: + prompt = decision.refine_prompt + + if not accepted: + return InterleaveTrace( + instruction=instruction, + attempts=attempts, + final_image=final_image, + success=False, + metadata={"failed_step_index": step_index}, + ) + + return InterleaveTrace( + instruction=instruction, + attempts=attempts, + final_image=final_image, + success=True, + metadata=dict(metadata or {}), + ) + + def _build_generation_request( + self, + prompt: str, + *, + input_image_path: str | None, + attempt_index: int = 0, + ) -> InterleaveEditRequest: + seed = self.seed + if seed is not None: + seed = seed + attempt_index + return InterleaveEditRequest( + prompt=prompt, + image=(encode_file_to_base64(input_image_path) if input_image_path else None), + width=self.width, + height=self.height, + seed=seed, + num_inference_steps=self.num_inference_steps, + guidance_scale=self.guidance_scale, + ) diff --git a/fastvideo/workflow/interleave_thinker/schema.py b/fastvideo/workflow/interleave_thinker/schema.py new file mode 100644 index 000000000..6d4c7f068 --- /dev/null +++ b/fastvideo/workflow/interleave_thinker/schema.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Schemas for InterleaveThinker-style orchestration.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field + + +class InterleaveEditRequest(BaseModel): + """Image edit/generation request used by Interleave orchestration backends. + + InterleaveThinker sends `num_inference_step` while FastVideo uses + `num_inference_steps`; accept both and let the plural form win when both are + provided. Unknown fields are tolerated so model-specific knobs can pass + through without forcing every backend to implement them immediately. + """ + + model_config = ConfigDict(extra="allow") + + prompt: str + image: str | None = None + negative_prompt: str | None = None + width: int | None = None + height: int | None = None + seed: int | None = None + num_inference_step: int | None = None + num_inference_steps: int | None = None + guidance_scale: float | None = None + true_cfg_scale: float | None = None + output_format: Literal["png", "jpeg", "jpg", "webp"] | None = "png" + enhance_prompt: bool | None = None + + def resolved_num_inference_steps(self) -> int | None: + return self.num_inference_steps if self.num_inference_steps is not None else self.num_inference_step + + +class InterleaveEditResponse(BaseModel): + success: bool + edited_image: str | None = None + file_path: str | None = None + prompt: str | None = None + inference_time_s: float | None = None + error: str | None = None + metadata: dict[str, Any] = Field(default_factory=dict) + + +@dataclass(slots=True) +class GeneratedImage: + prompt: str + image_base64: str | None = None + file_path: str | None = None + inference_time_s: float | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class PlannedInterleaveStep: + prompt: str + name: str | None = None + input_image_path: str | None = None + max_attempts: int = 2 + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class PlannerInput: + instruction: str + initial_image_path: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class CriticInput: + step: PlannedInterleaveStep + attempt_index: int + generated: GeneratedImage + previous_image_path: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class CriticDecision: + success: bool + refine_prompt: str | None = None + reason: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class InterleaveAttempt: + step_index: int + attempt_index: int + prompt: str + generated: GeneratedImage + decision: CriticDecision | None = None + + +@dataclass(slots=True) +class InterleaveTrace: + instruction: str + attempts: list[InterleaveAttempt] + final_image: GeneratedImage | None + success: bool + metadata: dict[str, Any] = field(default_factory=dict) diff --git a/fastvideo/workflow/interleave_thinker/trace.py b/fastvideo/workflow/interleave_thinker/trace.py new file mode 100644 index 000000000..49efd821f --- /dev/null +++ b/fastvideo/workflow/interleave_thinker/trace.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Serialization helpers for interleaved generation traces.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from fastvideo.workflow.interleave_thinker.schema import ( + CriticDecision, + GeneratedImage, + InterleaveAttempt, + InterleaveTrace, +) + + +def trace_to_dict( + trace: InterleaveTrace, + *, + include_images: bool = False, +) -> dict[str, Any]: + return { + "instruction": trace.instruction, + "success": trace.success, + "final_image": _generated_image_to_dict( + trace.final_image, + include_images=include_images, + ), + "attempts": [_attempt_to_dict( + attempt, + include_images=include_images, + ) for attempt in trace.attempts], + "metadata": dict(trace.metadata), + } + + +def save_trace( + trace: InterleaveTrace, + path: str | Path, + *, + include_images: bool = False, +) -> None: + output_path = Path(path) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps( + trace_to_dict( + trace, + include_images=include_images, + ), + indent=2, + sort_keys=True, + ) + "\n", + encoding="utf-8", + ) + + +def _attempt_to_dict( + attempt: InterleaveAttempt, + *, + include_images: bool, +) -> dict[str, Any]: + return { + "step_index": attempt.step_index, + "attempt_index": attempt.attempt_index, + "prompt": attempt.prompt, + "generated": _generated_image_to_dict( + attempt.generated, + include_images=include_images, + ), + "decision": _critic_decision_to_dict(attempt.decision), + } + + +def _generated_image_to_dict( + image: GeneratedImage | None, + *, + include_images: bool, +) -> dict[str, Any] | None: + if image is None: + return None + result = { + "prompt": image.prompt, + "file_path": image.file_path, + "inference_time_s": image.inference_time_s, + "metadata": dict(image.metadata), + } + if include_images: + result["image_base64"] = image.image_base64 + return result + + +def _critic_decision_to_dict(decision: CriticDecision | None) -> dict[str, Any] | None: + if decision is None: + return None + return { + "success": decision.success, + "refine_prompt": decision.refine_prompt, + "reason": decision.reason, + "metadata": dict(decision.metadata), + }