241 lines
9.3 KiB
Python
241 lines
9.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import inspect
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
|
|
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
LTX25_ORIGINAL_MODEL_ID = "Lightricks/LTX-2.5"
|
|
LTX25_DIFFUSERS_MODEL_ID = "Lightricks/LTX-2.5-Diffusers"
|
|
LTX25_DIFFUSERS_COMMIT = "7564fb016dabda0c943416190fc92398c50b1b20"
|
|
|
|
# Diffusers' public LTX-2.5 schedule intentionally excludes the scheduler's
|
|
# terminal zero. FlowMatchEulerDiscreteScheduler appends that internally.
|
|
LTX25_DISTILLED_SIGMAS = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875]
|
|
|
|
_LTX25_DISTILLED_GUIDANCE_DEFAULTS: dict[str, float] = {
|
|
"guidance_scale": 1.0,
|
|
"audio_guidance_scale": 1.0,
|
|
"stg_scale": 0.0,
|
|
"audio_stg_scale": 0.0,
|
|
"modality_scale": 1.0,
|
|
"audio_modality_scale": 1.0,
|
|
"guidance_rescale": 0.0,
|
|
"audio_guidance_rescale": 0.0,
|
|
}
|
|
|
|
|
|
def _supports_ltx25_diffusers() -> bool:
|
|
"""Check the standard-pipeline capabilities added with LTX-2.5."""
|
|
try:
|
|
import diffusers
|
|
import transformers
|
|
|
|
pipeline_class = diffusers.LTX2Pipeline
|
|
init_parameters = inspect.signature(pipeline_class.__init__).parameters
|
|
return (
|
|
"duration_head" in init_parameters
|
|
and "prompt_enhancer" in init_parameters
|
|
and hasattr(transformers, "Gemma4UnifiedForConditionalGeneration")
|
|
)
|
|
except (AttributeError, ImportError, RuntimeError):
|
|
return False
|
|
|
|
|
|
def is_ltx25_diffusers_model(od_config: OmniDiffusionConfig) -> bool:
|
|
"""Return whether config targets the official converted LTX-2.5 repo."""
|
|
normalized_model = str(od_config.model).rstrip("/\\").replace("\\", "/").lower()
|
|
if normalized_model == LTX25_DIFFUSERS_MODEL_ID.lower():
|
|
return True
|
|
|
|
# A downloaded checkpoint may be served from any directory name. The
|
|
# duration-head component is the pipeline-level marker introduced for
|
|
# LTX-2.5; older LTX-2.0/2.3 indexes do not contain it.
|
|
model_path = Path(str(od_config.model)).expanduser()
|
|
subfolder = od_config.diffusers_load_kwargs.get("subfolder")
|
|
if subfolder is not None:
|
|
model_path /= str(subfolder)
|
|
try:
|
|
with (model_path / "model_index.json").open(encoding="utf-8") as model_index_file:
|
|
model_index = json.load(model_index_file)
|
|
except (OSError, TypeError, ValueError):
|
|
return False
|
|
if not isinstance(model_index, dict):
|
|
return False
|
|
|
|
duration_head = model_index.get("duration_head")
|
|
return (
|
|
model_index.get("_class_name") == "LTX2Pipeline"
|
|
and isinstance(duration_head, list)
|
|
and bool(duration_head)
|
|
and duration_head[0] is not None
|
|
)
|
|
|
|
|
|
class BasePipelineUtils:
|
|
"""No-op hooks for pipeline-specific diffusers adapter behavior."""
|
|
|
|
def update_load_kwargs(self, od_config: OmniDiffusionConfig, load_kwargs: dict[str, Any]) -> None:
|
|
pass
|
|
|
|
def apply_post_load_updates(self, pipeline: DiffusionPipeline, od_config: OmniDiffusionConfig) -> None:
|
|
pass
|
|
|
|
def enable_vae_optimization(self, pipeline: DiffusionPipeline, optimization: str) -> None:
|
|
getattr(pipeline, f"enable_vae_{optimization}")()
|
|
|
|
def validate_runtime_sampling_params(self, sampling: OmniDiffusionSamplingParams) -> None:
|
|
pass
|
|
|
|
def update_call_kwargs(
|
|
self,
|
|
od_config: OmniDiffusionConfig,
|
|
call_kwargs: dict[str, Any],
|
|
) -> None:
|
|
pass
|
|
|
|
def normalize_output(
|
|
self,
|
|
pipeline: DiffusionPipeline,
|
|
od_config: OmniDiffusionConfig,
|
|
output: Any,
|
|
) -> Any:
|
|
return output
|
|
|
|
|
|
def validate_model_compatibility(od_config: OmniDiffusionConfig) -> None:
|
|
"""Reject known model/dependency combinations before loading weights."""
|
|
normalized_model = str(od_config.model).rstrip("/\\").replace("\\", "/").lower()
|
|
if normalized_model == LTX25_ORIGINAL_MODEL_ID.lower():
|
|
raise ValueError(
|
|
f"{LTX25_ORIGINAL_MODEL_ID} is a component-weight repository without a Diffusers "
|
|
f"model_index.json. Use the official converted repository {LTX25_DIFFUSERS_MODEL_ID} "
|
|
"with --diffusion-load-format diffusers."
|
|
)
|
|
|
|
if is_ltx25_diffusers_model(od_config) and not _supports_ltx25_diffusers():
|
|
raise ImportError(
|
|
f"{LTX25_DIFFUSERS_MODEL_ID} requires unreleased LTX-2.5 support from Diffusers and "
|
|
"Gemma 4 Unified support from Transformers. Install the source-preview dependencies with "
|
|
"`python -m pip install --upgrade 'transformers>=5.10.1' "
|
|
f"'diffusers @ git+https://github.com/huggingface/diffusers.git@{LTX25_DIFFUSERS_COMMIT}'`."
|
|
)
|
|
|
|
|
|
class LTX2PipelineUtils(BasePipelineUtils):
|
|
"""Diffusers call policy for the official default LTX-2.5 checkpoint."""
|
|
|
|
def enable_vae_optimization(self, pipeline: DiffusionPipeline, optimization: str) -> None:
|
|
pipeline_method = getattr(pipeline, f"enable_vae_{optimization}", None)
|
|
if callable(pipeline_method):
|
|
pipeline_method()
|
|
else:
|
|
getattr(pipeline.vae, f"enable_{optimization}")()
|
|
|
|
def update_call_kwargs(
|
|
self,
|
|
od_config: OmniDiffusionConfig,
|
|
call_kwargs: dict[str, Any],
|
|
) -> None:
|
|
if not is_ltx25_diffusers_model(od_config):
|
|
return
|
|
|
|
has_sigmas = call_kwargs.get("sigmas") is not None
|
|
has_timesteps = call_kwargs.get("timesteps") is not None
|
|
has_step_count = call_kwargs.get("num_inference_steps") is not None
|
|
if not has_sigmas and not has_timesteps:
|
|
if has_step_count:
|
|
logger.warning(
|
|
"LTX-2.5 distilled inference received num_inference_steps without sigmas. "
|
|
"Diffusers will use a generic schedule; omit num_inference_steps to use the "
|
|
"official distilled schedule."
|
|
)
|
|
else:
|
|
call_kwargs["sigmas"] = list(LTX25_DISTILLED_SIGMAS)
|
|
call_kwargs["num_inference_steps"] = len(LTX25_DISTILLED_SIGMAS)
|
|
|
|
for key, value in _LTX25_DISTILLED_GUIDANCE_DEFAULTS.items():
|
|
if call_kwargs.get(key) is None:
|
|
call_kwargs[key] = value
|
|
|
|
def normalize_output(
|
|
self,
|
|
pipeline: DiffusionPipeline,
|
|
od_config: OmniDiffusionConfig,
|
|
output: Any,
|
|
) -> Any:
|
|
if not is_ltx25_diffusers_model(od_config):
|
|
return output
|
|
|
|
frames = getattr(output, "frames", None)
|
|
audio = getattr(output, "audio", None)
|
|
if frames is None or audio is None:
|
|
return output
|
|
|
|
envelope: dict[str, Any] = {"payload": {"video": frames, "audio": audio}}
|
|
config = getattr(getattr(pipeline, "vocoder", None), "config", None)
|
|
value = (
|
|
config.get("output_sampling_rate")
|
|
if hasattr(config, "get")
|
|
else getattr(config, "output_sampling_rate", None)
|
|
)
|
|
try:
|
|
sample_rate = None if isinstance(value, bool) else int(value)
|
|
except (TypeError, ValueError):
|
|
sample_rate = None
|
|
if sample_rate is not None and sample_rate > 0:
|
|
envelope["metadata"] = {"audio": {"sample_rate": sample_rate}}
|
|
return envelope
|
|
|
|
|
|
class WanPipelineUtils(BasePipelineUtils):
|
|
def update_load_kwargs(self, od_config: OmniDiffusionConfig, load_kwargs: dict[str, Any]) -> None:
|
|
if od_config.boundary_ratio is not None:
|
|
load_kwargs["boundary_ratio"] = od_config.boundary_ratio
|
|
|
|
def apply_post_load_updates(self, pipeline: DiffusionPipeline, od_config: OmniDiffusionConfig) -> None:
|
|
if od_config.flow_shift is not None:
|
|
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
|
|
|
|
pipeline.scheduler = UniPCMultistepScheduler.from_config(
|
|
pipeline.scheduler.config, flow_shift=od_config.flow_shift
|
|
)
|
|
|
|
def validate_runtime_sampling_params(self, sampling: OmniDiffusionSamplingParams) -> None:
|
|
if sampling.boundary_ratio is not None:
|
|
raise ValueError(
|
|
"Boundary ratio is not supported at runtime with the diffusers backend for Wan models. Please set "
|
|
"it at model loading time using the `boundary_ratio` kwarg or `--diffusers-load-kwargs` JSON."
|
|
)
|
|
if sampling.extra_args.get("flow_shift") is not None:
|
|
raise ValueError(
|
|
"Flow shift is not supported at runtime with the diffusers backend for Wan models. Please set "
|
|
"it at model loading time using the `flow_shift` kwarg."
|
|
)
|
|
|
|
|
|
PIPELINE_UTILS_REGISTRY: dict[str, type[BasePipelineUtils]] = {
|
|
"LTX2Pipeline": LTX2PipelineUtils,
|
|
"WanPipeline": WanPipelineUtils,
|
|
"WanImageToVideoPipeline": WanPipelineUtils,
|
|
"WanVACEPipeline": WanPipelineUtils,
|
|
"WanVideoToVideoPipeline": WanPipelineUtils,
|
|
"WanAnimatePipeline": WanPipelineUtils,
|
|
}
|
|
|
|
|
|
def get_pipeline_utils(pipeline_class_name: str | None) -> BasePipelineUtils:
|
|
if pipeline_class_name is None:
|
|
return BasePipelineUtils()
|
|
pipeline_utils_cls = PIPELINE_UTILS_REGISTRY.get(pipeline_class_name, BasePipelineUtils)
|
|
return pipeline_utils_cls()
|