diff --git a/fastvideo/api/inference_resolution.py b/fastvideo/api/inference_resolution.py index 75dfbaf97..e1cda9454 100644 --- a/fastvideo/api/inference_resolution.py +++ b/fastvideo/api/inference_resolution.py @@ -31,7 +31,8 @@ from typing import Any, get_args import fastvideo.envs as envs from fastvideo.api.parser import parse_config -from fastvideo.api.resolution import ResolutionStep, ResolutionView, ResolvedGeneratorConfig, resolve_generator_config +from fastvideo.api.resolution import (ResolutionStep, ResolutionView, ResolvedGeneratorConfig, resolve_generator_config, + thaw) from fastvideo.api.schema import ExecutionMode, GeneratorConfig, WorkloadType from fastvideo.logger import init_logger @@ -303,6 +304,21 @@ def warn_deprecated_environment_variables(view: ResolutionView) -> dict[str, Any return {} +def torch_compile_kwargs(resolved_config: ResolvedGeneratorConfig) -> dict[str, Any]: + """The keyword arguments for ``torch.compile`` that ``engine.compile`` describes. + + ``backend``, ``fullgraph``, ``mode``, and ``dynamic`` are included when they are set; ``extras`` is merged on top. + """ + compile_config = resolved_config.engine.compile + kwargs: dict[str, Any] = {} + for key in ("backend", "fullgraph", "mode", "dynamic"): + value = getattr(compile_config, key) + if value is not None: + kwargs[key] = value + kwargs.update(thaw(compile_config.extras)) + return kwargs + + # Values that runtime code uses for typed fields that no earlier step decided. ``pipeline.ltx2.refine.*`` and # ``engine.offload.lazy_module_load`` are not listed: ``None`` there means "decide later" (the LTX-2 checkpoint defaults # and the device policy). @@ -508,6 +524,7 @@ def _non_default_fields(value: Any, default: Any) -> dict[str, Any]: __all__ = [ + "torch_compile_kwargs", "ENVIRONMENT_STEPS", "FLAT_INPUT_STEPS", "VALIDATION_STEPS", diff --git a/fastvideo/api/resolution.py b/fastvideo/api/resolution.py index cd92c0afa..727f3f3fd 100644 --- a/fastvideo/api/resolution.py +++ b/fastvideo/api/resolution.py @@ -139,6 +139,17 @@ def _freeze(value: Any) -> Any: return deepcopy(value) +def thaw(value: Any) -> Any: + """Mutable copy of a value read from a resolved config: mapping proxies become dicts, tuples become lists.""" + if isinstance(value, Mapping): + return {key: thaw(child) for key, child in value.items()} + if isinstance(value, tuple): + return [thaw(child) for child in value] + if isinstance(value, frozenset): + return {thaw(child) for child in value} + return deepcopy(value) + + def _split(path: str) -> list[str] | None: """Dotted path parts, or ``None`` when ``path`` is not a non-empty dotted string.""" if not isinstance(path, str) or not path or any(not part for part in path.split(".")): @@ -564,4 +575,5 @@ __all__ = [ "ResolvedRequest", "resolve_generation_request", "resolve_generator_config", + "thaw", ] diff --git a/fastvideo/pipelines/component_state.py b/fastvideo/pipelines/component_state.py new file mode 100644 index 000000000..1b311a897 --- /dev/null +++ b/fastvideo/pipelines/component_state.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Runtime state of a pipeline's components: where each one was loaded from, and whether it is resident. + +A pipeline owns one ``ComponentState``. ``ComposedPipelineBase`` records each component's path after loading it, and +attaches the state to every stage it registers (``PipelineStage.component_state``), so that a stage can reload a +component that it released to save memory. +""" +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass +class ComponentState: + """Mutable per-pipeline component bookkeeping that is not configuration.""" + + # Local path that each component was loaded from, keyed by module name. + model_paths: dict[str, str] = field(default_factory=dict) + # Whether each releasable component is currently loaded, keyed by module name. + model_loaded: dict[str, bool] = field(default_factory=lambda: { + "transformer": True, + "vae": True, + "upsampler": True, + })