7d6904aa94
Signed-off-by: 0z5a <Dezhen.lu@student.uni-tuebingen.de> Signed-off-by: 0z5a <dezhen.lu@uni-tuebingen.de> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de> Signed-off-by: princepride <wangzhipeng628@gmail.com> Co-authored-by: 0z5a <dezhen.lu@uni-tuebingen.de> Co-authored-by: princepride <wangzhipeng628@gmail.com>
60 lines
2.6 KiB
Python
60 lines
2.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Operator limits and colour options owned by SeedVR2."""
|
|
|
|
import os
|
|
import tempfile
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
COLOR_CORRECTION_METHODS = ("lab", "wavelet", "adain", "none")
|
|
DEFAULT_COLOR_CORRECTION_METHOD = "lab"
|
|
|
|
|
|
def _positive_int(name: str, default: int | None) -> int | None:
|
|
"""Read an operator-tuned limit, rejecting values that would disable it."""
|
|
value = os.environ.get(name)
|
|
if value is None:
|
|
return default
|
|
if not value.strip().isdigit() or not int(value):
|
|
raise ValueError(f"{name} must be a positive integer, got {value!r}")
|
|
return int(value)
|
|
|
|
|
|
environment_variables: dict[str, Callable[[], Any]] = {
|
|
# ================== SeedVR2 Restoration Env Vars ==================
|
|
# Admission budgets for the sharded serving profile. The defaults are
|
|
# calibrated for the smallest qualified device, so larger accelerators raise
|
|
# them here; setting them beyond what the device holds surfaces as a CUDA
|
|
# OOM during the forward pass. See docs/models/seedvr2.md.
|
|
"VLLM_OMNI_SEEDVR2_MAX_FRAMES": lambda: _positive_int("VLLM_OMNI_SEEDVR2_MAX_FRAMES", 257),
|
|
"VLLM_OMNI_SEEDVR2_SHARDED_FRAME_PIXELS": lambda: _positive_int(
|
|
"VLLM_OMNI_SEEDVR2_SHARDED_FRAME_PIXELS", 2560 * 1472
|
|
),
|
|
# Unset, the clip budget scales with device memory from the calibrated
|
|
# 32 GB value; see seedvr2.pipeline_seedvr2.sharded_budget.
|
|
"VLLM_OMNI_SEEDVR2_SHARDED_CLIP_PIXELS": lambda: _positive_int("VLLM_OMNI_SEEDVR2_SHARDED_CLIP_PIXELS", None),
|
|
# Bounds for the long-video restoration route.
|
|
"VLLM_OMNI_SEEDVR2_LONG_MAX_UPLOAD_BYTES": lambda: _positive_int(
|
|
"VLLM_OMNI_SEEDVR2_LONG_MAX_UPLOAD_BYTES", 8 * 1024**3
|
|
),
|
|
# Longest model window. The clip budget usually binds first; this caps it
|
|
# at 30 latent frames, the DiT's largest temporal attention window.
|
|
"VLLM_OMNI_SEEDVR2_LONG_MAX_WINDOW": lambda: _positive_int("VLLM_OMNI_SEEDVR2_LONG_MAX_WINDOW", 121),
|
|
"VLLM_OMNI_SEEDVR2_LONG_JOB_TTL_SECONDS": lambda: _positive_int("VLLM_OMNI_SEEDVR2_LONG_JOB_TTL_SECONDS", 3600),
|
|
"VLLM_OMNI_SEEDVR2_LONG_OUTPUT_DIR": lambda: os.environ.get(
|
|
"VLLM_OMNI_SEEDVR2_LONG_OUTPUT_DIR", tempfile.gettempdir()
|
|
),
|
|
}
|
|
|
|
|
|
def __getattr__(name):
|
|
# lazy evaluation of environment variables
|
|
if name in environment_variables:
|
|
return environment_variables[name]()
|
|
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
|
|
|
|
|
def __dir__():
|
|
return list(environment_variables.keys())
|