[feat]: add InterleaveThinker workflow core
This commit is contained in:
Mac Lee
@@ -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")
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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),
|
||||
}
|
||||
Reference in New Issue
Block a user