diff --git a/.agents/exploration/interleavethinker-fastvideo-integration.md b/.agents/exploration/interleavethinker-fastvideo-integration.md index 9efedf9c1..845d44a81 100644 --- a/.agents/exploration/interleavethinker-fastvideo-integration.md +++ b/.agents/exploration/interleavethinker-fastvideo-integration.md @@ -1841,3 +1841,128 @@ Conclusion: FastVideo image-generation model/config, then decide whether the planner and critic should stay resident together or be unloaded/reloaded around the generator for tighter GPU memory budgets. + +## Stage 5 Execution: Native Interleave Run CLI + +Status: validated, pending commit/push. + +Scope for this implementation slice: + +- Add a first-class `fastvideo interleave-run` command that runs the complete + native orchestration path from a config file. +- Reuse existing pieces rather than adding another execution framework: + - `InterleaveOrchestrator`; + - `FastVideoImageGeneratorBackend`; + - `NanoBananaImageGeneratorBackend`; + - `InterleaveThinkerPlannerProvider`; + - `InterleaveThinkerCriticProvider`; + - `SinglePromptPlanner` and `AcceptAllCritic` for lightweight smoke paths. +- Add a typed interleave run config with blocks for: + - `generator`: regular FastVideo `GeneratorConfig` for local generation; + - `request`: regular `GenerationRequest` defaults for per-step generation; + - `image_backend`: FastVideo or Nano Banana backend selection; + - `planner`: single-prompt or InterleaveThinker planner backend; + - `critic`: none, accept-all, or InterleaveThinker critic backend; + - `interleave`: instruction, optional initial image, output directory, trace + path, and trace image payload policy. +- Keep the command testable without loading GPU models by allowing runner tests + to inject a fake image backend. + +Expected files: + +- `fastvideo/entrypoints/interleave/config.py` +- `fastvideo/entrypoints/interleave/runner.py` +- `fastvideo/entrypoints/cli/interleave_run.py` +- `fastvideo/entrypoints/cli/main.py` +- `fastvideo/entrypoints/interleave/orchestrator.py` +- `fastvideo/entrypoints/interleave/__init__.py` +- `examples/interleave/interleave_run.yaml` +- `examples/interleave/README.md` +- `tests/local_tests/test_interleave_run_cli.py` +- this handoff file + +Validation plan: + +- Local syntax and `git diff --check`. +- Modal pytest for interleave entrypoint/provider/run CLI tests. +- Modal pre-commit on changed non-excluded files. +- Modal import/config smoke for `fastvideo interleave-run` using injected/fake + components in tests. +- A later GPU smoke should run real planner + real critic + real FastVideo + generator once a small practical model/config is selected and GPU memory + residency is measured. + +Implemented changes: + +- Added `fastvideo/entrypoints/interleave/config.py` with typed config blocks: + - `InterleaveRunConfig`; + - `InterleaveRunStateConfig`; + - `InterleaveImageBackendConfig`; + - `InterleavePlannerConfig`; + - `InterleaveCriticConfig`. +- Added `load_interleave_run_config(...)` with: + - FastVideo-style YAML/JSON loading; + - dotted CLI overrides for `interleave.`, `image_backend.`, `planner.`, + `critic.`, `request.`, and `generator.`; + - CLI convenience fields for prompt, input image, output directory, and trace + path; + - explicit-path binding for the nested `GenerationRequest` so per-step + generation defaults behave like existing `generate`/`serve` configs. +- Added `fastvideo/entrypoints/interleave/runner.py` with: + - `run_interleave_config(...)`; + - `build_planner(...)`; + - `build_critic(...)`; + - `build_image_backend(...)`; + - trace path resolution and cleanup of local `VideoGenerator` instances. +- Added `fastvideo interleave-run` in + `fastvideo/entrypoints/cli/interleave_run.py` and registered it in + `fastvideo/entrypoints/cli/main.py`. +- Updated `SinglePromptPlanner` to accept a configurable retry count. +- Exported the new config/runner API from `fastvideo.entrypoints.interleave`. +- Added `examples/interleave/interleave_run.yaml` and updated the README with + native CLI usage plus the real InterleaveThinker planner/critic config block. +- Added `tests/local_tests/test_interleave_run_cli.py` covering: + - config loading from YAML; + - CLI field overrides and dotted override normalization; + - rejection of unsupported override prefixes; + - full runner execution with an injected fake image backend; + - configurable fallback planner attempt count. + +Validation completed: + +- Local lightweight checks: + - `python -m py_compile` passed for the new/changed Python files. + - `git diff --check` passed. +- Modal focused pytest: + - App URL: `https://modal.com/apps/hao-ai-lab/main/ap-XyjjWxk63VMMAgyefDV9iz` + - Command: + `pytest tests/local_tests/test_interleave_entrypoint.py + tests/local_tests/test_interleave_model_providers.py + tests/local_tests/test_interleave_run_cli.py -q` + - Result: `14 passed, 14 warnings in 15.93s`. +- Modal pre-commit: + - App URL: `https://modal.com/apps/hao-ai-lab/main/ap-4lOvHURSxKqFmit8TqnhBe` + - Command: + `pre-commit run --files examples/interleave/README.md + examples/interleave/interleave_run.yaml + fastvideo/entrypoints/cli/interleave_run.py + fastvideo/entrypoints/cli/main.py + fastvideo/entrypoints/interleave/__init__.py + fastvideo/entrypoints/interleave/config.py + fastvideo/entrypoints/interleave/orchestrator.py + fastvideo/entrypoints/interleave/runner.py + tests/local_tests/test_interleave_run_cli.py` + - Result: yapf, ruff, codespell, mypy, filename check, and suggestion hooks + passed; PyMarkdown/actionlint skipped with no files to check. + +Conclusion: + +- FastVideo now has a native config-driven `interleave-run` entrypoint that can + wire planner, generator, and critic backends without hand-written scripts. +- The default example provides a lightweight FastVideo image-backend smoke path; + the same config can switch to the real InterleaveThinker planner and critic + checkpoints. +- Next step: run a GPU smoke with `fastvideo interleave-run` and a real + FastVideo image backend. Use fallback planner/accept-all critic first to + validate the CLI + generator path, then test real planner/critic residency or + a split-process `/edit` generator to avoid loading all models in one process. diff --git a/examples/interleave/README.md b/examples/interleave/README.md index 0e83878a8..b6aa178f2 100644 --- a/examples/interleave/README.md +++ b/examples/interleave/README.md @@ -4,6 +4,8 @@ This directory contains two InterleaveThinker-oriented entrypoints: - `flux2_klein_interleave_serve.yaml`: run FastVideo as an InterleaveThinker-compatible generator service. +- `interleave_run.yaml`: run the native planner -> generator -> critic loop + through `fastvideo interleave-run`. - `interleave_single_prompt.py`: run the native FastVideo interleave orchestrator with a fallback single-prompt planner and accept-all critic. @@ -51,6 +53,57 @@ export EDIT_MODEL_NAME=klein The same app also exposes `/generate` and `/v1/interleave/edit` aliases. +## Native Interleave Run CLI + +Run the native orchestrator from a config file: + +```bash +FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \ + fastvideo interleave-run --config examples/interleave/interleave_run.yaml +``` + +Override common per-run fields from the CLI: + +```bash +fastvideo interleave-run \ + --config examples/interleave/interleave_run.yaml \ + --prompt "create a clean product photo of a red ceramic mug" \ + --output-dir outputs/interleave_mug \ + --trace-path outputs/interleave_mug/trace.json +``` + +The example defaults to `planner.kind: single_prompt` and +`critic.kind: accept_all` so it can smoke-test the FastVideo image backend +without loading planner/critic checkpoints. To run the real InterleaveThinker +loop, switch those blocks to `interleave_thinker` and set: + +```yaml +planner: + kind: interleave_thinker + init_from: InterleaveThinker/InterleaveThinker-Planner-8B + processor_from: Qwen/Qwen3-VL-8B-Instruct + trainable: false + torch_dtype: bf16 + device_map: cuda:0 + attn_implementation: sdpa + max_new_tokens: 2048 + +critic: + kind: interleave_thinker + init_from: InterleaveThinker/Critic-SFT-8B + processor_from: Qwen/Qwen3-VL-8B-Instruct + trainable: false + torch_dtype: bf16 + device_map: cuda:0 + attn_implementation: sdpa + max_new_tokens: 512 +``` + +Loading the planner, critic, and image generator in one process can be memory +intensive. On tighter GPUs, run the `/edit` compatibility service separately or +load only the planner/critic in the native runner while the generator is served +out-of-process. + ## Nano Banana / Gemini API Backends The InterleaveThinker RL config can use closed-source Google models through diff --git a/examples/interleave/interleave_run.yaml b/examples/interleave/interleave_run.yaml new file mode 100644 index 000000000..317a66a20 --- /dev/null +++ b/examples/interleave/interleave_run.yaml @@ -0,0 +1,65 @@ +# Native FastVideo interleave-run example. +# +# This default runs a lightweight one-step trace with the fallback planner and +# accept-all critic. Switch planner.kind / critic.kind to interleave_thinker and +# set the model fields below to run the real InterleaveThinker planner/critic. + +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 + +image_backend: + kind: fastvideo + +planner: + kind: single_prompt + max_attempts_per_step: 1 + # Real planner example: + # kind: interleave_thinker + # init_from: InterleaveThinker/InterleaveThinker-Planner-8B + # processor_from: Qwen/Qwen3-VL-8B-Instruct + # trainable: false + # torch_dtype: bf16 + # device_map: cuda:0 + # attn_implementation: sdpa + # max_new_tokens: 2048 + +critic: + kind: accept_all + # Real critic example: + # kind: interleave_thinker + # init_from: InterleaveThinker/Critic-SFT-8B + # processor_from: Qwen/Qwen3-VL-8B-Instruct + # trainable: false + # torch_dtype: bf16 + # device_map: cuda:0 + # attn_implementation: sdpa + # max_new_tokens: 512 + +interleave: + output_dir: outputs/interleave_run + trace_path: outputs/interleave_run/trace.json + include_images_in_trace: false + +request: + prompt: a brushed steel espresso machine on a marble counter, morning window light + sampling: + num_frames: 1 + fps: 1 + height: 1024 + width: 1024 + num_inference_steps: 4 + guidance_scale: 1.0 + seed: 0 + output: + save_video: true + return_frames: false diff --git a/fastvideo/entrypoints/cli/interleave_run.py b/fastvideo/entrypoints/cli/interleave_run.py new file mode 100644 index 000000000..8597e1890 --- /dev/null +++ b/fastvideo/entrypoints/cli/interleave_run.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CLI entrypoint for native InterleaveThinker orchestration.""" + +from __future__ import annotations + +import argparse +import os +from typing import cast + +from fastvideo.entrypoints.cli.cli_types import CLISubcommand +from fastvideo.entrypoints.interleave.config import load_interleave_run_config +from fastvideo.entrypoints.interleave.runner import run_interleave_config +from fastvideo.logger import init_logger +from fastvideo.utils import FlexibleArgumentParser + +logger = init_logger(__name__) +_VALIDATED_INTERLEAVE_RUN_CONFIG_ATTR = "_fastvideo_validated_interleave_run_config" + + +class InterleaveRunSubcommand(CLISubcommand): + """Runs a native planner -> generator -> critic interleave trace.""" + + def __init__(self) -> None: + self.name = "interleave-run" + super().__init__() + + def cmd(self, args: argparse.Namespace) -> None: + run_config = getattr(args, _VALIDATED_INTERLEAVE_RUN_CONFIG_ATTR, None) + if run_config is None: + run_config = load_interleave_run_config( + args.config, + overrides=getattr(args, "_unknown", None), + prompt=args.prompt, + input_image=args.input_image, + output_dir=args.output_dir, + trace_path=args.trace_path, + ) + logger.info("CLI interleave run config: %s", run_config) + + result = run_interleave_config(run_config) + final_path = result.trace.final_image.file_path if result.trace.final_image is not None else None + if final_path: + print(f"Image: {final_path}") + print(f"Trace: {result.trace_path}") + print(f"Success: {result.trace.success}") + + def validate(self, args: argparse.Namespace) -> None: + if not args.config: + raise ValueError("fastvideo interleave-run requires --config PATH; " + "use a nested interleave 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_RUN_CONFIG_ATTR, + load_interleave_run_config( + args.config, + overrides=getattr(args, "_unknown", None), + prompt=args.prompt, + input_image=args.input_image, + output_dir=args.output_dir, + trace_path=args.trace_path, + ), + ) + + def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser: + parser = subparsers.add_parser( + "interleave-run", + help="Run native InterleaveThinker planner/generator/critic orchestration", + usage="fastvideo interleave-run --config CONFIG [--prompt TEXT] [--dotted.override VALUE]", + ) + parser.add_argument( + "--config", + type=str, + default="", + required=False, + help="Path to a nested interleave run config JSON or YAML file. Required.", + ) + parser.add_argument( + "--prompt", + type=str, + default=None, + help="Override interleave.instruction for this run.", + ) + parser.add_argument( + "--input-image", + type=str, + default=None, + help="Override interleave.initial_image_path for image-conditioned planning/editing.", + ) + parser.add_argument( + "--output-dir", + type=str, + default=None, + help="Override interleave.output_dir.", + ) + parser.add_argument( + "--trace-path", + type=str, + default=None, + help="Override interleave.trace_path.", + ) + return cast(FlexibleArgumentParser, parser) + + +def cmd_init() -> list[CLISubcommand]: + return [InterleaveRunSubcommand()] diff --git a/fastvideo/entrypoints/cli/main.py b/fastvideo/entrypoints/cli/main.py index c50fd142c..64996a56b 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_run import ( + cmd_init as interleave_run_cmd_init, ) from fastvideo.entrypoints.cli.interleave_serve import ( cmd_init as interleave_serve_cmd_init, ) from fastvideo.utils import FlexibleArgumentParser @@ -17,6 +19,7 @@ def cmd_init() -> list[CLISubcommand]: commands = [] commands.extend(generate_cmd_init()) commands.extend(serve_cmd_init()) + commands.extend(interleave_run_cmd_init()) commands.extend(interleave_serve_cmd_init()) commands.extend(router_serve_cmd_init()) commands.extend(bench_cmd_init()) @@ -36,7 +39,7 @@ def main() -> None: cmds[cmd.name] = cmd args, unknown = parser.parse_known_args() - if unknown and args.subparser not in {"generate", "serve", "interleave-serve"}: + if unknown and args.subparser not in {"generate", "serve", "interleave-run", "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 index 90a65471a..dfcbe08f2 100644 --- a/fastvideo/entrypoints/interleave/__init__.py +++ b/fastvideo/entrypoints/interleave/__init__.py @@ -1,6 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 """InterleaveThinker-compatible FastVideo app layer.""" +from fastvideo.entrypoints.interleave.config import ( + InterleaveCriticConfig, + InterleaveImageBackendConfig, + InterleavePlannerConfig, + InterleaveRunConfig, + InterleaveRunStateConfig, + load_interleave_run_config, + resolve_interleave_instruction, +) from fastvideo.entrypoints.interleave.generator import ( FastVideoImageGeneratorBackend, ImageGeneratorBackend, @@ -16,6 +25,10 @@ from fastvideo.entrypoints.interleave.providers import ( InterleaveThinkerCriticProvider, InterleaveThinkerPlannerProvider, ) +from fastvideo.entrypoints.interleave.runner import ( + InterleaveRunResult, + run_interleave_config, +) from fastvideo.entrypoints.interleave.schema import ( CriticDecision, CriticInput, @@ -42,9 +55,15 @@ __all__ = [ "GeneratedImage", "ImageGeneratorBackend", "InterleaveAttempt", + "InterleaveCriticConfig", "InterleaveEditRequest", "InterleaveEditResponse", + "InterleaveImageBackendConfig", "InterleaveOrchestrator", + "InterleavePlannerConfig", + "InterleaveRunConfig", + "InterleaveRunResult", + "InterleaveRunStateConfig", "InterleaveThinkerCriticProvider", "InterleaveThinkerPlannerProvider", "InterleaveTrace", @@ -53,6 +72,9 @@ __all__ = [ "PlannerProvider", "SinglePromptPlanner", "build_app", + "load_interleave_run_config", + "resolve_interleave_instruction", + "run_interleave_config", "save_trace", "trace_to_dict", ] diff --git a/fastvideo/entrypoints/interleave/config.py b/fastvideo/entrypoints/interleave/config.py new file mode 100644 index 000000000..f5149051d --- /dev/null +++ b/fastvideo/entrypoints/interleave/config.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Typed config for native interleaved generation runs.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Literal + +from fastvideo.api.overrides import apply_overrides, parse_cli_overrides +from fastvideo.api.parser import load_raw_config, parse_config +from fastvideo.api.request_metadata import bind_generation_request_raw +from fastvideo.api.schema import GenerationRequest, GeneratorConfig + + +@dataclass +class InterleaveRunStateConfig: + instruction: str | None = None + initial_image_path: str | None = None + output_dir: str = "outputs/interleave_run" + trace_path: str | None = None + include_images_in_trace: bool = False + + +@dataclass +class InterleaveImageBackendConfig: + kind: Literal["fastvideo", "nano_banana"] = "fastvideo" + output_dir: str | None = None + model: str = "gemini-3.1-flash-image" + api_key: str | None = None + base_url: str | None = None + aspect_ratio: str | None = None + image_size: str | None = None + max_attempts: int = 3 + retry_delay_s: float = 2.0 + + +@dataclass +class InterleavePlannerConfig: + kind: Literal["single_prompt", "interleave_thinker"] = "single_prompt" + init_from: str | None = None + processor_from: str | None = None + load_backend: bool = True + trainable: bool = False + image_dir: str = "" + torch_dtype: str = "auto" + device_map: Any | None = None + attn_implementation: str | None = None + trust_remote_code: bool = False + use_cache: bool = False + max_prompt_length: int = 16384 + max_response_length: int = 4096 + lora: dict[str, Any] | None = None + num_generations: int = 1 + temperature: float = 0.0 + top_p: float = 1.0 + max_new_tokens: int | None = None + max_attempts_per_step: int = 2 + + +@dataclass +class InterleaveCriticConfig: + kind: Literal["none", "accept_all", "interleave_thinker"] = "accept_all" + init_from: str | None = None + processor_from: str | None = None + load_backend: bool = True + trainable: bool = False + image_dir: str = "" + torch_dtype: str = "auto" + device_map: Any | None = None + attn_implementation: str | None = None + trust_remote_code: bool = False + use_cache: bool = False + max_prompt_length: int = 16384 + max_response_length: int = 4096 + lora: dict[str, Any] | None = None + num_generations: int = 1 + temperature: float = 0.0 + top_p: float = 1.0 + max_new_tokens: int | None = None + + +@dataclass +class InterleaveRunConfig: + interleave: InterleaveRunStateConfig = field(default_factory=InterleaveRunStateConfig) + image_backend: InterleaveImageBackendConfig = field(default_factory=InterleaveImageBackendConfig) + planner: InterleavePlannerConfig = field(default_factory=InterleavePlannerConfig) + critic: InterleaveCriticConfig = field(default_factory=InterleaveCriticConfig) + request: GenerationRequest = field(default_factory=GenerationRequest) + generator: GeneratorConfig | None = None + + +_INTERLEAVE_RUN_OVERRIDE_PREFIXES = ( + "interleave.", + "image_backend.", + "planner.", + "critic.", + "request.", + "generator.", +) + + +def load_interleave_run_config( + path: str | Path, + *, + overrides: list[str] | None = None, + prompt: str | None = None, + input_image: str | None = None, + output_dir: str | None = None, + trace_path: str | None = None, +) -> InterleaveRunConfig: + raw = load_raw_config(path) + raw = _apply_interleave_cli_fields( + raw, + prompt=prompt, + input_image=input_image, + output_dir=output_dir, + trace_path=trace_path, + ) + raw = _apply_interleave_overrides(raw, overrides) + config = parse_config(InterleaveRunConfig, raw) + bind_generation_request_raw( + config.request, + raw.get("request") if isinstance(raw.get("request"), Mapping) else {}, + ) + validate_interleave_run_config(config) + return config + + +def resolve_interleave_instruction(config: InterleaveRunConfig) -> str: + if config.interleave.instruction: + return config.interleave.instruction + if isinstance(config.request.prompt, str) and config.request.prompt: + return config.request.prompt + if isinstance(config.request.prompt, list) and len(config.request.prompt) == 1: + prompt = config.request.prompt[0] + if isinstance(prompt, str) and prompt: + return prompt + raise ValueError("interleave-run requires interleave.instruction or a single request.prompt") + + +def validate_interleave_run_config(config: InterleaveRunConfig) -> None: + resolve_interleave_instruction(config) + if config.image_backend.kind == "fastvideo" and config.generator is None: + raise ValueError("interleave-run with image_backend.kind=fastvideo requires a generator config") + if config.planner.kind == "interleave_thinker" and config.planner.max_new_tokens is not None: + _require_positive_int(config.planner.max_new_tokens, "planner.max_new_tokens") + _require_positive_int(config.planner.max_attempts_per_step, "planner.max_attempts_per_step") + if config.critic.kind == "interleave_thinker" and config.critic.max_new_tokens is not None: + _require_positive_int(config.critic.max_new_tokens, "critic.max_new_tokens") + + +def _apply_interleave_cli_fields( + raw: Mapping[str, Any], + *, + prompt: str | None, + input_image: str | None, + output_dir: str | None, + trace_path: str | None, +) -> dict[str, Any]: + merged = deepcopy(dict(raw)) + interleave = merged.setdefault("interleave", {}) + if not isinstance(interleave, dict): + raise ValueError("interleave must be a mapping") + if prompt is not None: + interleave["instruction"] = prompt + if input_image is not None: + interleave["initial_image_path"] = input_image + if output_dir is not None: + interleave["output_dir"] = output_dir + if trace_path is not None: + interleave["trace_path"] = trace_path + return merged + + +def _apply_interleave_overrides( + raw: Mapping[str, Any], + overrides: list[str] | None, +) -> dict[str, Any]: + if not overrides: + return deepcopy(dict(raw)) + + parsed = parse_cli_overrides(overrides) + for key in parsed: + if "." not in key: + raise ValueError("CLI overrides must use dotted config paths like --request.sampling.seed 42") + if not key.startswith(_INTERLEAVE_RUN_OVERRIDE_PREFIXES): + allowed = ", ".join(_INTERLEAVE_RUN_OVERRIDE_PREFIXES) + raise ValueError(f"Unsupported override path {key!r}. Allowed prefixes: {allowed}") + return apply_overrides(raw, parsed) + + +def _require_positive_int(value: int, path: str) -> None: + if value <= 0: + raise ValueError(f"{path} must be > 0; got {value}") + + +__all__ = [ + "InterleaveCriticConfig", + "InterleaveImageBackendConfig", + "InterleavePlannerConfig", + "InterleaveRunConfig", + "InterleaveRunStateConfig", + "load_interleave_run_config", + "resolve_interleave_instruction", + "validate_interleave_run_config", +] diff --git a/fastvideo/entrypoints/interleave/orchestrator.py b/fastvideo/entrypoints/interleave/orchestrator.py index 76822e56f..96c2f755f 100644 --- a/fastvideo/entrypoints/interleave/orchestrator.py +++ b/fastvideo/entrypoints/interleave/orchestrator.py @@ -39,12 +39,15 @@ class CriticProvider(Protocol): 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=1, + max_attempts=self.max_attempts, ) ] diff --git a/fastvideo/entrypoints/interleave/runner.py b/fastvideo/entrypoints/interleave/runner.py new file mode 100644 index 000000000..4d380205e --- /dev/null +++ b/fastvideo/entrypoints/interleave/runner.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Config-driven native InterleaveThinker orchestration runner.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from fastvideo.entrypoints.interleave.config import ( + InterleaveCriticConfig, + InterleaveImageBackendConfig, + InterleavePlannerConfig, + InterleaveRunConfig, + resolve_interleave_instruction, +) +from fastvideo.entrypoints.interleave.generator import ( + FastVideoImageGeneratorBackend, + ImageGeneratorBackend, + NanoBananaImageGeneratorBackend, +) +from fastvideo.entrypoints.interleave.orchestrator import ( + AcceptAllCritic, + CriticProvider, + InterleaveOrchestrator, + PlannerProvider, + SinglePromptPlanner, +) +from fastvideo.entrypoints.interleave.providers import ( + InterleaveThinkerCriticProvider, + InterleaveThinkerPlannerProvider, +) +from fastvideo.entrypoints.interleave.schema import InterleaveTrace +from fastvideo.entrypoints.interleave.trace import save_trace +from fastvideo.train.models.interleave_thinker import ( + InterleaveThinkerCriticModel, + InterleaveThinkerPlannerModel, +) + + +@dataclass(frozen=True) +class InterleaveRunResult: + trace: InterleaveTrace + trace_path: str + + +def run_interleave_config( + config: InterleaveRunConfig, + *, + image_backend: ImageGeneratorBackend | None = None, +) -> InterleaveRunResult: + """Run one native interleaved generation trace from a typed config.""" + + cleanup: Callable[[], None] = lambda: None + if image_backend is None: + image_backend, cleanup = build_image_backend(config) + + try: + orchestrator = InterleaveOrchestrator( + planner=build_planner(config.planner), + generator=image_backend, + critic=build_critic(config.critic), + ) + trace = orchestrator.run( + resolve_interleave_instruction(config), + initial_image_path=config.interleave.initial_image_path, + metadata={ + "image_backend": config.image_backend.kind, + "planner": config.planner.kind, + "critic": config.critic.kind, + }, + ) + trace_path = resolve_trace_path(config) + save_trace( + trace, + trace_path, + include_images=config.interleave.include_images_in_trace, + ) + return InterleaveRunResult( + trace=trace, + trace_path=str(trace_path), + ) + finally: + cleanup() + + +def resolve_trace_path(config: InterleaveRunConfig) -> Path: + if config.interleave.trace_path: + return Path(config.interleave.trace_path) + return Path(config.interleave.output_dir) / "trace.json" + + +def build_planner(config: InterleavePlannerConfig) -> PlannerProvider: + if config.kind == "single_prompt": + return SinglePromptPlanner(max_attempts=config.max_attempts_per_step) + + model = InterleaveThinkerPlannerModel( + **_actor_model_kwargs(config), + ) + return InterleaveThinkerPlannerProvider( + model, + num_generations=config.num_generations, + temperature=config.temperature, + top_p=config.top_p, + max_new_tokens=config.max_new_tokens or config.max_response_length, + max_attempts_per_step=config.max_attempts_per_step, + ) + + +def build_critic(config: InterleaveCriticConfig) -> CriticProvider | None: + if config.kind == "none": + return None + if config.kind == "accept_all": + return AcceptAllCritic() + + model = InterleaveThinkerCriticModel( + **_actor_model_kwargs(config), + ) + return InterleaveThinkerCriticProvider( + model, + num_generations=config.num_generations, + temperature=config.temperature, + top_p=config.top_p, + max_new_tokens=config.max_new_tokens or config.max_response_length, + ) + + +def build_image_backend(config: InterleaveRunConfig) -> tuple[ImageGeneratorBackend, Callable[[], None]]: + image_config = config.image_backend + output_dir = image_config.output_dir or config.interleave.output_dir + if image_config.kind == "nano_banana": + return ( + NanoBananaImageGeneratorBackend( + model=image_config.model, + api_key=image_config.api_key, + base_url=image_config.base_url, + output_dir=output_dir, + aspect_ratio=image_config.aspect_ratio, + image_size=image_config.image_size, + max_attempts=image_config.max_attempts, + retry_delay_s=image_config.retry_delay_s, + ), + lambda: None, + ) + return _build_fastvideo_image_backend(config, image_config, output_dir) + + +def _build_fastvideo_image_backend( + config: InterleaveRunConfig, + image_config: InterleaveImageBackendConfig, + output_dir: str, +) -> tuple[ImageGeneratorBackend, Callable[[], None]]: + del image_config + if config.generator is None: + raise ValueError("FastVideo image backend requires config.generator") + + from fastvideo import VideoGenerator + + generator = VideoGenerator.from_config(config.generator) + + def cleanup() -> None: + generator.shutdown() + + return ( + FastVideoImageGeneratorBackend( + generator, + output_dir=output_dir, + default_request=config.request, + ), + cleanup, + ) + + +def _actor_model_kwargs(config: InterleavePlannerConfig | InterleaveCriticConfig) -> dict[str, Any]: + kwargs: dict[str, Any] = { + "load_backend": config.load_backend, + "trainable": config.trainable, + "image_dir": config.image_dir, + "torch_dtype": config.torch_dtype, + "device_map": config.device_map, + "attn_implementation": config.attn_implementation, + "trust_remote_code": config.trust_remote_code, + "use_cache": config.use_cache, + "max_prompt_length": config.max_prompt_length, + "max_response_length": config.max_response_length, + "lora": config.lora, + } + if config.init_from is not None: + kwargs["init_from"] = config.init_from + if config.processor_from is not None: + kwargs["processor_from"] = config.processor_from + return kwargs + + +__all__ = [ + "InterleaveRunResult", + "build_critic", + "build_image_backend", + "build_planner", + "resolve_trace_path", + "run_interleave_config", +] diff --git a/tests/local_tests/test_interleave_run_cli.py b/tests/local_tests/test_interleave_run_cli.py new file mode 100644 index 000000000..680ed7036 --- /dev/null +++ b/tests/local_tests/test_interleave_run_cli.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import base64 +from pathlib import Path + +import pytest + +from fastvideo.entrypoints.interleave import ( + GeneratedImage, + InterleaveEditRequest, + load_interleave_run_config, + resolve_interleave_instruction, + run_interleave_config, +) +from fastvideo.entrypoints.interleave.orchestrator import SinglePromptPlanner +from fastvideo.entrypoints.interleave.schema import PlannerInput + + +def test_interleave_run_config_loads_prompt_and_request_defaults(tmp_path: Path) -> None: + config_path = _write_run_config(tmp_path) + + config = load_interleave_run_config(str(config_path)) + + assert resolve_interleave_instruction(config) == "draw a red mug" + assert config.generator is not None + assert config.generator.model_path == "black-forest-labs/FLUX.2-klein-4B" + assert config.request.sampling.width == 512 + assert config.request.sampling.num_inference_steps == 4 + + +def test_interleave_run_config_accepts_cli_fields_and_dotted_overrides(tmp_path: Path) -> None: + config_path = _write_run_config(tmp_path) + + config = load_interleave_run_config( + str(config_path), + prompt="draw a blue mug", + output_dir=str(tmp_path / "override_outputs"), + trace_path=str(tmp_path / "trace_override.json"), + overrides=[ + "--request.sampling.seed", + "99", + "--planner.max-attempts-per-step", + "3", + ], + ) + + assert resolve_interleave_instruction(config) == "draw a blue mug" + assert config.interleave.output_dir == str(tmp_path / "override_outputs") + assert config.interleave.trace_path == str(tmp_path / "trace_override.json") + assert config.request.sampling.seed == 99 + assert config.planner.max_attempts_per_step == 3 + + +def test_interleave_run_config_rejects_unknown_override_prefix(tmp_path: Path) -> None: + config_path = _write_run_config(tmp_path) + + with pytest.raises(ValueError, match="Unsupported override path"): + load_interleave_run_config( + str(config_path), + overrides=["--server.port", "9000"], + ) + + +def test_run_interleave_config_with_injected_backend_writes_trace(tmp_path: Path) -> None: + config_path = _write_run_config(tmp_path) + config = load_interleave_run_config(str(config_path)) + backend = _FakeImageBackend(tmp_path / "generated.png") + + result = run_interleave_config(config, image_backend=backend) + + assert result.trace.success is True + assert result.trace.final_image is not None + assert result.trace.final_image.file_path == str(tmp_path / "generated.png") + assert Path(result.trace_path).exists() + trace_text = Path(result.trace_path).read_text(encoding="utf-8") + assert "draw a red mug" in trace_text + assert "image_base64" not in trace_text + assert backend.requests[0].prompt == "draw a red mug" + + +def test_single_prompt_planner_uses_configured_attempt_count() -> None: + planner = SinglePromptPlanner(max_attempts=4) + + steps = list(planner.plan(PlannerInput(instruction="draw a red mug"))) + + assert len(steps) == 1 + assert steps[0].max_attempts == 4 + + +class _FakeImageBackend: + + def __init__(self, output_path: Path) -> None: + self.requests: list[InterleaveEditRequest] = [] + self.output_path = output_path + + def generate( + self, + request: InterleaveEditRequest, + *, + request_id: str | None = None, + ) -> GeneratedImage: + del request_id + self.requests.append(request) + self.output_path.write_bytes(b"fake-image") + return GeneratedImage( + prompt=request.prompt, + image_base64=base64.b64encode(b"fake-image").decode("utf-8"), + file_path=str(self.output_path), + metadata={"backend": "fake"}, + ) + + +def _write_run_config(tmp_path: Path) -> Path: + output_path = tmp_path / "generated.png" + config_path = tmp_path / "interleave_run.yaml" + config_path.write_text( + f""" +generator: + model_path: black-forest-labs/FLUX.2-klein-4B + engine: + num_gpus: 1 + pipeline: + workload_type: t2i + +image_backend: + kind: fastvideo + +planner: + kind: single_prompt + max_attempts_per_step: 1 + +critic: + kind: accept_all + +interleave: + output_dir: {tmp_path} + trace_path: {tmp_path / "trace.json"} + +request: + prompt: draw a red mug + sampling: + width: 512 + height: 512 + seed: 7 + num_inference_steps: 4 + extensions: + test_output_path: {output_path} +""", + encoding="utf-8", + ) + return config_path