[diffusion] feat: enable passing Cache‑DiT config for diffusers backend (#16662)
Signed-off-by: Chi <chixie.mcisaac@gmail.com> Signed-off-by: qimcis <chixie.mcisaac@gmail.com> Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
@@ -235,7 +235,9 @@ class ServerArgs:
|
||||
|
||||
# Attention
|
||||
attention_backend: str = None
|
||||
diffusers_attention_backend: str = None # for diffusers backend only
|
||||
cache_dit_config: str | dict[str, Any] | None = (
|
||||
None # cache-dit config for diffusers
|
||||
)
|
||||
|
||||
# Distributed executor backend
|
||||
nccl_port: Optional[int] = None
|
||||
@@ -452,15 +454,25 @@ class ServerArgs:
|
||||
"--attention-backend",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=[e.name.lower() for e in AttentionBackendEnum] + ["fa3", "fa4"],
|
||||
help="The attention backend to use. If not specified, the backend is automatically selected based on hardware and installed packages.",
|
||||
help=(
|
||||
"The attention backend to use. For SGLang-native pipelines, use "
|
||||
"values like fa, torch_sdpa, sage_attn, etc. For diffusers pipelines, "
|
||||
"use diffusers attention backend names such as flash, _flash_3_hub, "
|
||||
"sage, or xformers."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--diffusers-attention-backend",
|
||||
type=str,
|
||||
dest="attention_backend",
|
||||
default=None,
|
||||
help="Attention backend for diffusers pipelines (e.g., flash, _flash_3_hub, sage, xformers). "
|
||||
"See: https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends",
|
||||
help=argparse.SUPPRESS,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache-dit-config",
|
||||
type=str,
|
||||
default=ServerArgs.cache_dit_config,
|
||||
help="Path to a Cache-DiT YAML/JSON config. Enables cache-dit for diffusers backend.",
|
||||
)
|
||||
|
||||
# HuggingFace specific parameters
|
||||
@@ -999,7 +1011,7 @@ class ServerArgs:
|
||||
raise ValueError("pipeline_config is not set in ServerArgs")
|
||||
|
||||
self.pipeline_config.check_pipeline_config()
|
||||
if self.attention_backend is None:
|
||||
if self.attention_backend is None and self.backend != Backend.DIFFUSERS:
|
||||
self._set_default_attention_backend()
|
||||
|
||||
# parallelism
|
||||
|
||||
Reference in New Issue
Block a user