vllm_omni.diffusion.registry ¶

DiffusionModelRegistry `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 (items())
    }
)

logger `module-attribute` ¶

logger = init_logger(__name__)

get_diffusion_post_process_func ¶

get_diffusion_post_process_func(
    od_config: OmniDiffusionConfig,
)

get_diffusion_pre_process_func ¶

get_diffusion_pre_process_func(
    od_config: OmniDiffusionConfig,
)

initialize_model ¶

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 ¶

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,
) -> 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).