diff --git a/docs/inference/h3_community_lora.md b/docs/inference/h3_community_lora.md index 744450200..4e66903eb 100644 --- a/docs/inference/h3_community_lora.md +++ b/docs/inference/h3_community_lora.md @@ -109,6 +109,64 @@ the consumer profiles already use, for packed NVFP4 and for bf16. Pass `lora_pat adapter is attached when the transformer arrives: as a bf16 side branch beside packed weights, or merged into a dense bf16 transformer. You do not load the LoRA as a separate resident copy of the model. +## Per-request adapters on the OpenAI API + +`fastvideo serve` can hold an allow-list of adapters and apply one per request on `POST /v1/videos/sync` (the whole +MP4 comes back in the response) and on the async `POST /v1/videos` job route. Name each adapter and its strength +range in the serve config; see `examples/serving/openai_fasth3_nvfp4_lora_allowlist.yaml`: + +```yaml +generator: + model_path: /models/FastVideo-FastH3-8-Step-V2.1-NVFP4-Consumer + engine: + quantization: + transformer_quant: NVFP4 # required: per-request adapters need the packed NVFP4 side-branch layers + layer_profile: h3_dit_vsa +lora: + adapters: + - id: fal/MiniMax-H3-Realism-People-LoRA + path: fal/MiniMax-H3-Realism-People-LoRA + default_strength: 0.8 + min_strength: 0.0 + max_strength: 1.5 + license: Apache-2.0 + source_url: https://huggingface.co/fal/MiniMax-H3-Realism-People-LoRA +``` + +An allow-list cannot be combined with a startup adapter (`generator.pipeline.components.lora_path`); the server +refuses to start with both. Without an allow-list, the startup-adapter behavior is unchanged. + +`GET /v1/model_info` reports the allow-list in its `lora` block, never the load paths. A request selects one +adapter with the vLLM-style `lora` object: + +```bash +curl -s http://127.0.0.1:8000/v1/videos/sync \ + -H 'content-type: application/json' \ + -d '{"prompt": "...", "size": "1344x768", "num_frames": 124, "seed": 7, + "lora": {"name": "fal/MiniMax-H3-Realism-People-LoRA", "scale": 0.8}}' \ + -o out.mp4 +``` + +`name` (or `lora_name`, `adapter`) must be an allow-listed `id`. `scale` (or `lora_scale`) is the strength; it +defaults to the adapter's `default_strength` and must lie in `[min_strength, max_strength]`. A client-supplied `path` +is rejected. A bad selection gets HTTP 400 before any job is stored. + +Strength `0` is exactly the base model, with output identical to a request without `lora`. A request without `lora` +sets the last-used adapter to strength 0 instead of unloading it, so one request never inherits another's +adapter. The server applies the adapter under the same lock as the generation, so requests never interleave. If +a change fails, that request fails with the loader's message, and the server clears every LoRA layer back to the +base weights without reading an adapter file, so it never renders with an adapter the request did not select and +requests without `lora` keep working. + +Per-request adapters need side-branch LoRA layers, so the server refuses an allow-list unless the model is a packed +NVFP4 export (a `transformer/nvfp4_weights.safetensors` file with `transformer_quant: NVFP4`) or +`FASTVIDEO_LORA_SIDE_BRANCH=1` is set. At startup it also checks that every adapter resolves to a safetensors file +and is purely low-rank: adapters with dense replacement payloads (`.diff` and `.set_weight` tensors, as in the +[FastH3 adapter](fasth3-distilled.md)) can only be applied when the model loads, and the server refuses to start with +one. A side-branch layer holds one adapter at a time. Linears that a packed export leaves unquantized still merge the +adapter into their weights; switching to another adapter first restores each of them, including those the new +adapter does not cover. + ## Limitations - Block-pruned models (FastH3 Trim, 42 of 50 blocks) are not supported yet. Community LoRAs index the full diff --git a/examples/serving/openai_fasth3_nvfp4_lora_allowlist.yaml b/examples/serving/openai_fasth3_nvfp4_lora_allowlist.yaml new file mode 100644 index 000000000..481a01ac7 --- /dev/null +++ b/examples/serving/openai_fasth3_nvfp4_lora_allowlist.yaml @@ -0,0 +1,47 @@ +# OpenAI-compatible FastH3 V2.1 NVFP4 server with a per-request community LoRA allow-list. +# Do not combine with generator.pipeline.components.lora_path (a startup adapter); the server refuses both. +# Clients select an adapter with {"lora": {"name": "", "scale": 0.8}} on POST /v1/videos/sync. +generator: + model_path: /path/to/FastVideo-FastH3-8-Step-V2.1-NVFP4-Consumer + engine: + num_gpus: 1 + quantization: + transformer_quant: NVFP4 + layer_profile: h3_dit_vsa + offload: + dit: true + dit_layerwise: true + text_encoder: true + image_encoder: true + vae: true + pin_cpu_memory: true + compile: + enabled: false + pipeline: + workload_type: t2v + +server: + host: 0.0.0.0 + port: 8000 + output_dir: outputs/openai_fasth3_lora_allowlist + +lora: + adapters: + - id: fal/MiniMax-H3-Realism-People-LoRA + path: fal/MiniMax-H3-Realism-People-LoRA + default_strength: 0.8 + min_strength: 0.0 + max_strength: 1.5 + description: Photoreal people and faces for MiniMax-H3. + license: Apache-2.0 + source_url: https://huggingface.co/fal/MiniMax-H3-Realism-People-LoRA + +default_request: + negative_prompt: "" + sampling: + height: 768 + width: 1344 + num_frames: 124 + fps: 24 + # The V2.1 consumer checkpoint samples a 9-point sigma grid for its 8 DMD forwards. + num_inference_steps: 9 diff --git a/fastvideo/api/__init__.py b/fastvideo/api/__init__.py index cb4a35136..b89288570 100644 --- a/fastvideo/api/__init__.py +++ b/fastvideo/api/__init__.py @@ -11,6 +11,7 @@ from fastvideo.api.schema import ( GeneratorConfig, GpuPoolConfig, InputConfig, + LoRAAdapterConfig, OffloadConfig, OutputConfig, ParallelismConfig, @@ -23,6 +24,7 @@ from fastvideo.api.schema import ( RunConfig, SamplingConfig, ServeConfig, + ServeLoRAConfig, ServerConfig, StreamingConfig, WarmupConfig, @@ -84,6 +86,7 @@ __all__ = [ "GeneratorConfig", "GpuPoolConfig", "InputConfig", + "LoRAAdapterConfig", "MiniMaxH3Reference", "OffloadConfig", "OutputConfig", @@ -98,6 +101,7 @@ __all__ = [ "SamplingConfig", "SamplingParam", "ServeConfig", + "ServeLoRAConfig", "ServerConfig", "StreamingConfig", "VideoEvent", diff --git a/fastvideo/api/schema.py b/fastvideo/api/schema.py index 44222d157..2c8c03d37 100644 --- a/fastvideo/api/schema.py +++ b/fastvideo/api/schema.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations +import math from dataclasses import dataclass, field from typing import Any, Literal @@ -308,6 +309,50 @@ class StreamingConfig: safety: PromptSafetyConfig = field(default_factory=PromptSafetyConfig) +@dataclass +class LoRAAdapterConfig: + """One community LoRA that a server may apply to a request. + + ``id`` is the only name a client sees. ``path`` is the local file or directory, or the Hugging Face repo id, + that ``VideoGenerator.set_lora_adapter`` loads. Strength 0 is exactly the base model. + """ + id: str + path: str + default_strength: float = 0.8 + min_strength: float = 0.0 + max_strength: float = 2.0 + description: str = "" + license: str = "" + source_url: str = "" + + def __post_init__(self) -> None: + if not self.id.strip(): + raise ValueError("lora.adapters[].id must not be empty") + if not self.path.strip(): + raise ValueError(f"LoRA adapter {self.id!r} has an empty path") + if not all( + math.isfinite(strength) for strength in (self.min_strength, self.default_strength, self.max_strength)): + raise ValueError(f"LoRA adapter {self.id!r} has a non-finite strength") + if not self.min_strength <= self.default_strength <= self.max_strength: + raise ValueError(f"LoRA adapter {self.id!r} needs min_strength <= default_strength <= " + f"max_strength, got {self.min_strength}, {self.default_strength}, {self.max_strength}") + + +@dataclass +class ServeLoRAConfig: + """The LoRAs a request to the OpenAI server may select. + + No adapter loads at startup: the transformer gains LoRA layers when the first request selects one. + """ + adapters: list[LoRAAdapterConfig] = field(default_factory=list) + + def __post_init__(self) -> None: + ids = [adapter.id for adapter in self.adapters] + duplicates = sorted({adapter_id for adapter_id in ids if ids.count(adapter_id) > 1}) + if duplicates: + raise ValueError(f"lora.adapters repeats id(s): {', '.join(duplicates)}") + + @dataclass class ServeConfig: """Typed serve config loaded from ``fastvideo serve --config``. @@ -333,6 +378,16 @@ class ServeConfig: server: ServerConfig = field(default_factory=ServerConfig) default_request: GenerationRequest = field(default_factory=GenerationRequest) streaming: StreamingConfig | None = None + lora: ServeLoRAConfig | None = None + + def __post_init__(self) -> None: + if not self.lora or not self.lora.adapters: + return + if self.generator.pipeline.components.lora_path: + raise ValueError("`lora.adapters` cannot be combined with a startup adapter " + "(generator.pipeline.components.lora_path); use one or the other.") + if self.streaming is not None: + raise ValueError("`lora.adapters` applies to the OpenAI server and cannot be combined with `streaming:`.") __all__ = [ @@ -345,6 +400,7 @@ __all__ = [ "GeneratorConfig", "GpuPoolConfig", "InputConfig", + "LoRAAdapterConfig", "OffloadConfig", "OutputConfig", "ParallelismConfig", @@ -357,6 +413,7 @@ __all__ = [ "RunConfig", "SamplingConfig", "ServeConfig", + "ServeLoRAConfig", "ServerConfig", "StreamingConfig", "WarmupConfig", diff --git a/fastvideo/entrypoints/cli/serve.py b/fastvideo/entrypoints/cli/serve.py index 42eae20be..336b7ac85 100644 --- a/fastvideo/entrypoints/cli/serve.py +++ b/fastvideo/entrypoints/cli/serve.py @@ -60,6 +60,10 @@ class ServeSubcommand(CLISubcommand): serve_config.server.port, ) + from fastvideo.entrypoints.openai.lora import build_lora_runtime + lora_runtime = build_lora_runtime(serve_config) + lora_kwargs = {} if lora_runtime is None else {"lora_runtime": lora_runtime} + fastvideo_args = generator_config_to_fastvideo_args(serve_config.generator) run_server( fastvideo_args, @@ -73,6 +77,7 @@ class ServeSubcommand(CLISubcommand): roots=tuple(serve_config.server.media_roots), allow_private_urls=serve_config.server.allow_private_media_urls, ), + **lora_kwargs, ) def validate(self, args: argparse.Namespace) -> None: diff --git a/fastvideo/entrypoints/openai/api_server.py b/fastvideo/entrypoints/openai/api_server.py index baea043e2..dfe20c476 100644 --- a/fastvideo/entrypoints/openai/api_server.py +++ b/fastvideo/entrypoints/openai/api_server.py @@ -13,6 +13,7 @@ from fastapi.responses import JSONResponse from fastvideo.api.presets import validate_preset_selection from fastvideo.api.schema import DEFAULT_MEDIA_REFERENCE_MAX_BYTES, GenerationRequest +from fastvideo.entrypoints.openai.lora import LoraRuntime from fastvideo.entrypoints.openai.media_policy import MediaPolicy from fastvideo.entrypoints.openai.state import ( DEFAULT_OUTPUT_DIR, @@ -68,7 +69,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: logger.info("Initializing %s generation runtime for %s ...", app.state.runtime, args.model_path) factory = app.state.generator_factory generator = factory() if factory is not None else VideoGenerator.from_fastvideo_args(args) - serving_engine = OpenAIServingEngine(generator, app.state.video_request_validator) + serving_engine = OpenAIServingEngine(generator, app.state.video_request_validator, app.state.lora_runtime) logger.info("Generation runtime ready.") set_state( @@ -105,6 +106,7 @@ def create_app( runtime: str = "cuda", media_reference_max_bytes: int = DEFAULT_MEDIA_REFERENCE_MAX_BYTES, media_policy: MediaPolicy | None = None, + lora_runtime: LoraRuntime | None = None, ) -> FastAPI: """Build the FastAPI application with all routers mounted""" @@ -129,6 +131,7 @@ def create_app( app.state.generator_factory = generator_factory app.state.video_request_validator = video_request_validator app.state.runtime = runtime + app.state.lora_runtime = lora_runtime app.add_middleware( CORSMiddleware, @@ -222,6 +225,7 @@ def run_server( served_model_name: str | None = None, media_reference_max_bytes: int = DEFAULT_MEDIA_REFERENCE_MAX_BYTES, media_policy: MediaPolicy | None = None, + lora_runtime: LoraRuntime | None = None, ): """Create the app and run it with uvicorn""" os.environ.setdefault("FASTVIDEO_STAGE_LOGGING", "1") @@ -235,6 +239,7 @@ def run_server( served_model_name=served_model_name, media_reference_max_bytes=media_reference_max_bytes, media_policy=media_policy, + lora_runtime=lora_runtime, ) logger.info("Starting FastVideo server on %s:%d", host, port) diff --git a/fastvideo/entrypoints/openai/common_api.py b/fastvideo/entrypoints/openai/common_api.py index 27245349b..0371d0cea 100644 --- a/fastvideo/entrypoints/openai/common_api.py +++ b/fastvideo/entrypoints/openai/common_api.py @@ -7,7 +7,7 @@ from fastapi import APIRouter from fastapi.responses import ORJSONResponse from pydantic import BaseModel, Field -from fastvideo.entrypoints.openai.state import get_served_model_name, get_server_args +from fastvideo.entrypoints.openai.state import get_served_model_name, get_server_args, get_serving_engine from fastvideo.logger import init_logger router = APIRouter(prefix="/v1") @@ -58,12 +58,13 @@ async def retrieve_model(model: str): async def model_info(): """Get basic model information""" args = get_server_args() + lora_runtime = get_serving_engine().lora_runtime return { "model_path": args.model_path, "served_model_name": get_served_model_name(), - "lora": ({ + "lora": (lora_runtime.info() if lora_runtime is not None else { "name": args.lora_nickname, "path": args.lora_path, "scale": args.lora_strength, diff --git a/fastvideo/entrypoints/openai/lora.py b/fastvideo/entrypoints/openai/lora.py new file mode 100644 index 000000000..591f0f167 --- /dev/null +++ b/fastvideo/entrypoints/openai/lora.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-request community LoRA selection for the OpenAI-compatible video API. + +The server owns one generator and FastVideo keeps the active adapter on its pipeline, so a job applies its adapter +under the same lock as its generation. Strength 0 is exactly the base model, so a job without ``lora`` that follows +one with it runs at strength 0. +""" +from __future__ import annotations + +import math +import os +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from fastvideo import envs +from fastvideo.api.schema import GeneratorConfig, LoRAAdapterConfig, ServeConfig, ServeLoRAConfig +from fastvideo.fastvideo_args import WorkloadType +from fastvideo.layers.quantization.nvfp4_config import (H3_NVFP4_DIT_EXPORT_FILENAME, find_minimax_h3_nvfp4_dit_export) +from fastvideo.logger import init_logger +from fastvideo.models.loader.lora_patch import DenseLoRAPatch +from fastvideo.registry import get_model_info +from fastvideo.utils import maybe_download_lora, maybe_download_model + +logger = init_logger(__name__) + +_NAME_KEYS = ("name", "lora_name", "adapter") +_PATH_KEYS = ("path", "lora_path", "local_path") +_SCALE_KEYS = ("scale", "lora_scale") + + +class LoraRequestError(ValueError): + """A request's LoRA selection cannot be served; the message is the client-facing reason.""" + + +@dataclass(frozen=True) +class LoraSelection: + """A validated request selection with its strength resolved.""" + id: str + strength: float + + +def adapters_attach_as_side_branch(generator: GeneratorConfig) -> bool: + """Whether adapters attach beside the base weights, so a request can change them on the running server. + + That holds with ``FASTVIDEO_LORA_SIDE_BRANCH=1`` and for a packed NVFP4 export, whose linears have no dense + weight to merge into. The export is looked for where the transformer loads from: ``transformer_weights`` when + set, otherwise the pipeline's transformer partition (``transformer_ref`` for Ref2VA). A Hugging Face model id + is resolved to its snapshot first, fetching only the packed transformer file. + """ + if envs.FASTVIDEO_LORA_SIDE_BRANCH.get(): + return True + quantization = generator.engine.quantization + if quantization is None or quantization.transformer_quant != "NVFP4": + return False + weights = generator.pipeline.components.transformer_weights + if weights is not None: + return find_minimax_h3_nvfp4_dit_export([weights]) is not None + export = f"{_transformer_partition(generator)}/{H3_NVFP4_DIT_EXPORT_FILENAME}" + try: + model_dir = maybe_download_model(generator.model_path, revision=generator.revision, allow_patterns=[export]) + except ValueError: + logger.warning("Could not resolve %s to check for a packed NVFP4 export", generator.model_path, exc_info=True) + return False + return (Path(model_dir) / export).is_file() + + +def _transformer_partition(generator: GeneratorConfig) -> str: + """The model directory the selected pipeline loads its transformer from.""" + workload = generator.pipeline.workload_type + try: + pipeline_cls = get_model_info( + generator.model_path, + workload_type=None if workload is None else WorkloadType.from_string(workload), + override_pipeline_cls_name=generator.pipeline.components.override_pipeline_cls_name, + revision=generator.revision, + ).pipeline_cls + except ValueError: + logger.warning("Could not resolve the pipeline for %s; looking for the packed export in transformer/", + generator.model_path, + exc_info=True) + return "transformer" + return getattr(pipeline_cls, "_extra_config_module_map", {}).get("transformer", "transformer") + + +def check_adapter_file(adapter: LoRAAdapterConfig) -> None: + """Raise ``ValueError`` for an adapter the pipeline would refuse before changing any weight. + + The adapter must resolve to a safetensors file, and must carry no dense payload, which can only be applied + when the model loads. + """ + local_path = maybe_download_lora(adapter.path) + if not os.path.isfile(local_path): + raise ValueError(f"LoRA adapter {adapter.id!r} resolves to {local_path!r}, which is not a file") + if DenseLoRAPatch.from_adapter(local_path) is not None: + raise ValueError(f"LoRA adapter {adapter.id!r} carries dense replacement tensors; only low-rank adapters " + "can change on a running server") + + +def build_lora_runtime(serve_config: ServeConfig) -> LoraRuntime | None: + """The runtime for the serve config's allow-list, or ``None`` without one; raises ``ValueError`` at startup. + + Per-request adapters are served only where the quantized linears take adapters as side branches, so a switch + rewrites only the few linears a packed export leaves unquantized; the pipeline restores those from their original + weights. Every adapter is checked here so that an unusable one stops the server. + """ + if serve_config.lora is None or not serve_config.lora.adapters: + return None + if not adapters_attach_as_side_branch(serve_config.generator): + raise ValueError("Per-request LoRA adapters need side-branch LoRA layers: serve a packed NVFP4 export " + "(transformer_quant: NVFP4) or set FASTVIDEO_LORA_SIDE_BRANCH=1.") + runtime = LoraRuntime(serve_config.lora) + for adapter in serve_config.lora.adapters: + check_adapter_file(adapter) + return runtime + + +def _first(body: dict[str, Any], keys: tuple[str, ...]) -> Any: + return next((body[key] for key in keys if body.get(key) is not None), None) + + +def _adapter_info(adapter: LoRAAdapterConfig) -> dict[str, Any]: + """What a client may select; never the local load path.""" + entry: dict[str, Any] = { + "id": adapter.id, + "default_strength": adapter.default_strength, + "min_strength": adapter.min_strength, + "max_strength": adapter.max_strength, + } + for key in ("description", "license", "source_url"): + if getattr(adapter, key): + entry[key] = getattr(adapter, key) + return entry + + +@dataclass +class LoraRuntime: + """The adapters a request may select, and the adapter the generator currently carries.""" + config: ServeLoRAConfig + check_adapter: Callable[[LoRAAdapterConfig], None] = check_adapter_file + current_id: str | None = None + current_strength: float | None = None + # Set only when clearing after a failed change failed too; only a successful clear resets it. + unknown_state: bool = False + adapters: dict[str, LoRAAdapterConfig] = field(init=False) + loaded: set[str] = field(init=False, default_factory=set) + + def __post_init__(self) -> None: + self.adapters = {adapter.id: adapter for adapter in self.config.adapters} + + def info(self) -> dict[str, Any]: + """The ``lora`` block of ``GET /v1/model_info``.""" + return { + "supported": True, + "runtime_switching": True, + "base_strength": 0.0, + "adapters": [_adapter_info(adapter) for adapter in self.config.adapters], + } + + def resolve(self, body: dict[str, Any] | None) -> LoraSelection | None: + """Validate a request's ``lora`` object against the allow-list.""" + if body is None: + return None + if _first(body, _PATH_KEYS) is not None: + raise LoraRequestError("lora.path is not accepted; select an adapter by name from GET /v1/model_info.") + name = _first(body, _NAME_KEYS) + if name is None: + raise LoraRequestError("lora must provide a name") + adapter = self.adapters.get(str(name)) + if adapter is None: + known = ", ".join(sorted(self.adapters)) or "none" + raise LoraRequestError(f"Unknown LoRA adapter {str(name)!r}; this server serves: {known}.") + scale = _first(body, _SCALE_KEYS) + try: + strength = adapter.default_strength if scale is None else float(scale) + except (TypeError, ValueError) as error: + raise LoraRequestError(f"Invalid LoRA scale {scale!r}") from error + if not math.isfinite(strength) or not adapter.min_strength <= strength <= adapter.max_strength: + raise LoraRequestError(f"LoRA adapter {adapter.id!r} takes a scale from {adapter.min_strength:g} to " + f"{adapter.max_strength:g}; got {strength:g}.") + return LoraSelection(adapter.id, strength) + + def apply(self, generator: Any, selection: LoraSelection | None) -> None: + """Make the selection current on ``generator``; ``None`` sets the current adapter to strength 0. + + Raises when the generator cannot be cleared back to the base model after a failed change, so the server + never renders with an adapter nobody asked for. + """ + if self.unknown_state: + # The layers may hold part of an adapter, and a retried selection the pipeline already recorded would + # skip repairing them, so only a successful clear makes any request servable again. + self._clear(generator) + if self.unknown_state: + raise RuntimeError("A failed LoRA adapter change left this server's adapter state unknown, and " + "clearing the adapter to restore the base model failed too.") + if selection is None: + if self.current_id is None or self.current_strength == 0.0: + return + self._set(generator, self.current_id, 0.0) + return + if self.current_id == selection.id and self.current_strength == selection.strength: + return + self._set(generator, selection.id, selection.strength) + + def invalidate(self) -> None: + """Mark the adapter state unknown after a generation failed with an adapter set. + + Applying a held adapter to a transformer that just loaded happens inside generation, so a failed generation + may have dropped it; the next request clears before it selects anything. + """ + if self.current_id is not None: + self.unknown_state = True + + def _set(self, generator: Any, adapter_id: str, strength: float) -> None: + if adapter_id not in self.loaded: + try: + self.check_adapter(self.adapters[adapter_id]) + except ValueError as error: + logger.warning("LoRA adapter %r failed its pre-load check: %s", adapter_id, error) + raise LoraRequestError(f"LoRA adapter {adapter_id!r} cannot be loaded on this server.") from error + logger.info("OpenAI server applies LoRA adapter %r at strength %g", adapter_id, strength) + try: + generator.set_lora_adapter(adapter_id, self.adapters[adapter_id].path, strength=strength) + except (BrokenPipeError, EOFError): + # A dead worker: the serving engine marks itself unhealthy, and a clear would only reach the same worker. + raise + except Exception as error: + logger.exception("LoRA adapter %r could not be applied", adapter_id) + # A failed change may keep the previous adapter or leave a partial one; clearing gives the base model. + self._clear(generator) + raise LoraRequestError(f"LoRA adapter {adapter_id!r} could not be applied on this server.") from error + self.loaded.add(adapter_id) + self.current_id = adapter_id + self.current_strength = strength + self.unknown_state = False + + def _clear(self, generator: Any) -> None: + self.current_id = None + self.current_strength = None + try: + generator.clear_lora_adapter() + except (BrokenPipeError, EOFError): + self.unknown_state = True + raise + except Exception: + logger.exception("OpenAI server could not clear the LoRA adapter") + self.unknown_state = True + return + self.unknown_state = False diff --git a/fastvideo/entrypoints/openai/request_adapter.py b/fastvideo/entrypoints/openai/request_adapter.py index 356892d16..070587d51 100644 --- a/fastvideo/entrypoints/openai/request_adapter.py +++ b/fastvideo/entrypoints/openai/request_adapter.py @@ -9,6 +9,7 @@ import os from typing import Any from urllib.parse import urlparse +from fastvideo.entrypoints.openai.lora import LoraRequestError, LoraRuntime, LoraSelection from fastvideo.entrypoints.openai.media_policy import ( MediaPolicy, MediaPolicyError, @@ -312,8 +313,11 @@ def validate_model_and_lora( request: VideoGenerationRequest, args: FastVideoArgs, served_model_name: str, -) -> None: - """Validate vLLM-style model and LoRA selectors against startup state. + lora_runtime: LoraRuntime | None = None, +) -> LoraSelection | None: + """Validate vLLM-style model and LoRA selectors; return the per-request selection when an allow-list is set. + + With a ``lora_runtime`` allow-list, ``lora.name`` must be an allow-listed id and ``lora.scale`` its strength. FastVideo's published FastH3 adapters include dense replacement tensors in addition to low-rank factors. Those tensors are applied while the model is @@ -322,8 +326,14 @@ def validate_model_and_lora( """ validate_served_model_name(request.model, args, served_model_name) + if lora_runtime is not None: + try: + return lora_runtime.resolve(request.lora) + except LoraRequestError as error: + raise RequestAdaptationError(str(error)) from error + if request.lora is None: - return + return None if not args.lora_path: raise RequestAdaptationError( "This server has no startup LoRA. Configure generator.pipeline.components.lora_path before using " @@ -348,6 +358,7 @@ def validate_model_and_lora( if not math.isclose(scale_value, args.lora_strength, rel_tol=0.0, abs_tol=1e-8): raise RequestAdaptationError( f"Requested LoRA scale {scale_value:g} does not match startup strength {args.lora_strength:g}.") + return None def _apply_reference_inputs( @@ -431,9 +442,10 @@ def build_generation_request( served_model_name: str, output_dir: str, default_request: GenerationRequest | None = None, + lora_runtime: LoraRuntime | None = None, ) -> GenerationRequest: """Build one tracked FastVideo request using explicit-field precedence.""" - validate_model_and_lora(request, args, served_model_name) + validate_model_and_lora(request, args, served_model_name, lora_runtime) workload_type = getattr(args, "workload_type", "") workload_value = str(getattr(workload_type, "value", workload_type) or "") is_audio = workload_value.endswith("2a") diff --git a/fastvideo/entrypoints/openai/serving_engine.py b/fastvideo/entrypoints/openai/serving_engine.py index e01b1c99d..57171516b 100644 --- a/fastvideo/entrypoints/openai/serving_engine.py +++ b/fastvideo/entrypoints/openai/serving_engine.py @@ -8,6 +8,7 @@ from collections.abc import Awaitable, Callable from typing import Any, Protocol, TypeVar from fastvideo.api.schema import GenerationRequest +from fastvideo.entrypoints.openai.lora import LoraRuntime, LoraSelection from fastvideo.entrypoints.openai.protocol import VideoGenerationRequest _T = TypeVar("_T") @@ -35,8 +36,10 @@ class OpenAIServingEngine: def __init__(self, generator: ServingGenerator, - video_request_validator: Callable[[VideoGenerationRequest], None] | None = None) -> None: + video_request_validator: Callable[[VideoGenerationRequest], None] | None = None, + lora_runtime: LoraRuntime | None = None) -> None: self._generator = generator + self._lora_runtime = lora_runtime self._video_request_validator = video_request_validator self._generation_lock = asyncio.Lock() self._closed = False @@ -46,6 +49,10 @@ class OpenAIServingEngine: def generator(self) -> ServingGenerator: return self._generator + @property + def lora_runtime(self) -> LoraRuntime | None: + return self._lora_runtime + def validate_video_request(self, request: VideoGenerationRequest) -> None: if self._video_request_validator is not None: self._video_request_validator(request) @@ -77,18 +84,24 @@ class OpenAIServingEngine: request: GenerationRequest, *, on_start: Callable[[], Awaitable[None]] | None = None, + lora: LoraSelection | None = None, ) -> Any: """Generate one typed request without blocking the event loop.""" - return await self.run_serialized(self._generator.generate, request, on_start=on_start) + return await self.run_serialized(self._generator.generate, request, on_start=on_start, lora=lora) async def run_serialized( self, function: Callable[..., _T], *args: Any, on_start: Callable[[], Awaitable[None]] | None = None, + lora: LoraSelection | None = None, **kwargs: Any, ) -> _T: - """Run a synchronous pipeline operation under the serving lock.""" + """Run a synchronous pipeline operation under the serving lock. + + With a LoRA allow-list, the adapter change runs in the same locked call as ``function``; ``lora=None`` + renders the base model. + """ if self._closed: raise RuntimeError("FastVideo serving engine is shutting down") async with self._generation_lock: @@ -96,7 +109,8 @@ class OpenAIServingEngine: raise RuntimeError("FastVideo serving engine is shutting down") if on_start is not None: await on_start() - worker = asyncio.create_task(asyncio.to_thread(function, *args, **kwargs)) + call = function if self._lora_runtime is None else self._with_lora(function, lora) + worker = asyncio.create_task(asyncio.to_thread(call, *args, **kwargs)) try: return await asyncio.shield(worker) except asyncio.CancelledError: @@ -109,6 +123,22 @@ class OpenAIServingEngine: self._unhealthy_reason = str(error) raise + def _with_lora(self, function: Callable[..., _T], selection: LoraSelection | None) -> Callable[..., _T]: + runtime = self._lora_runtime + assert runtime is not None + + def call(*args: Any, **kwargs: Any) -> _T: + runtime.apply(self._generator, selection) + try: + return function(*args, **kwargs) + except (BrokenPipeError, EOFError): + raise + except Exception: + runtime.invalidate() + raise + + return call + async def run_async_serialized(self, function: Callable[[], Awaitable[_T]]) -> _T: """Run an async operation under the same pipeline lock.""" if self._closed: diff --git a/fastvideo/entrypoints/openai/video_api.py b/fastvideo/entrypoints/openai/video_api.py index 92f481819..0ec8099b7 100644 --- a/fastvideo/entrypoints/openai/video_api.py +++ b/fastvideo/entrypoints/openai/video_api.py @@ -18,6 +18,7 @@ from starlette.datastructures import UploadFile from fastvideo.api.compat import explicit_request_updates, request_to_sampling_param from fastvideo.api.schema import GenerationRequest +from fastvideo.entrypoints.openai.lora import LoraSelection from fastvideo.entrypoints.openai.protocol import ( VideoDeleteResponse, VideoGenerationRequest, @@ -213,6 +214,7 @@ async def _run_generation( request_id: str, generation_request: GenerationRequest, materialized_media: list[str] | None = None, + lora: LoraSelection | None = None, ) -> None: started = 0.0 video_path = generation_request.output.output_path @@ -229,7 +231,7 @@ async def _run_generation( ) try: - result = await get_serving_engine().generate(generation_request, on_start=mark_started) + result = await get_serving_engine().generate(generation_request, on_start=mark_started, lora=lora) if isinstance(result, list): if not result: raise RuntimeError("FastVideo returned no generation results") @@ -354,11 +356,13 @@ async def _parse_video_request(raw_request: Request) -> VideoGenerationRequest: raise HTTPException(status_code=400, detail=f"Invalid request body: {error}") from error -async def _adapt_request(request_id: str, request: VideoGenerationRequest) -> tuple[GenerationRequest, list[str]]: +async def _adapt_request(request_id: str, + request: VideoGenerationRequest) -> tuple[GenerationRequest, list[str], LoraSelection | None]: materialized_media: list[str] = [] try: get_serving_engine().validate_video_request(request) - validate_model_and_lora(request, get_server_args(), get_served_model_name()) + selection = validate_model_and_lora(request, get_server_args(), get_served_model_name(), + get_serving_engine().lora_runtime) materialized_media = await prepare_reference_media(request_id, request, get_output_dir(), @@ -371,8 +375,9 @@ async def _adapt_request(request_id: str, request: VideoGenerationRequest) -> tu served_model_name=get_served_model_name(), output_dir=get_output_dir(), default_request=get_default_request(), + lora_runtime=get_serving_engine().lora_runtime, ) - return generation_request, materialized_media + return generation_request, materialized_media, selection except Exception as error: discard_materialized_media(materialized_media) raise HTTPException(status_code=400, detail=str(error)) from error @@ -384,14 +389,14 @@ async def create_video(raw_request: Request) -> VideoResponse: """Create an asynchronous video generation job.""" request = await _parse_video_request(raw_request) request_id = f"video_gen_{generate_request_id()}" - generation_request, materialized_media = await _adapt_request(request_id, request) + generation_request, materialized_media, lora = await _adapt_request(request_id, request) try: job = _make_video_job(request_id, request, generation_request) await VIDEO_STORE.upsert(request_id, job) except Exception: discard_materialized_media(materialized_media) raise - task = asyncio.create_task(_run_generation(request_id, generation_request, materialized_media), + task = asyncio.create_task(_run_generation(request_id, generation_request, materialized_media, lora), name=f"video-job-{request_id}") task.add_done_callback(lambda _: discard_materialized_media(materialized_media)) _track_video_job(request_id, task) @@ -403,12 +408,12 @@ async def create_video_sync(raw_request: Request) -> FileResponse: """Generate synchronously and return the generated media with vLLM headers.""" request = await _parse_video_request(raw_request) request_id = f"video_sync-{generate_request_id()}" - generation_request, materialized_media = await _adapt_request(request_id, request) + generation_request, materialized_media, lora = await _adapt_request(request_id, request) started = time.perf_counter() response: FileResponse | None = None try: try: - result = await get_serving_engine().generate(generation_request) + result = await get_serving_engine().generate(generation_request, lora=lora) except Exception as error: logger.exception("Sync video generation failed for %s", request_id) raise HTTPException(status_code=500, detail=f"Video generation failed: {error}") from error diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 74a06e0a2..dc52f444c 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -1528,6 +1528,10 @@ class VideoGenerator: accumulate: bool = False) -> None: self.executor.set_lora_adapter(lora_nickname, lora_path, strength=strength, accumulate=accumulate) + def clear_lora_adapter(self) -> None: + """Return the transformer to its base weights (plus any constructor adapter), without reading an adapter.""" + self.executor.clear_lora_adapter() + def unmerge_lora_weights(self) -> None: """ Use unmerged weights for inference to produce videos that align with diff --git a/fastvideo/layers/lora/linear.py b/fastvideo/layers/lora/linear.py index 31b21d531..a65ef28c0 100644 --- a/fastvideo/layers/lora/linear.py +++ b/fastvideo/layers/lora/linear.py @@ -303,9 +303,10 @@ class BaseLayerWithLoRA(nn.Module): raise RuntimeError("Training-mode LoRA layers do not retain an inference unmerge snapshot") if isinstance(self.base_layer.weight, DTensor): device = self.base_layer.weight.data.device - self.base_layer.weight = nn.Parameter(self.cpu_weight.to(device, non_blocking=True)) + self.base_layer.weight = nn.Parameter(self.cpu_weight.to(device, non_blocking=True, copy=True)) else: - self.base_layer.weight.data = self.cpu_weight.data.to(self.base_layer.weight, non_blocking=True) + # copy=True: on the snapshot's own device .to() would alias it, and the next merge would write into it. + self.base_layer.weight.data = self.cpu_weight.data.to(self.base_layer.weight, non_blocking=True, copy=True) self.merged = False diff --git a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py index 4ad6c466c..5a6e38c19 100644 --- a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -295,6 +295,10 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): super().unmerge_lora_weights() self._drop_pinned_host_copies() + def clear_lora_adapter(self) -> None: + super().clear_lora_adapter() + self._drop_pinned_host_copies() + def __init__(self, *args: Any, **kwargs: Any) -> None: self._ref2va = getattr(self, "_ref2va_default", False) self._denoise_stages_ready = False diff --git a/fastvideo/pipelines/lora_pipeline.py b/fastvideo/pipelines/lora_pipeline.py index 172969d07..e482ae651 100644 --- a/fastvideo/pipelines/lora_pipeline.py +++ b/fastvideo/pipelines/lora_pipeline.py @@ -41,6 +41,10 @@ def _merge_mode_modules(pipeline: Any, transformer_modules: dict[str, nn.Module] """ lora_layers = getattr(pipeline, "lora_layers", None) or {} for transformer_name, transformer in transformer_modules.items(): + if is_lazy_module(transformer) and not transformer.is_materialized: + # Nothing resident to rewrite, and touching the proxy would load it, or reload it from inside its own + # materialize transform. A reload re-applies the adapter to fresh layers. + continue transformer_layers = lora_layers.get(transformer_name) if transformer_layers is None: yield from transformer.modules() @@ -221,9 +225,15 @@ class LoRAPipeline(ComposedPipelineBase): lora_min_match_fraction: float = 0.5 _constructor_dense_lora_path: str | None = None _setting_constructor_adapter: bool = False + # The adapter and strength last set at runtime; a released lazy transformer gets it back when it loads again. + _runtime_adapter: tuple[str, float] | None = None # Sequential load (the 32 GB and other cards that cannot hold the text encoder and the DiT at once) omits # the transformer from the first load. The adapter is applied when that module arrives. _constructor_adapter_pending: bool = False + # Set once inference LoRA state exists, so ``add_module`` wraps a transformer that arrives later. + _tracks_late_transformers: bool = False + # Adapters loaded while the transformer was not in memory; their coverage is checked when it loads. + _coverage_deferred: set[str] def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) @@ -239,6 +249,7 @@ class LoRAPipeline(ComposedPipelineBase): self.cur_adapter_name = "" self.cur_adapter_path = "" self.cur_adapter_strength = 1.0 + self._coverage_deferred = set() # build list of trainable transformers for transformer_name in self.trainable_transformer_names: if (transformer_name in self.modules and self.modules[transformer_name] is not None): @@ -316,25 +327,35 @@ class LoRAPipeline(ComposedPipelineBase): self.convert_to_lora_layers() if not any(is_lazy_module(module) for module in self.trainable_transformer_modules.values()): self._apply_constructor_adapter() + self._tracks_late_transformers = not self.training_mode def add_module(self, module_name: str, module: Any) -> None: super().add_module(module_name, module) - if not getattr(self, "_constructor_adapter_pending", False): - return if module is None or module_name not in self.trainable_transformer_names: return - self.trainable_transformer_modules[module_name] = module - if is_lazy_module(module): - if not self.lora_initialized: - self.convert_to_lora_layers() + if not (self._tracks_late_transformers or self._constructor_adapter_pending): return - if self.lora_initialized: - self._convert_one_transformer(module_name, module) - else: + if self.trainable_transformer_modules.get(module_name) is module: + return + # Sequential load brings the DiT in after construction, and an adapter may already have been set. + self.trainable_transformer_modules[module_name] = module + if not self.lora_initialized and not self._constructor_adapter_pending: + # The first ``set_lora_adapter`` wraps it. + return + if not self.lora_initialized: self.convert_to_lora_layers() + elif is_lazy_module(module): + self._defer_lora_until_loaded(module_name, module) + else: + self._convert_one_transformer(module_name, module) + if is_lazy_module(module): + # Its materialize transform applies the adapters when it loads. + return self._apply_constructor_adapter() - self._constructor_adapter_pending = False - logger.info("LoRA %s attached to %s after it loaded", self.lora_path, module_name) + self._reapply_runtime_adapter() + if self._constructor_adapter_pending: + self._constructor_adapter_pending = False + logger.info("LoRA %s attached to %s after it loaded", self.lora_path, module_name) def is_target_layer(self, module_name: str) -> bool: if self.lora_target_modules is None: @@ -351,21 +372,22 @@ class LoRAPipeline(ComposedPipelineBase): del lora_path return state_dict - def _check_adapter_coverage(self, lora_nickname: str, lora_path: str | None) -> None: + def _check_adapter_coverage(self, lora_nickname: str, lora_path: str | None) -> bool: """Fail before merging when an adapter reaches no layer, or too few of them. Without this the adapter loads, generation runs, and the result is an unadapted model with a - warning per layer buried in the log. + warning per layer buried in the log. Returns False when the check must wait for the transformer to load. """ adapter = self.lora_adapters[lora_nickname] if not adapter and lora_path is not None and lora_path == self._constructor_dense_lora_path: # A dense-only adapter (.diff/.set_weight) has no low-rank layers; DenseLoRAPatch applied it at load. - return - released = sorted(set(self.trainable_transformer_modules) - set(self.lora_layers)) + return True + released = sorted((set(self.trainable_transformer_modules) | set(self.trainable_transformer_names)) - + set(self.lora_layers)) if released: - # A released lazy transformer has no LoRA layers until it loads again, so nothing can be matched yet. - logger.info("Skipping the LoRA coverage check for %r: %s not loaded", lora_path, ", ".join(released)) - return + # A released or not yet loaded transformer has no LoRA layers, so nothing can be matched yet. + logger.info("Deferring the LoRA coverage check for %r: %s not loaded", lora_path, ", ".join(released)) + return False lora_a = {key[:-len(".lora_A")] for key in adapter if key.endswith(".lora_A")} lora_b = {key[:-len(".lora_B")] for key in adapter if key.endswith(".lora_B")} # A module with only one factor is never applied, so it counts as unmatched, not as covered. @@ -380,7 +402,7 @@ class LoRAPipeline(ComposedPipelineBase): matched = len(adapter_modules) - len(unmatched) fraction = matched / len(adapter_modules) if adapter_modules else 0.0 if fraction >= self.lora_min_match_fraction and matched > 0: - return + return True shown = unmatched[:20] more = f" (+{len(unmatched) - len(shown)} more)" if len(unmatched) > len(shown) else "" raise ValueError( @@ -452,6 +474,42 @@ class LoRAPipeline(ComposedPipelineBase): finally: self._setting_constructor_adapter = False + def _reapply_runtime_adapter(self) -> None: + """Apply the runtime adapter, from memory, to the fresh LoRA layers of a reloaded lazy transformer.""" + if self._runtime_adapter is None: + return + name, strength = self._runtime_adapter + if name not in self.lora_adapters: + return + if self.lora_strict_coverage and name in self._coverage_deferred: + path = self.lora_adapter_paths.get(name) + try: + checked = self._check_adapter_coverage(name, path) + except ValueError: + # Drop it so a later request reloads and checks it again instead of silently running without it. + self.lora_adapters.pop(name, None) + self.lora_adapter_paths.pop(name, None) + self._coverage_deferred.discard(name) + self._runtime_adapter = None + raise + if checked: + self._coverage_deferred.discard(name) + self.cur_adapter_name = "" + self.cur_adapter_path = "" + self.set_lora_adapter(name, strength=strength) + + def _transformer_name_mappings(self) -> tuple[dict, dict]: + """The transformer's checkpoint and LoRA key mappings, read without loading it. + + A deferred transformer, or one sequential load has not brought in yet, is read from the DiT arch config, + which the model class mirrors, so loading an adapter never loads the transformer early. + """ + transformer = self.modules.get("transformer") + if transformer is not None and not (is_lazy_module(transformer) and not transformer.is_materialized): + return transformer.param_names_mapping, transformer.lora_param_names_mapping + arch_config = self.fastvideo_args.pipeline_config.dit_config.arch_config + return arch_config.param_names_mapping, arch_config.lora_param_names_mapping + def _convert_one_transformer(self, transformer_name: str, transformer_module: nn.Module) -> None: excluded_lora_layers = self._exclude_lora_layers_for(transformer_name, transformer_module) # Fresh instance after a lazy rematerialize must not keep the previous @@ -505,6 +563,28 @@ class LoRAPipeline(ComposedPipelineBase): converted_count += 1 logger.info("Converted %d layers to LoRA layers", converted_count) + def _defer_lora_until_loaded(self, transformer_name: str, transformer_module: Any) -> None: + """Wrap a lazy transformer and apply the adapters each time it loads; drop its layers when it is released.""" + + def _drop_lora_refs(*, _name: str = transformer_name) -> None: + self.lora_layers.pop(_name, None) + self.cur_adapter_name = "" + self.cur_adapter_path = "" + + def _lora_after_load(module: nn.Module, *, _name: str = transformer_name) -> nn.Module: + try: + self._convert_one_transformer(_name, module) + self._apply_constructor_adapter() + self._reapply_runtime_adapter() + except Exception: + # A failed materialize runs no release callback, and these layers would keep the transformer alive. + _drop_lora_refs() + raise + return module + + transformer_module.add_release_callback(_drop_lora_refs) + transformer_module.set_materialize_transform(_lora_after_load) + def convert_to_lora_layers(self) -> None: """ Unified method to convert the transformer to a LoRA transformer. @@ -517,19 +597,7 @@ class LoRAPipeline(ComposedPipelineBase): transformer_module, ) in self.trainable_transformer_modules.items(): if is_lazy_module(transformer_module): - - def _drop_lora_refs(*, _name: str = transformer_name) -> None: - self.lora_layers.pop(_name, None) - self.cur_adapter_name = "" - self.cur_adapter_path = "" - - def _lora_after_load(module: nn.Module, *, _name: str = transformer_name) -> nn.Module: - self._convert_one_transformer(_name, module) - self._apply_constructor_adapter() - return module - - transformer_module.add_release_callback(_drop_lora_refs) - transformer_module.set_materialize_transform(_lora_after_load) + self._defer_lora_until_loaded(transformer_name, transformer_module) continue self._convert_one_transformer(transformer_name, transformer_module) @@ -591,9 +659,9 @@ class LoRAPipeline(ComposedPipelineBase): lora_state_dict = self.convert_adapter_state_dict(load_file(lora_local_path), lora_path) # Map the hf layer names to our custom layer names - param_names_mapping_fn = get_param_names_mapping(self.modules["transformer"].param_names_mapping) - lora_param_names_mapping_fn = get_param_names_mapping( - self.modules["transformer"].lora_param_names_mapping) + param_names_mapping, lora_param_names_mapping = self._transformer_name_mappings() + param_names_mapping_fn = get_param_names_mapping(param_names_mapping) + lora_param_names_mapping_fn = get_param_names_mapping(lora_param_names_mapping) # Extract alpha values and weights in a single pass to_merge_params: defaultdict[Hashable, dict[Any, Any]] = (defaultdict(dict)) @@ -634,8 +702,10 @@ class LoRAPipeline(ComposedPipelineBase): if target_name in self.lora_adapters[lora_nickname]: raise ValueError(f"Target name {target_name} already exists in lora_adapters[{lora_nickname}]") self.lora_adapters[lora_nickname][target_name] = weight.to(self.device) - if self.lora_strict_coverage: - self._check_adapter_coverage(lora_nickname, lora_path) + if self.lora_strict_coverage and not self._check_adapter_coverage(lora_nickname, lora_path): + self._coverage_deferred.add(lora_nickname) + else: + self._coverage_deferred.discard(lora_nickname) except Exception: if previous_adapter is None: self.lora_adapters.pop(lora_nickname, None) @@ -650,8 +720,8 @@ class LoRAPipeline(ComposedPipelineBase): if (not adapter_updated and self.cur_adapter_name == lora_nickname and self.cur_adapter_strength == strength and not accumulate): return - self.cur_adapter_name = lora_nickname - self.cur_adapter_strength = strength + # Recorded only once every layer holds it, so retrying after a failure part-way through changes them again. + self.cur_adapter_name = "" # Merge the new adapter adapted_count = 0 @@ -698,11 +768,16 @@ class LoRAPipeline(ComposedPipelineBase): consumed.update((lora_A_name, lora_B_name, lora_alpha_name)) else: untouched.append(name) + if not accumulate: + # An earlier adapter may be merged here; disabling alone would keep its delta. + layer.unmerge_lora_weights() layer.disable_lora = True # Repack inside the offload scope, while the block's parameters are resident. repack_scope = [module] if module is not None else list(layers.values()) _convert_quantized_weights_after_lora_change(repack_scope, skip=_all_side_branch(layers.values(), repack_scope)) + self.cur_adapter_name = lora_nickname + self.cur_adapter_strength = strength if untouched and rank == 0: logger.warning( "LoRA adapter %s has no weights for %d wrapped layers; LoRA is disabled on them (e.g. %s).", @@ -727,6 +802,8 @@ class LoRAPipeline(ComposedPipelineBase): lora_path) if unmatched: logger.warning("LoRA adapter %s: %d weights did not reach a layer", lora_path, len(unmatched)) + if not self._setting_constructor_adapter and not accumulate: + self._runtime_adapter = (lora_nickname, strength) def merge_lora_weights(self) -> None: """Merge LoRA weights, then repack the packed quantized buffers for the merged state. @@ -755,6 +832,30 @@ class LoRAPipeline(ComposedPipelineBase): _convert_quantized_weights_after_lora_change(repack_scope, skip=_all_side_branch(layers.values(), repack_scope)) + def clear_lora_adapter(self) -> None: + """Return every LoRA layer to its base weights without reading an adapter, then restore the constructor adapter. + + Merged layers get back the weights snapshotted when they were wrapped, and side-branch layers drop their + factors, so the result does not depend on which adapter, if any, the layers hold or how a change failed. + Loaded adapters stay in memory for a later ``set_lora_adapter``. + """ + for transformer_lora_layers in self.lora_layers.values(): + for module, layers in transformer_lora_layers.lora_layers_by_block(): + with _get_hook_ctx(module): + for layer in layers.values(): + layer.disable_lora = False + layer.unmerge_lora_weights() + layer.lora_A = None + layer.lora_B = None + # Repack inside the offload scope, while the block's parameters are resident. + repack_scope = [module] if module is not None else list(layers.values()) + _convert_quantized_weights_after_lora_change(repack_scope, + skip=_all_side_branch(layers.values(), repack_scope)) + self.cur_adapter_name = "" + self.cur_adapter_path = "" + self._runtime_adapter = None + self._apply_constructor_adapter() + def unmerge_lora_weights(self) -> None: """Unmerge LoRA weights, then repack the packed quantized buffers for the unmerged state. diff --git a/fastvideo/tests/api/test_cli_translation.py b/fastvideo/tests/api/test_cli_translation.py index 6226ff654..d1ba803be 100644 --- a/fastvideo/tests/api/test_cli_translation.py +++ b/fastvideo/tests/api/test_cli_translation.py @@ -502,6 +502,54 @@ def test_serve_subcommand_dispatches_via_typed_config(tmp_path, monkeypatch): assert captured["output_dir"] == "serve-outputs/" +def test_serve_subcommand_forwards_lora_allow_list(tmp_path, monkeypatch): + config_path = tmp_path / "serve-lora.yaml" + config_path.write_text( + "generator:\n" + " model_path: serve-model\n" + "lora:\n" + " adapters:\n" + " - id: realism\n" + " path: org/realism-lora\n", + encoding="utf-8", + ) + args, _ = _parse_serve_args(["--config", str(config_path)]) + captured: dict[str, object] = {} + + def fake_run_server(fastvideo_args, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr( + "fastvideo.entrypoints.cli.serve.generator_config_to_fastvideo_args", + lambda config: SimpleNamespace(model_path=config.model_path), + ) + monkeypatch.setattr(api_server, "run_server", fake_run_server) + monkeypatch.setattr("fastvideo.entrypoints.openai.lora.check_adapter_file", lambda adapter: None) + + with envs.FASTVIDEO_LORA_SIDE_BRANCH.override(True): + ServeSubcommand().cmd(args) + + assert list(captured["lora_runtime"].adapters) == ["realism"] + + +def test_serve_cmd_refuses_allow_list_with_startup_adapter_override(tmp_path): + config_path = tmp_path / "serve-lora-conflict.yaml" + config_path.write_text( + "generator:\n" + " model_path: serve-model\n" + "lora:\n" + " adapters:\n" + " - id: realism\n" + " path: org/realism-lora\n", + encoding="utf-8", + ) + args, _ = _parse_serve_args( + ["--config", str(config_path), "--generator.pipeline.components.lora_path", "/x/adapter.safetensors"]) + + with pytest.raises(ValueError, match="startup adapter"): + ServeSubcommand().cmd(args) + + def test_serve_subcommand_forwards_default_request(tmp_path, monkeypatch): config_path = tmp_path / "serve-default-request.yaml" config_path.write_text( diff --git a/fastvideo/tests/entrypoints/test_openai_lora.py b/fastvideo/tests/entrypoints/test_openai_lora.py new file mode 100644 index 000000000..2ba5802bc --- /dev/null +++ b/fastvideo/tests/entrypoints/test_openai_lora.py @@ -0,0 +1,608 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-request community LoRA on the OpenAI video API (CPU only, fake generator).""" +from __future__ import annotations + +import json +import time +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +from starlette.testclient import TestClient + +from fastvideo import envs +from fastvideo.api.parser import parse_config +from fastvideo.api.schema import ( + ComponentConfig, + EngineConfig, + GeneratorConfig, + LoRAAdapterConfig, + PipelineSelection, + QuantizationConfig, + ServeConfig, + ServeLoRAConfig, +) +from fastvideo.entrypoints.openai.api_server import create_app +from fastvideo.entrypoints.openai import lora as lora_module +from fastvideo.entrypoints.openai.lora import LoraRuntime, adapters_attach_as_side_branch, build_lora_runtime +from fastvideo.entrypoints.openai.stores import VIDEO_STORE +from fastvideo.layers.quantization.nvfp4_config import H3_NVFP4_DIT_EXPORT_FILENAME + +ARTIFACT = b"fake-video" +REALISM_PATH = "fal/MiniMax-H3-Realism-People-LoRA" + + +class _FakeGenerator: + """Records each adapter change, clear and generation, in the order they arrive.""" + + def __init__(self, failing: set[tuple[str, float | None]] | None = None, clear_fails: bool = False) -> None: + self.failing = failing or set() + self.clear_fails = clear_fails + self.held: set[str] = set() + self.events: list[tuple[Any, ...]] = [] + + def clear_lora_adapter(self) -> None: + self.events.append(("clear",)) + if self.clear_fails: + raise RuntimeError("cannot clear") + + def set_lora_adapter(self, nickname: str, path: str | None = None, strength: float = 1.0) -> None: + self.events.append(("set", nickname, path, strength)) + if path is None and nickname not in self.held: + raise ValueError(f"Adapter {nickname} not found in the pipeline.") + if (nickname, strength) in self.failing or (nickname, None) in self.failing: + raise RuntimeError(f"cannot load {nickname}") + self.held.add(nickname) + + def generate(self, request: Any) -> SimpleNamespace: + self.events.append(("generate", request.prompt)) + output = Path(request.output.output_path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_bytes(ARTIFACT) + return SimpleNamespace(video_path=str(output), generation_time=0.01) + + def shutdown(self) -> None: + pass + + +def _adapters() -> ServeLoRAConfig: + return ServeLoRAConfig(adapters=[ + LoRAAdapterConfig(id="realism", path=REALISM_PATH, default_strength=0.8, max_strength=1.5, + description="People realism", license="Apache-2.0", source_url="https://example.test/r"), + LoRAAdapterConfig(id="orbit", path="pablodawson/Orbit", default_strength=1.0, min_strength=0.2, + max_strength=1.5), + ]) + + +def _check(adapter: LoRAAdapterConfig, unusable: frozenset[str]) -> None: + if adapter.id in unusable: + raise ValueError(f"{adapter.id} is missing") + + +def _args(lora_path: str | None = None) -> SimpleNamespace: + return SimpleNamespace(model_path="m", lora_path=lora_path, lora_nickname="startup", lora_strength=0.5, + override_pipeline_cls_name=None) + + +@pytest.fixture +def make_client(tmp_path: Path): + clients: list[TestClient] = [] + + def make(generator: _FakeGenerator, *, allow_list: bool = True, lora_path: str | None = None, + unusable: frozenset[str] = frozenset()) -> TestClient: + runtime = LoraRuntime(_adapters(), check_adapter=lambda adapter: _check(adapter, unusable)) if allow_list else None + app = create_app(_args(lora_path), str(tmp_path / "out"), generator_factory=lambda: generator, + lora_runtime=runtime) + client = TestClient(app) + client.__enter__() + clients.append(client) + return client + + yield make + for client in clients: + client.__exit__(None, None, None) + import asyncio + asyncio.run(VIDEO_STORE.clear()) + + +def _sync(client: TestClient, prompt: str = "a fox", **lora: Any): + body: dict[str, Any] = {"prompt": prompt} + if lora: + body["lora"] = lora + return client.post("/v1/videos/sync", json=body) + + +def _job(client: TestClient, prompt: str = "a fox", **lora: Any) -> dict[str, Any]: + body: dict[str, Any] = {"prompt": prompt} + if lora: + body["lora"] = lora + response = client.post("/v1/videos", json=body) + assert response.status_code == 200, response.text + video_id = response.json()["id"] + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + job = client.get(f"/v1/videos/{video_id}").json() + if job["status"] in ("completed", "failed"): + return job + time.sleep(0.02) + raise AssertionError("job did not finish") + + + + +def _serve_raw(**extra: Any) -> dict[str, Any]: + raw: dict[str, Any] = {"generator": {"model_path": "/models/h3"}} + raw.update(extra) + return raw + + +def test_allow_list_parses_with_defaults_and_metadata() -> None: + config = parse_config(ServeConfig, _serve_raw(lora={"adapters": [ + {"id": "realism", "path": REALISM_PATH, "license": "Apache-2.0"}, + ]})) + + assert config.lora is not None + adapter = config.lora.adapters[0] + assert (adapter.default_strength, adapter.min_strength, adapter.max_strength) == (0.8, 0.0, 2.0) + assert adapter.license == "Apache-2.0" + + +def test_serve_config_has_no_allow_list_by_default() -> None: + assert parse_config(ServeConfig, _serve_raw()).lora is None + + +@pytest.mark.parametrize( + ("adapter", "message"), + [ + ({"id": " ", "path": "p"}, "must not be empty"), + ({"id": "a", "path": ""}, "empty path"), + ({"id": "a", "path": "p", "default_strength": 3.0}, "min_strength <= default_strength"), + ({"id": "a", "path": "p", "min_strength": 1.0, "default_strength": 0.8}, "min_strength <= default_strength"), + ({"id": "a", "path": "p", "max_strength": float("inf")}, "non-finite"), + ], +) +def test_allow_list_rejects_invalid_adapter(adapter: dict[str, Any], message: str) -> None: + with pytest.raises(Exception, match=message): + parse_config(ServeConfig, _serve_raw(lora={"adapters": [adapter]})) + + +def test_allow_list_rejects_duplicate_ids() -> None: + adapters = [{"id": "a", "path": "p"}, {"id": "a", "path": "q"}] + with pytest.raises(Exception, match="repeats id"): + parse_config(ServeConfig, _serve_raw(lora={"adapters": adapters})) + + +def test_allow_list_with_startup_adapter_is_refused() -> None: + raw = _serve_raw(lora={"adapters": [{"id": "a", "path": "p"}]}) + raw["generator"]["pipeline"] = {"components": {"lora_path": "/x/adapter.safetensors"}} + + with pytest.raises(Exception, match="startup adapter"): + parse_config(ServeConfig, raw) + + +def test_startup_adapter_without_allow_list_still_parses() -> None: + raw = _serve_raw() + raw["generator"]["pipeline"] = {"components": {"lora_path": "/x/adapter.safetensors"}} + + assert parse_config(ServeConfig, raw).generator.pipeline.components.lora_path == "/x/adapter.safetensors" + + +@pytest.mark.parametrize( + ("quant", "packed", "side_branch_env", "expected"), + [("NVFP4", True, False, True), ("NVFP4", False, False, False), (None, True, False, False), + ("FP8", True, False, False), (None, False, True, True)], +) +def test_runtime_switching_matches_side_branch_attachment(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, + quant: str | None, packed: bool, side_branch_env: bool, + expected: bool) -> None: + monkeypatch.setattr(lora_module, "get_model_info", lambda *args, **kwargs: SimpleNamespace(pipeline_cls=object)) + transformer = tmp_path / "transformer" + transformer.mkdir() + if packed: + (transformer / H3_NVFP4_DIT_EXPORT_FILENAME).touch() + quantization = None if quant is None else QuantizationConfig(transformer_quant=quant) + config = GeneratorConfig(model_path=str(tmp_path), engine=EngineConfig(quantization=quantization)) + + with envs.FASTVIDEO_LORA_SIDE_BRANCH.override(side_branch_env): + assert adapters_attach_as_side_branch(config) is expected + + + + +@pytest.mark.parametrize("route", ["/v1/videos", "/v1/videos/sync"]) +@pytest.mark.parametrize( + ("lora", "message"), + [ + ({"name": "nope"}, "Unknown LoRA adapter"), + ({"name": "realism", "scale": 2.0}, "takes a scale from"), + ({"name": "orbit", "scale": 0.0}, "takes a scale from"), + ({"name": "realism", "scale": "high"}, "Invalid LoRA scale"), + ({"name": "realism", "path": "/etc/passwd"}, "not accepted"), + ({"path": "/etc/passwd"}, "not accepted"), + ({"lora_path": REALISM_PATH, "name": "realism"}, "not accepted"), + ({"scale": 0.5}, "must provide a name"), + ], +) +def test_bad_selection_is_400_before_any_job_is_stored(make_client, route: str, lora: dict[str, Any], + message: str) -> None: + generator = _FakeGenerator() + client = make_client(generator) + + response = client.post(route, json={"prompt": "a fox", "lora": lora}) + + assert response.status_code == 400 + assert message in response.json()["error"]["message"] + assert generator.events == [] + assert client.get("/v1/videos").json()["data"] == [] + + +def test_selector_aliases_and_default_strength(make_client) -> None: + generator = _FakeGenerator() + client = make_client(generator) + + assert _sync(client, lora_name="realism").status_code == 200 + assert _sync(client, adapter="orbit", lora_scale=0.5).status_code == 200 + + assert [event for event in generator.events if event[0] == "set"] == [ + ("set", "realism", REALISM_PATH, 0.8), + ("set", "orbit", "pablodawson/Orbit", 0.5), + ] + + +def test_selector_without_allow_list_keeps_startup_behaviour(make_client) -> None: + generator = _FakeGenerator() + client = make_client(generator, allow_list=False) + + response = _sync(client, name="realism") + + assert response.status_code == 400 + assert "no startup LoRA" in response.json()["error"]["message"] + assert _sync(client).status_code == 200 + assert [event[0] for event in generator.events] == ["generate"] + + +def test_startup_adapter_selector_without_allow_list_is_unchanged(make_client) -> None: + client = make_client(_FakeGenerator(), allow_list=False, lora_path="/x/adapter.safetensors") + + assert _sync(client, name="startup", scale=0.5).status_code == 200 + other = _sync(client, name="realism") + assert other.status_code == 400 + assert "not the startup adapter" in other.json()["error"]["message"] + + + + +def test_sync_applies_adapter_before_generate(make_client) -> None: + generator = _FakeGenerator() + client = make_client(generator) + + response = _sync(client, "first", name="realism", scale=0.8) + + assert response.status_code == 200 + assert response.content == ARTIFACT + assert generator.events == [("set", "realism", REALISM_PATH, 0.8), ("generate", "first")] + + +def test_job_applies_adapter_before_generate(make_client) -> None: + generator = _FakeGenerator() + client = make_client(generator) + + job = _job(client, "first", name="realism", scale=0.8) + + assert job["status"] == "completed" + assert generator.events == [("set", "realism", REALISM_PATH, 0.8), ("generate", "first")] + + +@pytest.mark.parametrize("run", [_sync, _job], ids=["sync", "job"]) +def test_request_without_lora_restores_base_after_an_adapter(make_client, run) -> None: + generator = _FakeGenerator() + client = make_client(generator) + + run(client, "styled", name="realism", scale=0.8) + run(client, "plain") + run(client, "plain again") + + assert generator.events == [ + ("set", "realism", REALISM_PATH, 0.8), + ("generate", "styled"), + ("set", "realism", REALISM_PATH, 0.0), + ("generate", "plain"), + ("generate", "plain again"), + ] + + +def test_base_only_server_never_loads_an_adapter(make_client) -> None: + generator = _FakeGenerator() + client = make_client(generator) + + _sync(client) + _sync(client) + + assert generator.events == [("generate", "a fox"), ("generate", "a fox")] + + +def test_same_selection_is_not_reapplied(make_client) -> None: + generator = _FakeGenerator() + client = make_client(generator) + + _sync(client, name="realism", scale=0.8) + _sync(client, name="realism", scale=0.8) + + assert [event[0] for event in generator.events] == ["set", "generate", "generate"] + + +def test_scale_zero_sets_the_adapter_to_zero(make_client) -> None: + generator = _FakeGenerator() + client = make_client(generator) + + assert _sync(client, name="realism", scale=0).status_code == 200 + + assert generator.events[0] == ("set", "realism", REALISM_PATH, 0.0) + + +def test_failed_change_fails_the_request_and_does_not_generate(make_client) -> None: + generator = _FakeGenerator(failing={("realism", None)}) + client = make_client(generator) + + sync = _sync(client, name="realism", scale=0.8) + job = _job(client, name="realism", scale=0.8) + + assert sync.status_code == 500 + assert "could not be applied on this server" in sync.json()["error"]["message"] + assert "cannot load" not in sync.json()["error"]["message"] + assert job["status"] == "failed" + assert all(event[0] != "generate" for event in generator.events) + + +def test_failed_change_never_leaves_an_unrequested_adapter(make_client) -> None: + generator = _FakeGenerator(failing={("orbit", 1.0)}) + client = make_client(generator) + + _sync(client, "styled", name="realism", scale=0.8) + failed = _sync(client, "orbit", name="orbit") + plain = _sync(client, "plain") + + assert failed.status_code == 500 + assert plain.status_code == 200 + # The failed change clears every layer back to the base model, so "plain" needs no further change. + assert generator.events[-3:] == [("set", "orbit", "pablodawson/Orbit", 1.0), ("clear",), ("generate", "plain")] + + +def test_failed_pre_load_check_leaves_the_base_model_in_place(make_client) -> None: + generator = _FakeGenerator() + client = make_client(generator, unusable=frozenset({"orbit"})) + + failed = _sync(client, "orbit", name="orbit") + plain = _sync(client, "plain") + + assert failed.status_code == 500 + assert "cannot be loaded on this server" in failed.json()["error"]["message"] + assert plain.status_code == 200 + assert generator.events == [("generate", "plain")] + + +def test_failed_first_load_leaves_the_base_model_serving(make_client) -> None: + generator = _FakeGenerator(failing={("realism", None)}) + client = make_client(generator) + + failed = _sync(client, name="realism") + plain = _sync(client, "plain") + again = _sync(client, "plain again") + + assert failed.status_code == 500 + assert plain.status_code == 200 and again.status_code == 200 + assert generator.events == [("set", "realism", REALISM_PATH, 0.8), ("clear",), ("generate", "plain"), + ("generate", "plain again")] + + +def test_a_failed_change_whose_clear_also_fails_refuses_until_a_clear_succeeds(make_client) -> None: + generator = _FakeGenerator(failing={("orbit", None)}, clear_fails=True) + client = make_client(generator) + assert _sync(client, "styled", name="realism").status_code == 200 + + assert _sync(client, name="orbit").status_code == 500 + refused = _sync(client, "plain") + generator.clear_fails = False + recovered = _sync(client, "plain again") + + assert refused.status_code == 500 + assert "clearing the adapter" in refused.json()["error"]["message"] + assert ("generate", "plain") not in generator.events + assert recovered.status_code == 200 + assert generator.events[-2:] == [("clear",), ("generate", "plain again")] + + +def test_a_retried_selection_after_an_unknown_state_clears_before_it_changes_the_layers(make_client) -> None: + generator = _FakeGenerator(failing={("orbit", None)}, clear_fails=True) + client = make_client(generator) + assert _sync(client, name="orbit").status_code == 500 + + refused = _sync(client, name="orbit", scale=1.0) + generator.failing.clear() + generator.clear_fails = False + retried = _sync(client, "styled", name="orbit", scale=1.0) + + orbit = ("set", "orbit", "pablodawson/Orbit", 1.0) + assert refused.status_code == 500 + assert retried.status_code == 200 + assert generator.events == [orbit, ("clear",), ("clear",), ("clear",), orbit, ("generate", "styled")] + + +def test_a_failed_generation_with_an_adapter_clears_before_the_next_request(make_client) -> None: + generator = _FakeGenerator() + client = make_client(generator) + generate = generator.generate + + def fail_once(request: Any) -> SimpleNamespace: + generator.generate = generate # type: ignore[method-assign] + raise RuntimeError("adapter reaches 0 of 1 modules") + + generator.generate = fail_once # type: ignore[method-assign] + assert _sync(client, "styled", name="realism").status_code == 500 + assert _sync(client, "styled again", name="realism").status_code == 200 + + realism = ("set", "realism", REALISM_PATH, 0.8) + assert generator.events == [realism, ("clear",), realism, ("generate", "styled again")] + + +def test_a_selection_after_a_failed_change_applies_normally(make_client) -> None: + generator = _FakeGenerator(failing={("realism", 0.8)}) + client = make_client(generator) + + assert _sync(client, name="realism", scale=0.5).status_code == 200 + assert _sync(client, name="realism", scale=0.8).status_code == 500 + assert _sync(client, name="realism", scale=0.5).status_code == 200 + + assert generator.events[-4:] == [("set", "realism", REALISM_PATH, 0.8), ("clear",), + ("set", "realism", REALISM_PATH, 0.5), ("generate", "a fox")] + + +def test_model_info_lists_adapters_without_paths(make_client) -> None: + client = make_client(_FakeGenerator()) + + body = client.get("/v1/model_info").json() + + assert body["lora"] == { + "supported": True, + "runtime_switching": True, + "base_strength": 0.0, + "adapters": [ + {"id": "realism", "default_strength": 0.8, "min_strength": 0.0, "max_strength": 1.5, + "description": "People realism", "license": "Apache-2.0", "source_url": "https://example.test/r"}, + {"id": "orbit", "default_strength": 1.0, "min_strength": 0.2, "max_strength": 1.5}, + ], + } + assert REALISM_PATH not in json.dumps(body) + + +def test_model_info_without_allow_list_keeps_startup_info(make_client) -> None: + assert make_client(_FakeGenerator(), allow_list=False).get("/v1/model_info").json()["lora"] is None + + +def test_model_info_startup_adapter_is_unchanged(make_client) -> None: + client = make_client(_FakeGenerator(), allow_list=False, lora_path="/x/adapter.safetensors") + + assert client.get("/v1/model_info").json()["lora"] == {"name": "startup", "path": "/x/adapter.safetensors", + "scale": 0.5} + + +def _serve_config(tmp_path: Path, *, quant: str | None = "NVFP4", packed: bool = True) -> ServeConfig: + transformer = tmp_path / "model" / "transformer" + transformer.mkdir(parents=True) + if packed: + (transformer / H3_NVFP4_DIT_EXPORT_FILENAME).touch() + quantization = None if quant is None else QuantizationConfig(transformer_quant=quant) + generator = GeneratorConfig(model_path=str(tmp_path / "model"), engine=EngineConfig(quantization=quantization)) + return ServeConfig(generator=generator, lora=_adapters()) + + +def test_build_lora_runtime_checks_every_adapter(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + checked: list[str] = [] + monkeypatch.setattr(lora_module, "check_adapter_file", lambda adapter: checked.append(adapter.id)) + + assert build_lora_runtime(_serve_config(tmp_path)) is not None + assert checked == ["realism", "orbit"] + + +def test_build_lora_runtime_refuses_merged_lora_layers(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(lora_module, "check_adapter_file", lambda adapter: None) + + with envs.FASTVIDEO_LORA_SIDE_BRANCH.override(False), pytest.raises(ValueError, match="side-branch"): + build_lora_runtime(_serve_config(tmp_path, quant=None, packed=False)) + + +def test_build_lora_runtime_refuses_an_unusable_adapter(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + def refuse(adapter: LoRAAdapterConfig) -> None: + raise ValueError(f"LoRA adapter {adapter.id!r} carries dense replacement tensors") + + monkeypatch.setattr(lora_module, "check_adapter_file", refuse) + + with pytest.raises(ValueError, match="dense replacement"): + build_lora_runtime(_serve_config(tmp_path)) + + +def test_build_lora_runtime_is_none_without_an_allow_list(tmp_path: Path) -> None: + config = _serve_config(tmp_path) + + assert build_lora_runtime(ServeConfig(generator=config.generator)) is None + + +@pytest.mark.parametrize(("file_name", "dense", "message"), [ + ("adapter.safetensors", True, "dense replacement"), + (None, False, "not a file"), +]) +def test_check_adapter_file_refuses_what_the_pipeline_would(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, + file_name: str | None, dense: bool, message: str) -> None: + target = tmp_path + if file_name is not None: + target = tmp_path / file_name + target.touch() + monkeypatch.setattr(lora_module, "maybe_download_lora", lambda path: str(target)) + monkeypatch.setattr(lora_module.DenseLoRAPatch, "from_adapter", lambda path: object() if dense else None) + + with pytest.raises(ValueError, match=message): + lora_module.check_adapter_file(LoRAAdapterConfig(id="x", path="org/x")) + + +@pytest.mark.parametrize(("packed_dir", "expected"), [("transformer_ref", True), ("transformer", False)]) +def test_side_branch_check_looks_in_the_pipelines_transformer_partition(tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + packed_dir: str, expected: bool) -> None: + (tmp_path / packed_dir).mkdir() + (tmp_path / packed_dir / H3_NVFP4_DIT_EXPORT_FILENAME).touch() + ref2va = type("Ref2VAPipeline", (), {"_extra_config_module_map": {"transformer": "transformer_ref"}}) + monkeypatch.setattr(lora_module, "get_model_info", lambda *args, **kwargs: SimpleNamespace(pipeline_cls=ref2va)) + config = GeneratorConfig(model_path=str(tmp_path), + engine=EngineConfig(quantization=QuantizationConfig(transformer_quant="NVFP4"))) + + with envs.FASTVIDEO_LORA_SIDE_BRANCH.override(False): + assert adapters_attach_as_side_branch(config) is expected + + +@pytest.mark.parametrize(("file_name", "expected"), [(H3_NVFP4_DIT_EXPORT_FILENAME, True), + ("diffusion_pytorch_model.safetensors", False)]) +def test_side_branch_check_follows_transformer_weights(tmp_path: Path, file_name: str, expected: bool) -> None: + weights = tmp_path / "weights" + weights.mkdir() + (weights / file_name).touch() + config = GeneratorConfig(model_path=str(tmp_path / "model"), + engine=EngineConfig(quantization=QuantizationConfig(transformer_quant="NVFP4")), + pipeline=PipelineSelection(components=ComponentConfig(transformer_weights=str(weights)))) + + with envs.FASTVIDEO_LORA_SIDE_BRANCH.override(False): + assert adapters_attach_as_side_branch(config) is expected + + +def test_side_branch_check_resolves_a_hub_model_id(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + (tmp_path / "transformer").mkdir() + (tmp_path / "transformer" / H3_NVFP4_DIT_EXPORT_FILENAME).touch() + requested: list[tuple[str, list[str]]] = [] + + def resolve(model_path: str, revision: str | None, allow_patterns: list[str]) -> str: + requested.append((model_path, revision, allow_patterns)) + return str(tmp_path) + + monkeypatch.setattr(lora_module, "maybe_download_model", resolve) + monkeypatch.setattr(lora_module, "get_model_info", lambda *args, **kwargs: SimpleNamespace(pipeline_cls=object)) + config = GeneratorConfig(model_path="org/model", + revision="v2", + engine=EngineConfig(quantization=QuantizationConfig(transformer_quant="NVFP4"))) + + with envs.FASTVIDEO_LORA_SIDE_BRANCH.override(False): + assert adapters_attach_as_side_branch(config) is True + assert requested == [("org/model", "v2", [f"transformer/{H3_NVFP4_DIT_EXPORT_FILENAME}"])] + + +def test_a_dead_worker_during_an_adapter_change_marks_the_server_unhealthy(make_client) -> None: + generator = _FakeGenerator() + client = make_client(generator) + + def dead_worker(*args: Any, **kwargs: Any) -> None: + raise BrokenPipeError("worker 0 exited") + + generator.set_lora_adapter = dead_worker # type: ignore[method-assign] + + assert _sync(client, name="realism").status_code == 500 + assert client.get("/health").status_code == 503 + assert ("clear",) not in generator.events diff --git a/fastvideo/tests/layers/test_lora_side_branch.py b/fastvideo/tests/layers/test_lora_side_branch.py index 3222fe19e..d1d164492 100644 --- a/fastvideo/tests/layers/test_lora_side_branch.py +++ b/fastvideo/tests/layers/test_lora_side_branch.py @@ -473,3 +473,16 @@ def test_attention_only_adapter_skips_repack_only_once_the_whole_block_is_packed assert not _all_side_branch(layers, [block]) ffn.register_buffer("_nvfp4_weight", torch.zeros(1), persistent=False) assert _all_side_branch(layers, [block]) + + +def test_unmerge_after_a_second_merge_restores_the_original_weight() -> None: + """Restoring must copy the snapshot: aliasing it on its own device lets the next merge overwrite it.""" + base = _dense_linear() + original = base.weight.detach().clone() + layer = get_lora_layer(base) + layer.set_lora_weights(*_factors(1), lora_alpha=8, strength=1.0) + layer.set_lora_weights(*_factors(2), lora_alpha=8, strength=1.0) + + layer.unmerge_lora_weights() + + assert torch.equal(layer.base_layer.weight, original) diff --git a/fastvideo/tests/loader/test_h3_community_lora.py b/fastvideo/tests/loader/test_h3_community_lora.py index 6a861f935..2f36a1302 100644 --- a/fastvideo/tests/loader/test_h3_community_lora.py +++ b/fastvideo/tests/loader/test_h3_community_lora.py @@ -219,6 +219,7 @@ def _pipeline_with_layers(names: list[str], adapter_modules: list[str], min_frac return SimpleNamespace(lora_adapters={"default": adapter}, lora_layers={"transformer": layers}, lora_min_match_fraction=min_fraction, + trainable_transformer_names=["transformer"], trainable_transformer_modules={"transformer": object()}) @@ -238,7 +239,7 @@ def test_coverage_check_waits_for_a_released_lazy_transformer() -> None: pipeline = _pipeline_with_layers([], ["transformer_blocks.0.attn.to_q"]) # A released lazy DiT drops its LoRA layers, so a valid adapter has nothing to match yet. pipeline.lora_layers.pop("transformer") - LoRAPipeline._check_adapter_coverage(pipeline, "default", "adapter.safetensors") + assert not LoRAPipeline._check_adapter_coverage(pipeline, "default", "adapter.safetensors") def test_coverage_check_rejects_a_mostly_unmatched_adapter() -> None: @@ -307,6 +308,7 @@ def test_failed_coverage_restores_the_previously_registered_adapter(monkeypatch: lora_adapters={"default": previous}, lora_adapter_paths={"default": "a.safetensors"}, lora_layers={"transformer": layers}, + trainable_transformer_names=["transformer"], trainable_transformer_modules={}, modules={"transformer": SimpleNamespace(param_names_mapping={}, lora_param_names_mapping={})}, cur_adapter_name="default", @@ -322,6 +324,7 @@ def test_failed_coverage_restores_the_previously_registered_adapter(monkeypatch: convert_adapter_state_dict=lambda state, path: state, ) pipeline._check_adapter_coverage = lambda *a: LoRAPipeline._check_adapter_coverage(pipeline, *a) + pipeline._transformer_name_mappings = lambda: LoRAPipeline._transformer_name_mappings(pipeline) with pytest.raises(ValueError, match="reaches 0 of 1"): LoRAPipeline.set_lora_adapter(pipeline, "default", "b.safetensors") assert pipeline.lora_adapters["default"] is previous @@ -382,6 +385,9 @@ class _RecordingLoRALayer(torch.nn.Module): def set_lora_weights(self, *args, **kwargs) -> None: self.applied = True + def unmerge_lora_weights(self) -> None: + pass + def test_layers_without_adapter_weights_log_one_summary_warning(monkeypatch: pytest.MonkeyPatch) -> None: """An adapter that trains a subset of the wrapped linears warns once, not once per layer.""" diff --git a/fastvideo/tests/pipelines/test_lora_adapter_switch.py b/fastvideo/tests/pipelines/test_lora_adapter_switch.py new file mode 100644 index 000000000..9bac068f9 --- /dev/null +++ b/fastvideo/tests/pipelines/test_lora_adapter_switch.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Switching runtime adapters restores every merged layer the new adapter does not cover (CPU only).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +import fastvideo.envs as envs +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.lora import linear as lora_linear +from fastvideo.layers.lora.linear import get_lora_layer +from fastvideo.pipelines import lora_pipeline +from fastvideo.pipelines.lora_pipeline import LoRAPipeline + +DIM = 8 +RANK = 2 + + +class _Pipe(LoRAPipeline): + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + del fastvideo_args + + +def _factors(seed: int) -> dict[str, torch.Tensor]: + generator = torch.Generator().manual_seed(seed) + return {"lora_A": torch.randn(RANK, DIM, generator=generator), "lora_B": torch.randn(DIM, RANK, generator=generator)} + + +def _adapter(layers: list[str], seed: int) -> dict[str, torch.Tensor]: + return {f"{layer}.{key}": value for layer in layers for key, value in _factors(seed).items()} + + +@pytest.fixture +def setup(monkeypatch: pytest.MonkeyPatch, env_overrides) -> tuple[_Pipe, torch.nn.Module]: + env_overrides.enter_context(envs.FASTVIDEO_LORA_SIDE_BRANCH.override(False)) + monkeypatch.setattr(lora_pipeline.dist, "get_rank", lambda: 0) + monkeypatch.setattr(lora_linear, "get_local_torch_device", lambda: torch.device("cpu")) + torch.manual_seed(0) + transformer = torch.nn.Module() + transformer.first = ReplicatedLinear(DIM, DIM, bias=False, params_dtype=torch.float32) + transformer.second = ReplicatedLinear(DIM, DIM, bias=False, params_dtype=torch.float32) + pipe = _Pipe.__new__(_Pipe) + attributes = dict(lora_adapters={"a": _adapter(["first", "second"], 1), "b": _adapter(["first"], 2)}, + lora_adapter_paths={}, lora_layers={}, cur_adapter_name="", cur_adapter_path="", + cur_adapter_strength=1.0, lora_initialized=False, lora_path=None, + fastvideo_args=SimpleNamespace(training_mode=False), device=torch.device("cpu"), + trainable_transformer_modules={"transformer": transformer}) + for name, value in attributes.items(): + setattr(pipe, name, value) + + def convert_one(name: str, module: torch.nn.Module) -> None: + module.first, module.second = get_lora_layer(module.first), get_lora_layer(module.second) + layers = {"first": module.first, "second": module.second} + pipe.lora_layers[name] = SimpleNamespace(lora_layers_by_block=lambda: iter([(None, layers)])) + + pipe._convert_one_transformer = convert_one # type: ignore[method-assign] + return pipe, transformer + + +def test_switching_restores_a_merged_layer_the_new_adapter_does_not_cover(setup) -> None: + pipeline, transformer = setup + original = transformer.second.weight.detach().clone() + + pipeline.set_lora_adapter("a", strength=1.0) + second = transformer.second + assert not second.side_branch and second.merged + assert not torch.equal(second.base_layer.weight, original) + + pipeline.set_lora_adapter("b", strength=1.0) + assert torch.equal(second.base_layer.weight, original) + assert not second.merged and second.disable_lora + + pipeline.set_lora_adapter("b", strength=0.0) + assert torch.equal(second.base_layer.weight, original) + + +def test_accumulating_keeps_the_earlier_adapter_where_the_new_one_has_nothing(setup) -> None: + pipeline, transformer = setup + pipeline.set_lora_adapter("a", strength=1.0) + merged = transformer.second.base_layer.weight.detach().clone() + + pipeline.set_lora_adapter("b", strength=1.0, accumulate=True) + + assert torch.equal(transformer.second.base_layer.weight, merged) + + +def test_clear_restores_base_weights_after_an_adapter(setup) -> None: + pipeline, transformer = setup + original = {name: getattr(transformer, name).weight.detach().clone() for name in ("first", "second")} + + pipeline.set_lora_adapter("a", strength=1.0) + pipeline.clear_lora_adapter() + + for name in ("first", "second"): + layer = getattr(transformer, name) + assert torch.equal(layer.base_layer.weight, original[name]) + assert layer.lora_A is None and not layer.merged + assert pipeline._runtime_adapter is None and pipeline.cur_adapter_name == "" + pipeline.set_lora_adapter("b", strength=1.0) + assert transformer.first.merged and torch.equal(transformer.second.base_layer.weight, original["second"]) + + +def test_clear_restores_base_weights_after_a_change_that_failed_part_way(setup) -> None: + pipeline, transformer = setup + original = {name: getattr(transformer, name).weight.detach().clone() for name in ("first", "second")} + pipeline.set_lora_adapter("a", strength=1.0) + # B's factors for "second" have the wrong shape, so the change fails after "first" already switched to B. + pipeline.lora_adapters["b"] = {**_adapter(["first"], 2), "second.lora_A": torch.randn(RANK, DIM + 1), + "second.lora_B": torch.randn(DIM, RANK)} + + with pytest.raises(RuntimeError): + pipeline.set_lora_adapter("b", strength=1.0) + pipeline.clear_lora_adapter() + + for name in ("first", "second"): + assert torch.equal(getattr(transformer, name).base_layer.weight, original[name]) + + +def test_clear_after_a_partial_failure_in_a_block_mixing_side_branch_and_merged_layers(monkeypatch, env_overrides) -> None: + from fastvideo.tests.layers.test_lora_side_branch import _packed_linear + + env_overrides.enter_context(envs.FASTVIDEO_LORA_SIDE_BRANCH.override(False)) + monkeypatch.setattr(lora_pipeline.dist, "get_rank", lambda: 0) + monkeypatch.setattr(lora_linear, "get_local_torch_device", lambda: torch.device("cpu")) + transformer = torch.nn.Module() + transformer.packed = _packed_linear() + transformer.dense = ReplicatedLinear(_packed_linear().input_size, _packed_linear().output_size, bias=False, + params_dtype=torch.float32) + in_dim, out_dim = transformer.packed.input_size, transformer.packed.output_size + dense_original = transformer.dense.weight.detach().clone() + generator = torch.Generator().manual_seed(3) + + def factors(rows_in: int) -> dict[str, torch.Tensor]: + return {"lora_A": torch.randn(RANK, rows_in, generator=generator), + "lora_B": torch.randn(out_dim, RANK, generator=generator)} + + pipe = _Pipe.__new__(_Pipe) + adapter_a = {f"{layer}.{key}": value for layer in ("packed", "dense") for key, value in factors(in_dim).items()} + # B applies to "packed" first, then fails on "dense" because its factors have the wrong input size. + adapter_b = {**{f"packed.{key}": value for key, value in factors(in_dim).items()}, + **{f"dense.{key}": value for key, value in factors(in_dim + 1).items()}} + for name, value in dict(lora_adapters={"a": adapter_a, "b": adapter_b}, lora_adapter_paths={}, lora_layers={}, + cur_adapter_name="", cur_adapter_path="", cur_adapter_strength=1.0, lora_initialized=False, + lora_path=None, fastvideo_args=SimpleNamespace(training_mode=False), + device=torch.device("cpu"), trainable_transformer_modules={"transformer": transformer}).items(): + setattr(pipe, name, value) + + def convert_one(name: str, module: torch.nn.Module) -> None: + module.packed, module.dense = get_lora_layer(module.packed), get_lora_layer(module.dense) + layers = {"packed": module.packed, "dense": module.dense} + pipe.lora_layers[name] = SimpleNamespace(lora_layers_by_block=lambda: iter([(None, layers)])) + + pipe._convert_one_transformer = convert_one # type: ignore[method-assign] + pipe.set_lora_adapter("a", strength=1.0) + assert transformer.packed.side_branch and transformer.dense.merged + with pytest.raises(RuntimeError): + pipe.set_lora_adapter("b", strength=1.0) + + pipe.clear_lora_adapter() + + assert transformer.packed.lora_A is None + assert transformer.dense.lora_A is None and not transformer.dense.merged + assert torch.equal(transformer.dense.base_layer.weight, dense_original) diff --git a/fastvideo/tests/pipelines/test_lora_late_transformer.py b/fastvideo/tests/pipelines/test_lora_late_transformer.py new file mode 100644 index 000000000..c6011415d --- /dev/null +++ b/fastvideo/tests/pipelines/test_lora_late_transformer.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +"""A runtime adapter set before the transformer is in memory reaches it when it arrives (CPU only).""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +import fastvideo.envs as envs +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.lora import linear as lora_linear +from fastvideo.layers.lora.linear import get_lora_layer +from fastvideo.pipelines import lora_pipeline +from fastvideo.pipelines.lazy_module import LazyModule +from fastvideo.pipelines.lora_pipeline import LoRAPipeline + +DIM = 8 +RANK = 2 + + +class _Pipe(LoRAPipeline): + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + del fastvideo_args + + +class _Transformer(torch.nn.Module): + param_names_mapping: dict = {} + lora_param_names_mapping: dict = {} + + def __init__(self) -> None: + super().__init__() + self.layer = ReplicatedLinear(DIM, DIM, bias=False, params_dtype=torch.float32) + + +@pytest.fixture +def adapter_file(tmp_path: Path) -> str: + path = tmp_path / "adapter.safetensors" + save_file({"layer.lora_A.weight": torch.randn(RANK, DIM), "layer.lora_B.weight": torch.randn(DIM, RANK)}, str(path)) + return str(path) + + +@pytest.fixture +def make_pipeline(monkeypatch: pytest.MonkeyPatch, env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_LORA_SIDE_BRANCH.override(True)) + monkeypatch.setattr(lora_pipeline.dist, "get_rank", lambda: 0) + monkeypatch.setattr(lora_pipeline, "maybe_download_lora", lambda path: path) + monkeypatch.setattr(lora_linear, "get_local_torch_device", lambda: torch.device("cpu")) + + def make(modules: dict[str, object]) -> _Pipe: + pipeline = _Pipe.__new__(_Pipe) + arch_config = SimpleNamespace(param_names_mapping={}, lora_param_names_mapping={}) + fastvideo_args = SimpleNamespace(training_mode=False, + pipeline_config=SimpleNamespace(dit_config=SimpleNamespace( + arch_config=arch_config))) + attributes = dict(modules=modules, lora_adapters={}, lora_adapter_paths={}, lora_layers={}, + cur_adapter_name="", cur_adapter_path="", cur_adapter_strength=1.0, lora_initialized=False, + lora_path=None, fastvideo_args=fastvideo_args, training_mode=False, + device=torch.device("cpu"), _tracks_late_transformers=True, _coverage_deferred=set(), + trainable_transformer_modules=dict(modules)) + for name, value in attributes.items(): + setattr(pipeline, name, value) + + def convert_one(name: str, module: torch.nn.Module) -> None: + module.layer = get_lora_layer(module.layer) + pipeline.lora_layers[name] = SimpleNamespace( + lora_layers_by_block=lambda: iter([(None, {"layer": module.layer})]), + all_lora_layers=lambda: iter([("layer", module.layer)])) + + pipeline._convert_one_transformer = convert_one # type: ignore[method-assign] + return pipeline + + return make + + +def test_adapter_set_before_sequential_load_reaches_the_transformer(make_pipeline, adapter_file: str) -> None: + pipeline = make_pipeline({}) + + pipeline.set_lora_adapter("realism", adapter_file, strength=0.8) + transformer = _Transformer() + pipeline.add_module("transformer", transformer) + + assert transformer.layer.lora_A is not None + assert transformer.layer.lora_strength == 0.8 + + +def test_transformer_loaded_before_any_adapter_is_wrapped_by_the_first_one(make_pipeline, adapter_file: str) -> None: + pipeline = make_pipeline({}) + transformer = _Transformer() + pipeline.add_module("transformer", transformer) + + pipeline.set_lora_adapter("realism", adapter_file, strength=0.5) + + assert transformer.layer.lora_A is not None + assert transformer.layer.lora_strength == 0.5 + + +def test_loading_an_adapter_does_not_load_a_deferred_transformer(make_pipeline, adapter_file: str) -> None: + transformer = LazyModule("transformer", _Transformer) + pipeline = make_pipeline({"transformer": transformer}) + + pipeline.set_lora_adapter("realism", adapter_file, strength=0.8) + + assert not transformer.is_materialized + assert transformer.materialize().layer.lora_strength == 0.8 + + +def test_a_change_that_fails_part_way_is_not_recorded(make_pipeline, adapter_file: str, + monkeypatch: pytest.MonkeyPatch) -> None: + transformer = _Transformer() + pipeline = make_pipeline({"transformer": transformer}) + pipeline.set_lora_adapter("realism", adapter_file, strength=0.8) + layer_cls = type(transformer.layer) + set_weights = layer_cls.set_lora_weights + + def fail(*args, **kwargs) -> None: + raise RuntimeError("device lost") + + monkeypatch.setattr(layer_cls, "set_lora_weights", fail) + with pytest.raises(RuntimeError, match="device lost"): + pipeline.set_lora_adapter("realism", strength=0.3) + monkeypatch.setattr(layer_cls, "set_lora_weights", set_weights) + + pipeline.set_lora_adapter("realism", strength=0.3) + + assert transformer.layer.lora_strength == 0.3 + + +def test_coverage_deferred_by_sequential_load_is_checked_when_the_transformer_arrives( + make_pipeline, tmp_path: Path) -> None: + foreign = tmp_path / "foreign.safetensors" + save_file({"unet.0.lora_A.weight": torch.randn(RANK, DIM), "unet.0.lora_B.weight": torch.randn(DIM, RANK)}, + str(foreign)) + pipeline = make_pipeline({}) + pipeline.lora_strict_coverage = True + pipeline.set_lora_adapter("foreign", str(foreign), strength=0.8) + + transformer = _Transformer() + with pytest.raises(ValueError, match="reaches 0 of 1"): + pipeline.add_module("transformer", transformer) + + assert "foreign" not in pipeline.lora_adapters and pipeline._runtime_adapter is None + assert transformer.layer.lora_A is None + + +def test_a_lazy_transformer_whose_deferred_coverage_fails_is_not_kept_alive(make_pipeline, tmp_path: Path) -> None: + foreign = tmp_path / "foreign.safetensors" + save_file({"unet.0.lora_A.weight": torch.randn(RANK, DIM), "unet.0.lora_B.weight": torch.randn(DIM, RANK)}, + str(foreign)) + transformer = LazyModule("transformer", _Transformer) + pipeline = make_pipeline({"transformer": transformer}) + pipeline.lora_strict_coverage = True + pipeline.set_lora_adapter("foreign", str(foreign), strength=0.8) + + with pytest.raises(ValueError, match="reaches 0 of 1"): + transformer.materialize() + + assert not transformer.is_materialized + assert "transformer" not in pipeline.lora_layers diff --git a/fastvideo/tests/pipelines/test_lora_lazy_reload.py b/fastvideo/tests/pipelines/test_lora_lazy_reload.py new file mode 100644 index 000000000..edfcb4ab5 --- /dev/null +++ b/fastvideo/tests/pipelines/test_lora_lazy_reload.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +"""A runtime LoRA adapter survives the release and reload of a lazy transformer (CPU only).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +import fastvideo.envs as envs +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.lora.linear import get_lora_layer +from fastvideo.pipelines import lora_pipeline +from fastvideo.pipelines.lazy_module import LazyModule +from fastvideo.pipelines.lora_pipeline import LoRAPipeline + +DIM = 8 +RANK = 2 + + +class _Pipe(LoRAPipeline): + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + del fastvideo_args + + +class _Transformer(torch.nn.Module): + + def __init__(self) -> None: + super().__init__() + self.layer = ReplicatedLinear(DIM, DIM, bias=False, params_dtype=torch.bfloat16) + + +def _lazy_pipeline() -> tuple[_Pipe, LazyModule]: + pipeline = _Pipe.__new__(_Pipe) + transformer = LazyModule("transformer", _Transformer) + adapter = {"layer.lora_A": torch.randn(RANK, DIM), "layer.lora_B": torch.randn(DIM, RANK)} + attributes = dict(lora_adapters={"realism": adapter}, lora_adapter_paths={}, lora_layers={}, cur_adapter_name="", + cur_adapter_path="", cur_adapter_strength=1.0, lora_initialized=False, lora_path=None, + fastvideo_args=SimpleNamespace(training_mode=False), device=torch.device("cpu"), + trainable_transformer_modules={"transformer": transformer}) + for name, value in attributes.items(): + setattr(pipeline, name, value) + + def convert_one(name: str, module: torch.nn.Module) -> None: + module.layer = get_lora_layer(module.layer) + pipeline.lora_layers[name] = SimpleNamespace(lora_layers_by_block=lambda: iter([(None, {"layer": module.layer})])) + + pipeline._convert_one_transformer = convert_one # type: ignore[method-assign] + return pipeline, transformer + + +def test_runtime_adapter_is_reapplied_after_a_lazy_reload(monkeypatch, env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_LORA_SIDE_BRANCH.override(True)) + monkeypatch.setattr(lora_pipeline.dist, "get_rank", lambda: 0) + pipeline, transformer = _lazy_pipeline() + pipeline.convert_to_lora_layers() + + first = transformer.materialize() + pipeline.set_lora_adapter("realism", strength=0.8) + assert first.layer.lora_A is not None + assert transformer.release() + + second = transformer.materialize() + assert second is not first + assert second.layer.lora_A is not None + assert second.layer.lora_strength == 0.8 + + + +def test_adapter_change_on_a_released_transformer_does_not_load_it(monkeypatch, env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_LORA_SIDE_BRANCH.override(True)) + monkeypatch.setattr(lora_pipeline.dist, "get_rank", lambda: 0) + pipeline, transformer = _lazy_pipeline() + pipeline.convert_to_lora_layers() + transformer.materialize() + pipeline.set_lora_adapter("realism", strength=0.8) + assert transformer.release() + + pipeline.set_lora_adapter("realism", strength=0.0) + + assert not transformer.is_materialized + assert transformer.materialize().layer.lora_strength == 0.0 + + +def test_a_cleared_adapter_is_not_reapplied_on_reload(monkeypatch, env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_LORA_SIDE_BRANCH.override(True)) + monkeypatch.setattr(lora_pipeline.dist, "get_rank", lambda: 0) + pipeline, transformer = _lazy_pipeline() + pipeline.convert_to_lora_layers() + transformer.materialize() + pipeline.set_lora_adapter("realism", strength=0.8) + + pipeline.clear_lora_adapter() + assert transformer.materialize().layer.lora_A is None + assert transformer.release() + + assert transformer.materialize().layer.lora_A is None diff --git a/fastvideo/worker/executor.py b/fastvideo/worker/executor.py index 17f9ce471..555a1c98e 100644 --- a/fastvideo/worker/executor.py +++ b/fastvideo/worker/executor.py @@ -120,6 +120,13 @@ class Executor(ABC): """ raise NotImplementedError + def clear_lora_adapter(self) -> None: + """Return every worker's LoRA layers to the base weights.""" + responses: list[dict[str, Any]] = self.collective_rpc("clear_lora_adapter", kwargs={}) + for i, response in enumerate(responses): + if response["status"] != "lora_adapter_cleared": + raise RuntimeError(f"Worker {i} failed to clear the LoRA adapter: {response['status']}") + @abstractmethod def merge_lora_weights(self) -> None: """ diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index ed86afa5b..7a80d8a9e 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -216,6 +216,12 @@ class Worker: return {"status": "lora_adapter_set"} return {"status": "failed: pipeline is not a LoRAPipeline"} + def clear_lora_adapter(self) -> dict[str, Any]: + if isinstance(self.pipeline, LoRAPipeline): + self.pipeline.clear_lora_adapter() + return {"status": "lora_adapter_cleared"} + return {"status": "failed: pipeline is not a LoRAPipeline"} + def unmerge_lora_weights(self) -> dict[str, Any]: if isinstance(self.pipeline, LoRAPipeline): self.pipeline.unmerge_lora_weights()