module-attribute ¶
DiffusionModelRegistry = _ModelRegistry(
{
model_arch: _LazyRegisteredModel(
module_name=f"vllm_omni.diffusion.models.{mod_folder}.{mod_relname}",
class_name=cls_name,
)
for model_arch, (
mod_folder,
mod_relname,
cls_name,
) in _DIFFUSION_MODELS.items()
}
)
get_diffusion_ir_op_priority_func(
od_config: OmniDiffusionConfig,
)
get_diffusion_post_process_func(
od_config: OmniDiffusionConfig,
)
get_diffusion_prefix_cache_func(
od_config: OmniDiffusionConfig,
)
Load optional model preparation for native multimodal KV identities.
initialize_model(od_config: OmniDiffusionConfig) -> Module
Initialize a diffusion model from the registry.
This function: 1. Loads the model class from the registry 2. Instantiates the model with the config 3. Configures VAE optimization settings 4. Applies sequence parallelism if enabled (similar to diffusers' enable_parallelism)
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
od_config | OmniDiffusionConfig | The OmniDiffusion configuration. | required |
Returns:
| Type | Description |
|---|---|
Module | The initialized pipeline model. |
Raises:
| Type | Description |
|---|---|
ValueError | If the model class is not found in the registry. |
register_diffusion_model(
model_arch: str,
module_name: str,
class_name: str,
pre_process_func_name: str | None = None,
post_process_func_name: str | None = None,
ir_op_priority_func_name: str | None = None,
action_post_process_func_name: str | None = None,
prefix_cache_func_name: str | None = None,
) -> None
Register a diffusion model pipeline from an out-of-tree plugin.
This can be used to add new model architectures or to replace an existing built-in pipeline with a platform-optimised implementation (same model_arch key).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model_arch | str | Architecture name (e.g. | required |
module_name | str | Fully qualified module path (e.g. | required |
class_name | str | Class name within module_name. | required |
pre_process_func_name | str | None | Optional name of the pre-process function located in module_name. Pass | None |
post_process_func_name | str | None | Optional name of the post-process function located in module_name. Pass | None |
ir_op_priority_func_name | str | None | Optional name of the IR op priority merge function located in module_name. Pass | None |
action_post_process_func_name | str | None | Deprecated compatibility-only keyword for out-of-tree plugins. Action postprocess hooks are no longer registered separately; move action handling into | None |
prefix_cache_func_name | str | None | Optional factory for a CPU cache-input hook, called after preprocessing only when paged prefix caching is enabled. It fills existing DiffusionKVRequest cache inputs in place. | None |