From eb33639cbd5a93fead2b869484f92c7cccf6bb34 Mon Sep 17 00:00:00 2001 From: Mac Lee Date: Thu, 18 Jun 2026 02:40:54 +0000 Subject: [PATCH] [feat] add InterleaveThinker compatibility service --- ...interleavethinker-fastvideo-integration.md | 108 ++++++++-- examples/interleave/README.md | 43 ++++ .../flux2_klein_interleave_serve.yaml | 29 +++ fastvideo/entrypoints/cli/interleave_serve.py | 78 +++++++ fastvideo/entrypoints/cli/main.py | 5 +- fastvideo/entrypoints/interleave/__init__.py | 46 +++++ fastvideo/entrypoints/interleave/generator.py | 172 ++++++++++++++++ .../entrypoints/interleave/orchestrator.py | 182 +++++++++++++++++ fastvideo/entrypoints/interleave/schema.py | 108 ++++++++++ fastvideo/entrypoints/interleave/server.py | 154 ++++++++++++++ .../local_tests/test_interleave_entrypoint.py | 192 ++++++++++++++++++ 11 files changed, 1104 insertions(+), 13 deletions(-) create mode 100644 examples/interleave/README.md create mode 100644 examples/interleave/flux2_klein_interleave_serve.yaml create mode 100644 fastvideo/entrypoints/cli/interleave_serve.py create mode 100644 fastvideo/entrypoints/interleave/__init__.py create mode 100644 fastvideo/entrypoints/interleave/generator.py create mode 100644 fastvideo/entrypoints/interleave/orchestrator.py create mode 100644 fastvideo/entrypoints/interleave/schema.py create mode 100644 fastvideo/entrypoints/interleave/server.py create mode 100644 tests/local_tests/test_interleave_entrypoint.py diff --git a/.agents/exploration/interleavethinker-fastvideo-integration.md b/.agents/exploration/interleavethinker-fastvideo-integration.md index b30db8b1c..e733073b7 100644 --- a/.agents/exploration/interleavethinker-fastvideo-integration.md +++ b/.agents/exploration/interleavethinker-fastvideo-integration.md @@ -99,12 +99,12 @@ Testing plan: - [x] Created branch `interleavethinker-fastvideo` from committed `main`. - [x] Read root `AGENTS.md`, `fastvideo/AGENTS.md`, onboarding, codebase map, skills index, workflows list, lessons list, and exploration template. -- [ ] Inspect `fastvideo/entrypoints/cli`, `VideoGenerator`, schema, serve, and +- [x] Inspect `fastvideo/entrypoints/cli`, `VideoGenerator`, schema, serve, and streaming prompt provider patterns. -- [ ] Decide exact CLI/service shape based on existing FastVideo conventions. -- [ ] Implement first minimal app/service layer. -- [ ] Add unit tests with fakes. -- [ ] Run Modal-backed validation. +- [x] Decide exact CLI/service shape based on existing FastVideo conventions. +- [x] Implement first minimal app/service layer. +- [x] Add unit tests with fakes. +- [x] Run Modal-backed validation. - [ ] Commit first coherent checkpoint. ## Findings @@ -117,6 +117,35 @@ Testing plan: under `fastvideo/train`. - Existing prompt orchestration under `fastvideo/entrypoints/streaming/prompt` provides a good provider/fallback style to mirror for planner/critic clients. +- Existing OpenAI image routes expose `/v1/images` and `/v1/images/edits`, + but InterleaveThinker reward code posts JSON to `/edit` and expects + `{"success": true, "edited_image": ""}`. Add a compatibility + app instead of changing the OpenAI-compatible API surface. +- Use `ServeConfig` for the new compatibility server rather than inventing a + new top-level config. `fastvideo interleave-serve --config ` + can reuse `generator`, `server`, and `default_request`. +- The first code slice will include: + - `fastvideo/entrypoints/interleave/schema.py` for Pydantic request/response + models and dataclasses shared by the orchestrator. + - `fastvideo/entrypoints/interleave/generator.py` for `VideoGenerator` request + translation plus base64 image IO. + - `fastvideo/entrypoints/interleave/orchestrator.py` for planner/critic loop + protocols and a minimal retry/refine implementation. + - `fastvideo/entrypoints/interleave/server.py` for `/health`, `/edit`, and + `/generate` alias endpoints. + - `fastvideo/entrypoints/cli/interleave_serve.py` and CLI registration. + - example config/docs under `examples/interleave/`. + - pure tests under `tests/local_tests/` using fakes. +- Implemented the package and CLI registration: + - `fastvideo/entrypoints/interleave/__init__.py` + - `fastvideo/entrypoints/interleave/schema.py` + - `fastvideo/entrypoints/interleave/generator.py` + - `fastvideo/entrypoints/interleave/orchestrator.py` + - `fastvideo/entrypoints/interleave/server.py` + - `fastvideo/entrypoints/cli/interleave_serve.py` + - `fastvideo/entrypoints/cli/main.py` + - `examples/interleave/README.md` + - `examples/interleave/flux2_klein_interleave_serve.yaml` ## Mistakes / Dead Ends @@ -127,6 +156,13 @@ Testing plan: - A first relative-path `apply_patch` attempt was rejected because it would have targeted the original checkout. All manual edits must use absolute paths under `/tmp/fastvideo-interleavethinker`. +- First Modal validation attempt failed before tests. Cause: `--apply-local-patch` + included `.agents/exploration/interleavethinker-fastvideo-integration.md` as a + tracked modification relative to the local-only checkpoint commit + `7e70f859`, but the Modal job checked out upstream `main` + `633d39356804e63478d242611e992dc8e1af3caa`, where that file does not exist. + Fix: rerun Modal with `--patch-paths` limited to code, examples, and tests, + excluding `.agents`. ## Current Handoff Notes @@ -142,16 +178,64 @@ git log --oneline -5 Before editing code, continue reading: -- `fastvideo/entrypoints/cli/main.py` -- `fastvideo/entrypoints/cli/serve.py` -- `fastvideo/entrypoints/openai/image_api.py` -- `fastvideo/entrypoints/openai/video_api.py` -- `fastvideo/entrypoints/streaming/prompt/enhancer.py` -- `fastvideo/entrypoints/video_generator.py` -- `fastvideo/api/schema.py` +- Already inspected: + - `fastvideo/entrypoints/cli/main.py` + - `fastvideo/entrypoints/cli/serve.py` + - `fastvideo/entrypoints/cli/generate.py` + - `fastvideo/entrypoints/cli/inference_config.py` + - `fastvideo/entrypoints/openai/image_api.py` + - `fastvideo/entrypoints/openai/video_api.py` + - `fastvideo/entrypoints/openai/api_server.py` + - `fastvideo/entrypoints/openai/protocol.py` + - `fastvideo/entrypoints/streaming/prompt/enhancer.py` + - `fastvideo/entrypoints/streaming/server.py` + - `fastvideo/entrypoints/video_generator.py` + - `fastvideo/api/schema.py` + - `fastvideo/api/compat.py` Do not use or modify the dirty `/home/toolbox/FastVideo` checkout. +Next immediate steps: + +1. Add pure tests: + - request schema accepts `num_inference_step` and `num_inference_steps`. + - backend builds a `GenerationRequest` with one image frame and optional + decoded input image. + - server returns InterleaveThinker-compatible JSON with a fake backend. + - orchestrator retries with critic-provided `refine_prompt`. +2. Run formatting/pre-commit on changed files if possible. +3. Run focused tests on Modal using + `fastvideo/tests/modal/launch_l40s_job.py` from this worktree, not from the + original checkout. + - Important: pass `--git-repo https://github.com/hao-ai-lab/FastVideo.git` + and `--git-commit 633d39356804e63478d242611e992dc8e1af3caa`. + - Important: pass `--patch-paths` excluding `.agents/...` until the next + local commit is pushed anywhere reachable by Modal. +4. Commit the first code checkpoint after tests or at least after static import + checks. + +Validation completed: + +- First Modal attempt: + - Command: `pytest tests/local_tests/test_interleave_entrypoint.py -q && + pre-commit run --files ...` + - Failed before tests because `.agents` handoff diff could not apply to + upstream `main`. +- Second Modal attempt: + - Command: `pytest tests/local_tests/test_interleave_entrypoint.py -q && + pre-commit run --files fastvideo/entrypoints/cli/main.py + fastvideo/entrypoints/cli/interleave_serve.py + fastvideo/entrypoints/interleave/__init__.py + fastvideo/entrypoints/interleave/schema.py + fastvideo/entrypoints/interleave/generator.py + fastvideo/entrypoints/interleave/orchestrator.py + fastvideo/entrypoints/interleave/server.py` + - Modal app URL shown by CLI: + `https://modal.com/apps/hao-ai-lab/main/ap-kB8zs3gMtKj66loYptASBw` + - Result: `4 passed, 14 warnings in 22.54s`. + - Pre-commit result: `yapf`, `ruff`, `codespell`, `mypy`, filename check all + passed. PyMarkdown/actionlint skipped with no files to check. + ## Proposed Standardization If this integration lands cleanly, promote this exploration into a workflow or diff --git a/examples/interleave/README.md b/examples/interleave/README.md new file mode 100644 index 000000000..1f6600927 --- /dev/null +++ b/examples/interleave/README.md @@ -0,0 +1,43 @@ +# FastVideo Interleave Service + +This example starts a FastVideo-backed image service compatible with +InterleaveThinker-style generator calls. + +InterleaveThinker reward code posts JSON to `/edit`: + +```json +{ + "image": "", + "prompt": "make the sky sunset orange", + "num_inference_step": 4, + "guidance_scale": 1.0, + "width": 1024, + "height": 1024 +} +``` + +FastVideo returns: + +```json +{ + "success": true, + "edited_image": "", + "file_path": "/abs/path/to/output.png" +} +``` + +Run: + +```bash +FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \ + fastvideo interleave-serve --config examples/interleave/flux2_klein_interleave_serve.yaml +``` + +Then point InterleaveThinker at: + +```bash +export EDIT_API_ENDPOINT=http://:8011 +export EDIT_MODEL_NAME=klein +``` + +The same app also exposes `/generate` and `/v1/interleave/edit` aliases. diff --git a/examples/interleave/flux2_klein_interleave_serve.yaml b/examples/interleave/flux2_klein_interleave_serve.yaml new file mode 100644 index 000000000..f2e8ff1fe --- /dev/null +++ b/examples/interleave/flux2_klein_interleave_serve.yaml @@ -0,0 +1,29 @@ +generator: + model_path: black-forest-labs/FLUX.2-klein-4B + engine: + num_gpus: 1 + use_fsdp_inference: false + offload: + dit: false + vae: true + text_encoder: true + pin_cpu_memory: false + pipeline: + workload_type: t2i + +server: + host: 0.0.0.0 + port: 8011 + output_dir: outputs/interleave_flux2_klein + +default_request: + sampling: + num_frames: 1 + fps: 1 + height: 1024 + width: 1024 + num_inference_steps: 4 + guidance_scale: 1.0 + output: + save_video: true + return_frames: false diff --git a/fastvideo/entrypoints/cli/interleave_serve.py b/fastvideo/entrypoints/cli/interleave_serve.py new file mode 100644 index 000000000..c773656a3 --- /dev/null +++ b/fastvideo/entrypoints/cli/interleave_serve.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CLI entrypoint for the InterleaveThinker-compatible image service.""" + +from __future__ import annotations + +import argparse +import os +from typing import cast + +from fastvideo.api.compat import generator_config_to_fastvideo_args +from fastvideo.entrypoints.cli.cli_types import CLISubcommand +from fastvideo.entrypoints.cli.inference_config import build_serve_config +from fastvideo.entrypoints.interleave.server import run_server +from fastvideo.logger import init_logger +from fastvideo.utils import FlexibleArgumentParser + +logger = init_logger(__name__) +_VALIDATED_INTERLEAVE_SERVE_CONFIG_ATTR = "_fastvideo_validated_interleave_serve_config" + + +class InterleaveServeSubcommand(CLISubcommand): + """Starts an InterleaveThinker-compatible FastVideo image service.""" + + def __init__(self) -> None: + self.name = "interleave-serve" + super().__init__() + + def cmd(self, args: argparse.Namespace) -> None: + serve_config = getattr(args, _VALIDATED_INTERLEAVE_SERVE_CONFIG_ATTR, None) + if serve_config is None: + serve_config = build_serve_config( + args, + overrides=getattr(args, "_unknown", None), + ) + + logger.info("CLI interleave serve config: %s", serve_config) + fastvideo_args = generator_config_to_fastvideo_args(serve_config.generator) + run_server( + fastvideo_args, + host=serve_config.server.host, + port=serve_config.server.port, + output_dir=serve_config.server.output_dir, + default_request=serve_config.default_request, + ) + + def validate(self, args: argparse.Namespace) -> None: + if not args.config: + raise ValueError("fastvideo interleave-serve requires --config PATH; " + "use a nested serve config plus optional dotted overrides") + if not os.path.exists(args.config): + raise ValueError(f"Config file not found: {args.config}") + setattr( + args, + _VALIDATED_INTERLEAVE_SERVE_CONFIG_ATTR, + build_serve_config( + args, + overrides=getattr(args, "_unknown", None), + ), + ) + + def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser: + parser = subparsers.add_parser( + "interleave-serve", + help="Start an InterleaveThinker-compatible image service", + usage="fastvideo interleave-serve --config SERVE_CONFIG [--dotted.override VALUE]", + ) + parser.add_argument( + "--config", + type=str, + default="", + required=False, + help="Path to a nested serve config JSON or YAML file. Required.", + ) + return cast(FlexibleArgumentParser, parser) + + +def cmd_init() -> list[CLISubcommand]: + return [InterleaveServeSubcommand()] diff --git a/fastvideo/entrypoints/cli/main.py b/fastvideo/entrypoints/cli/main.py index 0f2b74b8a..c50fd142c 100644 --- a/fastvideo/entrypoints/cli/main.py +++ b/fastvideo/entrypoints/cli/main.py @@ -2,6 +2,8 @@ # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/main.py from fastvideo.entrypoints.cli.cli_types import CLISubcommand from fastvideo.entrypoints.cli.generate import cmd_init as generate_cmd_init +from fastvideo.entrypoints.cli.interleave_serve import ( + cmd_init as interleave_serve_cmd_init, ) from fastvideo.utils import FlexibleArgumentParser from fastvideo.entrypoints.cli.router_serve import ( cmd_init as router_serve_cmd_init, ) @@ -15,6 +17,7 @@ def cmd_init() -> list[CLISubcommand]: commands = [] commands.extend(generate_cmd_init()) commands.extend(serve_cmd_init()) + commands.extend(interleave_serve_cmd_init()) commands.extend(router_serve_cmd_init()) commands.extend(bench_cmd_init()) commands.extend(eval_cmd_init()) @@ -33,7 +36,7 @@ def main() -> None: cmds[cmd.name] = cmd args, unknown = parser.parse_known_args() - if unknown and args.subparser not in {"generate", "serve"}: + if unknown and args.subparser not in {"generate", "serve", "interleave-serve"}: parser.error(f"unrecognized arguments: {' '.join(unknown)}") args._unknown = unknown if args.subparser in cmds: diff --git a/fastvideo/entrypoints/interleave/__init__.py b/fastvideo/entrypoints/interleave/__init__.py new file mode 100644 index 000000000..050a30780 --- /dev/null +++ b/fastvideo/entrypoints/interleave/__init__.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +"""InterleaveThinker-compatible FastVideo app layer.""" + +from fastvideo.entrypoints.interleave.generator import ( + FastVideoImageGeneratorBackend, + ImageGeneratorBackend, +) +from fastvideo.entrypoints.interleave.orchestrator import ( + AcceptAllCritic, + CriticProvider, + InterleaveOrchestrator, + PlannerProvider, + SinglePromptPlanner, +) +from fastvideo.entrypoints.interleave.schema import ( + CriticDecision, + CriticInput, + GeneratedImage, + InterleaveAttempt, + InterleaveEditRequest, + InterleaveEditResponse, + InterleaveTrace, + PlannedInterleaveStep, + PlannerInput, +) +from fastvideo.entrypoints.interleave.server import build_app + +__all__ = [ + "AcceptAllCritic", + "CriticDecision", + "CriticInput", + "CriticProvider", + "FastVideoImageGeneratorBackend", + "GeneratedImage", + "ImageGeneratorBackend", + "InterleaveAttempt", + "InterleaveEditRequest", + "InterleaveEditResponse", + "InterleaveOrchestrator", + "InterleaveTrace", + "PlannedInterleaveStep", + "PlannerInput", + "PlannerProvider", + "SinglePromptPlanner", + "build_app", +] diff --git a/fastvideo/entrypoints/interleave/generator.py b/fastvideo/entrypoints/interleave/generator.py new file mode 100644 index 000000000..07c21c800 --- /dev/null +++ b/fastvideo/entrypoints/interleave/generator.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastVideo generator adapter for InterleaveThinker-style image calls.""" + +from __future__ import annotations + +import base64 +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.entrypoints.interleave.schema import ( + GeneratedImage, + InterleaveEditRequest, +) + + +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, + ) + + +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) diff --git a/fastvideo/entrypoints/interleave/orchestrator.py b/fastvideo/entrypoints/interleave/orchestrator.py new file mode 100644 index 000000000..76822e56f --- /dev/null +++ b/fastvideo/entrypoints/interleave/orchestrator.py @@ -0,0 +1,182 @@ +# 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.entrypoints.interleave.generator import ( + ImageGeneratorBackend, + encode_file_to_base64, +) +from fastvideo.entrypoints.interleave.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 plan(self, request: PlannerInput) -> Sequence[PlannedInterleaveStep]: + return [ + PlannedInterleaveStep( + prompt=request.instruction, + input_image_path=request.initial_image_path, + max_attempts=1, + ) + ] + + +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, + ) + 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, + ) -> InterleaveEditRequest: + 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=self.seed, + num_inference_steps=self.num_inference_steps, + guidance_scale=self.guidance_scale, + ) diff --git a/fastvideo/entrypoints/interleave/schema.py b/fastvideo/entrypoints/interleave/schema.py new file mode 100644 index 000000000..cf8315728 --- /dev/null +++ b/fastvideo/entrypoints/interleave/schema.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Schemas for InterleaveThinker-style orchestration and compatibility APIs.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field + + +class InterleaveEditRequest(BaseModel): + """JSON body accepted by the InterleaveThinker-compatible ``/edit`` route. + + 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 the endpoint can sit behind + several generator clients without rejecting model-specific knobs it does not + use yet. + """ + + 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/entrypoints/interleave/server.py b/fastvideo/entrypoints/interleave/server.py new file mode 100644 index 000000000..95ae75992 --- /dev/null +++ b/fastvideo/entrypoints/interleave/server.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +"""InterleaveThinker-compatible FastVideo image generation service.""" + +from __future__ import annotations + +import asyncio +import uuid +from contextlib import asynccontextmanager +from collections.abc import AsyncIterator + +import uvicorn +from fastapi import FastAPI +from fastapi.middleware.cors import CORSMiddleware + +from fastvideo.api.schema import GenerationRequest +from fastvideo.entrypoints.interleave.generator import ( + FastVideoImageGeneratorBackend, + ImageGeneratorBackend, +) +from fastvideo.entrypoints.interleave.schema import ( + InterleaveEditRequest, + InterleaveEditResponse, +) +from fastvideo.entrypoints.video_generator import VideoGenerator +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +DEFAULT_HOST = "0.0.0.0" +DEFAULT_PORT = 8011 +DEFAULT_OUTPUT_DIR = "outputs/interleave" + + +@asynccontextmanager +async def lifespan(app: FastAPI) -> AsyncIterator[None]: + generator = None + if app.state.backend is None: + args: FastVideoArgs = app.state.fastvideo_args + logger.info("Loading Interleave FastVideo backend from %s ...", args.model_path) + generator = VideoGenerator.from_fastvideo_args(args) + app.state.backend = FastVideoImageGeneratorBackend( + generator, + output_dir=app.state.output_dir, + default_request=app.state.default_request, + ) + logger.info("Interleave FastVideo backend loaded.") + + yield + + if generator is not None: + logger.info("Shutting down Interleave FastVideo backend ...") + generator.shutdown() + + +def build_app( + *, + fastvideo_args: FastVideoArgs | None = None, + output_dir: str = DEFAULT_OUTPUT_DIR, + default_request: GenerationRequest | None = None, + backend: ImageGeneratorBackend | None = None, +) -> FastAPI: + if backend is None and fastvideo_args is None: + raise ValueError("build_app requires either `backend` or `fastvideo_args`") + + app = FastAPI( + title="FastVideo Interleave Compatibility API", + version="0.1.0", + lifespan=lifespan, + ) + app.state.fastvideo_args = fastvideo_args + app.state.output_dir = output_dir + app.state.default_request = default_request + app.state.backend = backend + + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + + @app.get("/health") + async def health(): + return {"status": "ok"} + + @app.post("/edit", response_model=InterleaveEditResponse) + async def edit(request: InterleaveEditRequest) -> InterleaveEditResponse: + return await _run_edit(app, request) + + @app.post("/generate", response_model=InterleaveEditResponse) + async def generate(request: InterleaveEditRequest) -> InterleaveEditResponse: + return await _run_edit(app, request) + + @app.post("/v1/interleave/edit", response_model=InterleaveEditResponse) + async def interleave_edit(request: InterleaveEditRequest) -> InterleaveEditResponse: + return await _run_edit(app, request) + + return app + + +async def _run_edit( + app: FastAPI, + request: InterleaveEditRequest, +) -> InterleaveEditResponse: + request_id = uuid.uuid4().hex + loop = asyncio.get_running_loop() + backend: ImageGeneratorBackend = app.state.backend + try: + generated = await loop.run_in_executor(None, lambda: backend.generate(request, request_id=request_id)) + except Exception as exc: # noqa: BLE001 + logger.exception("Interleave edit request %s failed", request_id) + return InterleaveEditResponse( + success=False, + prompt=request.prompt, + error=str(exc), + metadata={"request_id": request_id}, + ) + + metadata = dict(generated.metadata) + metadata.setdefault("request_id", request_id) + return InterleaveEditResponse( + success=True, + edited_image=generated.image_base64, + file_path=generated.file_path, + prompt=generated.prompt, + inference_time_s=generated.inference_time_s, + metadata=metadata, + ) + + +def run_server( + fastvideo_args: FastVideoArgs, + *, + host: str = DEFAULT_HOST, + port: int = DEFAULT_PORT, + output_dir: str = DEFAULT_OUTPUT_DIR, + default_request: GenerationRequest | None = None, +) -> None: + app = build_app( + fastvideo_args=fastvideo_args, + output_dir=output_dir, + default_request=default_request, + ) + logger.info("Starting FastVideo Interleave server on %s:%d", host, port) + logger.info("Model: %s", fastvideo_args.model_path) + uvicorn.run( + app, + host=host, + port=port, + log_level="info", + timeout_keep_alive=300, + ) diff --git a/tests/local_tests/test_interleave_entrypoint.py b/tests/local_tests/test_interleave_entrypoint.py new file mode 100644 index 000000000..8d53e748a --- /dev/null +++ b/tests/local_tests/test_interleave_entrypoint.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import base64 +from pathlib import Path + +from starlette.testclient import TestClient + +from fastvideo.api.compat import ( + explicit_request_updates, + legacy_generate_call_to_request, +) +from fastvideo.api.results import GenerationResult +from fastvideo.entrypoints.interleave.generator import ( + FastVideoImageGeneratorBackend, +) +from fastvideo.entrypoints.interleave.orchestrator import InterleaveOrchestrator +from fastvideo.entrypoints.interleave.schema import ( + CriticDecision, + GeneratedImage, + InterleaveEditRequest, + PlannedInterleaveStep, +) +from fastvideo.entrypoints.interleave.server import build_app + + +def test_interleave_edit_request_accepts_singular_step_field() -> None: + request = InterleaveEditRequest( + prompt="a ceramic cup on a table", + num_inference_step=4, + ) + + assert request.resolved_num_inference_steps() == 4 + + plural_request = InterleaveEditRequest( + prompt="a ceramic cup on a table", + num_inference_step=4, + num_inference_steps=8, + ) + assert plural_request.resolved_num_inference_steps() == 8 + + +def test_fastvideo_backend_translates_edit_request(tmp_path: Path) -> None: + class FakeGenerator: + def __init__(self) -> None: + self.requests = [] + + def generate(self, request): + self.requests.append(request) + updates = explicit_request_updates(request) + output_path = Path(updates["output_path"]) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_bytes(b"fake-png") + return GenerationResult( + prompt=request.prompt, + video_path=str(output_path), + generation_time=0.25, + ) + + default_request = legacy_generate_call_to_request( + "unused", + None, + legacy_kwargs={ + "height": 512, + "width": 768, + "num_inference_steps": 4, + }, + ) + fake = FakeGenerator() + backend = FastVideoImageGeneratorBackend( + fake, + output_dir=str(tmp_path), + default_request=default_request, + ) + input_b64 = base64.b64encode(b"input-image").decode("utf-8") + + generated = backend.generate( + InterleaveEditRequest( + prompt="turn it into a watercolor", + image=input_b64, + width=1024, + seed=7, + ), + request_id="abc123", + ) + + assert generated.prompt == "turn it into a watercolor" + assert generated.image_base64 == base64.b64encode(b"fake-png").decode("utf-8") + assert generated.file_path is not None + assert generated.file_path.endswith("abc123.png") + + updates = explicit_request_updates(fake.requests[0]) + assert updates["num_frames"] == 1 + assert updates["fps"] == 1 + assert updates["height"] == 512 + assert updates["width"] == 1024 + assert updates["num_inference_steps"] == 4 + assert updates["seed"] == 7 + assert updates["save_video"] is True + assert updates["return_frames"] is False + assert Path(updates["image_path"]).read_bytes() == b"input-image" + + +def test_interleave_server_returns_compatible_edit_response() -> None: + class FakeBackend: + def __init__(self) -> None: + self.requests = [] + + def generate(self, request, *, request_id=None): + self.requests.append((request, request_id)) + return GeneratedImage( + prompt=request.prompt, + image_base64="ZmFrZQ==", + file_path="/tmp/fake.png", + inference_time_s=0.1, + metadata={"backend": "fake"}, + ) + + backend = FakeBackend() + app = build_app(backend=backend) + + with TestClient(app) as client: + response = client.post( + "/edit", + json={ + "prompt": "add a red umbrella", + "num_inference_step": 4, + }, + ) + + assert response.status_code == 200 + payload = response.json() + assert payload["success"] is True + assert payload["edited_image"] == "ZmFrZQ==" + assert payload["file_path"] == "/tmp/fake.png" + assert payload["metadata"]["backend"] == "fake" + assert backend.requests[0][0].prompt == "add a red umbrella" + assert backend.requests[0][1] + + +def test_interleave_orchestrator_retries_with_refined_prompt() -> None: + class FakePlanner: + def plan(self, request): + return [ + PlannedInterleaveStep( + prompt=request.instruction, + max_attempts=2, + ) + ] + + class FakeGenerator: + def __init__(self) -> None: + self.prompts = [] + + def generate(self, request, *, request_id=None): + del request_id + self.prompts.append(request.prompt) + return GeneratedImage( + prompt=request.prompt, + image_base64=base64.b64encode(request.prompt.encode("utf-8")).decode("utf-8"), + file_path=f"/tmp/{len(self.prompts)}.png", + ) + + class RefiningCritic: + def __init__(self) -> None: + self.calls = 0 + + def review(self, request): + self.calls += 1 + if self.calls == 1: + return CriticDecision( + success=False, + refine_prompt="refined prompt", + reason="first attempt missed the instruction", + ) + return CriticDecision(success=True) + + generator = FakeGenerator() + orchestrator = InterleaveOrchestrator( + planner=FakePlanner(), + generator=generator, + critic=RefiningCritic(), + ) + + trace = orchestrator.run("initial prompt") + + assert trace.success is True + assert generator.prompts == ["initial prompt", "refined prompt"] + assert len(trace.attempts) == 2 + assert trace.final_image is not None + assert trace.final_image.prompt == "refined prompt"