Davids048andClaude Fable 5.1
a5aa64ab6b
VideoGenerator.from_pretrained(model_path, config) takes the same nested
settings as from_config; the 24 flat keywords, from_pretrained_kwargs_to_config,
FROM_PRETRAINED_KWARGS, and every flat_field in the schema are deleted.
nvfp4_fa4 is the typed field engine.attention.nvfp4_fa4, applied by a
resolution step. 65 call sites and the docs use the nested form; the three
kwargs golden cases are config cases with byte-identical results.
Every value is decided during resolution; no runtime code calls
with_override. The device offload policy (unified memory, MPS, layerwise
conflicts, lazy module load) runs as resolution steps in the main process
(fastvideo/api/device_policy.py), and the worker runs on the config it
receives. The LTX-2 refine defaults from model_index.json and the MiniMax-H3
schedule from fastvideo_inference.json are resolution steps
(fastvideo/api/checkpoint_defaults.py) that read local checkpoints or download
only those files; the H3 pipeline validates the schedule against the loaded
schedulers but no longer writes it. The preprocessing entry point passes the
downloaded local path as run state instead of a config override. Teacher and
critic models load with override_transformer_cls_name as a load argument.
The test isolation turns the device policy off, as it blocks downloads, so
that the goldens do not depend on the machine.
Against the 3f6893a0 goldens, every field still matches except boundary_ratio
and ltx2_vae_tiling (DESIGN section 9) and the fields that no longer exist.
The launcher and trainer YAML verifications report 0 unexplained differences.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
134 lines
3.9 KiB
Python
134 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Per-GPU worker subprocess entry for :class:`SubprocessGpuPool`.
|
|
|
|
The pool manages binding, lifecycle, and message dispatch in the parent
|
|
process. The worker constructs its :class:`VideoGenerator` from a typed
|
|
:class:`GeneratorConfig`, runs the two-segment warmup so both
|
|
initial-segment and continuation-branch compile graphs are hot, and
|
|
then loops on the job queue.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import multiprocessing as mp
|
|
import queue
|
|
from typing import Any
|
|
|
|
from fastvideo.api.schema import (
|
|
GeneratorConfig,
|
|
GenerationRequest,
|
|
InputConfig,
|
|
OutputConfig,
|
|
SamplingConfig,
|
|
WarmupConfig,
|
|
)
|
|
from fastvideo.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# Synthetic warmup dimensions: small enough to keep boot fast, big enough
|
|
# to exercise the real shape-dependent compile paths. Keep in sync with
|
|
# WarmupConfig if these become user-tunable.
|
|
_WARMUP_NUM_FRAMES = 8
|
|
_WARMUP_HEIGHT = 256
|
|
_WARMUP_WIDTH = 256
|
|
_WARMUP_NUM_INFERENCE_STEPS = 1
|
|
|
|
|
|
def worker_main(
|
|
*,
|
|
gpu_id: int,
|
|
worker_id: str,
|
|
generator_config: GeneratorConfig,
|
|
warmup_config: WarmupConfig,
|
|
job_queue: mp.Queue,
|
|
result_queue: mp.Queue,
|
|
shutdown_event: Any,
|
|
) -> None: # pragma: no cover - exercised via integration only
|
|
"""Per-worker subprocess entry.
|
|
|
|
Runs inside the child spawned by ``SubprocessGpuPool``. Blocking
|
|
``VideoGenerator`` construction + generation happens here, not in
|
|
the parent's event loop.
|
|
"""
|
|
import fastvideo.envs as envs
|
|
|
|
envs.set_external("CUDA_VISIBLE_DEVICES", str(gpu_id))
|
|
try:
|
|
from fastvideo import VideoGenerator
|
|
|
|
generator = VideoGenerator.from_config(generator_config)
|
|
if warmup_config.enabled:
|
|
_warmup_worker(generator, warmup_config)
|
|
result_queue.put({"kind": "ready", "worker_id": worker_id})
|
|
except Exception as exc:
|
|
result_queue.put({"kind": "error", "error": repr(exc)})
|
|
return
|
|
|
|
while not shutdown_event.is_set():
|
|
try:
|
|
item = job_queue.get(timeout=0.5)
|
|
except queue.Empty:
|
|
continue
|
|
if item is None:
|
|
break
|
|
job_id = item["job_id"]
|
|
request = item["request"]
|
|
try:
|
|
result = generator.generate(request)
|
|
result_queue.put({
|
|
"kind": "result",
|
|
"job_id": job_id,
|
|
"result": result,
|
|
})
|
|
except Exception as exc:
|
|
result_queue.put({
|
|
"kind": "error",
|
|
"job_id": job_id,
|
|
"error": repr(exc),
|
|
})
|
|
|
|
|
|
def _warmup_worker(
|
|
generator: Any,
|
|
warmup_config: WarmupConfig,
|
|
) -> None:
|
|
"""Run two synthetic generations so both compile branches are primed.
|
|
|
|
Segment 1 is a fresh start (no continuation state) and exercises
|
|
the initial-segment graph. Segment 2 feeds segment 1's continuation
|
|
state back in so the conditioning branch is also compiled before
|
|
the first user request lands.
|
|
"""
|
|
sampling = SamplingConfig(
|
|
num_frames=_WARMUP_NUM_FRAMES,
|
|
height=_WARMUP_HEIGHT,
|
|
width=_WARMUP_WIDTH,
|
|
num_inference_steps=_WARMUP_NUM_INFERENCE_STEPS,
|
|
)
|
|
seg1 = GenerationRequest(
|
|
prompt=warmup_config.prompt,
|
|
sampling=sampling,
|
|
inputs=InputConfig(),
|
|
output=OutputConfig(save_video=False, return_frames=False, return_state=True),
|
|
)
|
|
seg1_result = generator.generate(seg1)
|
|
|
|
seg2 = GenerationRequest(
|
|
prompt=warmup_config.prompt,
|
|
sampling=sampling,
|
|
inputs=InputConfig(),
|
|
output=OutputConfig(save_video=False, return_frames=False),
|
|
state=_extract_continuation_state(seg1_result),
|
|
)
|
|
generator.generate(seg2)
|
|
|
|
|
|
def _extract_continuation_state(result: Any) -> Any:
|
|
state = getattr(result, "state", None)
|
|
if state is None and isinstance(result, dict):
|
|
state = result.get("state")
|
|
return state
|
|
|
|
|
|
__all__ = ["worker_main"]
|