Common¶
This page documents the common utilities and base classes used by all RBLN Diffusion model pipelines.
Supported NPU
Diffusers pipelines run on ATOM™ only. Rebel100™ does not support them yet.
API Reference¶
RBLNDiffusionMixin ¶
Bases: _PipelineBase
RBLNDiffusionMixin provides essential functionalities for compiling Stable Diffusion pipeline components to run on RBLN NPUs. This mixin class serves as a base for implementing RBLN-compatible Stable Diffusion pipelines. It contains shared logic for handling the core components of Stable Diffusion.
The class has no base class at runtime. For type checking only it declares DiffusionPipeline as its base, because every concrete pipeline mixes it in next to a DiffusionPipeline subclass and that is what self.load_config, self.vae and the super() calls resolve to.
To use this mixin:
- Create a new pipeline class that inherits from both this mixin and the original StableDiffusionPipeline.
- Define the required _submodules and _optional_submodules class variable listing the components to be compiled.
Example
Class Variables
_submodules: List of submodule names that should be compiled (typically ["text_encoder", "unet", "vae"])_optional_submodules: List of submodule names compiled without inheriting RBLNModel (typically ["safety_checker"])
Methods:
| Name | Description |
|---|---|
from_pretrained | Creates and optionally compiles a model from a pretrained checkpoint |
Notes
- When
export=True, all compatible submodules will be compiled for NPU inference - The compilation config can be customized per submodule by including submodule names as keys in rbln_config
from_pretrained(model_id, *, export=None, model_save_dir=None, rbln_config=None, lora_ids=None, lora_weights_names=None, lora_scales=None, **kwargs) classmethod ¶
Load a pretrained diffusion pipeline from a model checkpoint, with optional compilation for RBLN NPUs.
It supports various diffusion pipelines including Stable Diffusion, Kandinsky, ControlNet, and other diffusers-based models.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model_id | `str` | The model ID or path to the pretrained model to load. Can be either:
| required |
export | bool | None | If True, takes a PyTorch model from | None |
model_save_dir | str | PathLike[str] | None | Directory to save the compiled model artifacts. Only used when | None |
rbln_config | dict[str, Any] | RBLNModelConfig | None | Configuration options for RBLN compilation. Can include settings for specific submodules such as | None |
lora_ids | str | list[str] | None | LoRA adapter ID(s) to load and apply before compilation. LoRA weights are fused into the model weights during compilation. Only used when | None |
lora_weights_names | str | list[str] | None | Names of specific LoRA weight files to load, corresponding to lora_ids. Only used when | None |
lora_scales | float | list[float] | None | Scaling factor(s) to apply to the LoRA adapter(s). Only used when | None |
kwargs | Any | Additional arguments to pass to the underlying diffusion pipeline constructor or the RBLN compilation process. These may include parameters specific to individual submodules or the particular diffusion pipeline being used. | {} |
Returns:
| Type | Description |
|---|---|
Self | A compiled or loaded diffusion pipeline that can be used for inference on RBLN NPU. The returned object is an instance of the class that called this method, inheriting from RBLNDiffusionMixin. |
RBLNDiffusionMixinConfig ¶
Bases: RBLNModelConfig
Configuration class for RBLN diffusion pipelines.
__init__(cls_name=None, create_runtimes=None, device=None, device_map=None, activate_profiler=None, npu=None, num_devices=None, timeout=None, optimum_rbln_version=None, dtype=None, _compile_cfgs=None, *, optimize_host_memory=None, **kwargs) ¶
Initialize a RBLN model configuration with runtime options and compile configurations.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
cls_name | str | None | The class name of the configuration. Defaults to the current class name. | None |
create_runtimes | bool | None | Whether to create RBLN runtimes. Defaults to True. | None |
device | int | list[int] | None | The device(s) to load the model onto. Can be a single device ID or a list. | None |
device_map | dict[str, int | list[int]] | None | Mapping from compiled model names to device IDs. | None |
activate_profiler | bool | None | Whether to activate the profiler for performance analysis. | None |
npu | str | None | The NPU device name to use for compilation. | None |
num_devices | int | None | Number of devices to distribute the model across. | None |
timeout | int | None | The timeout for the runtime in seconds. If it isn't provided, it will be set to 60 by default. | None |
optimum_rbln_version | str | None | The optimum-rbln version used for this configuration. | None |
dtype | str | dtype | None | The data type to use for the model. | None |
_compile_cfgs | list[RBLNCompileConfig] | List of compilation configurations for the model. | None |
kwargs | Any | Additional keyword arguments. | {} |
Raises:
| Type | Description |
|---|---|
ValueError | If unexpected keyword arguments are provided. |
from_pretrained(path, rbln_config=None, return_unused_kwargs=False, **kwargs) classmethod ¶
Load a RBLNModelConfig from a path.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
path | str | Path to the RBLNModelConfig file or directory containing the config file. | required |
rbln_config | dict[str, Any] | None | Additional configuration to override. | None |
return_unused_kwargs | bool | Whether to return unused kwargs. | False |
kwargs | Any | Additional keyword arguments to override configuration values. Keys starting with 'rbln_' will have the prefix removed and be used to update the configuration. | {} |
Returns:
| Name | Type | Description |
|---|---|---|
RBLNModelConfig | Self | tuple[Self, dict[str, Any]] | The loaded configuration instance. |
Note
This method loads the configuration from the specified path and applies any provided overrides. If the loaded configuration class doesn't match the expected class, a warning will be logged.
Examples:
initialize_submodule_config(submodule_config=None, force_kwargs=False, **kwargs) ¶
Build the config of a submodule from the parent's runtime options and kwargs.
Returns the submodule config instance when its class is known (an instance was passed, or cls_name was given), otherwise the merged kwargs dict for the model class to resolve later.
get_load_overrides() ¶
Extract the options that must cross the load boundary, for passing to from_pretrained.
Returns the explicitly-set runtime options (device, device_map, create_runtimes, activate_profiler, timeout) and the subclass-specific non-save attributes (load-behavior flags such as _load_visual_runtime, which are never serialized so they can only travel with the caller), recursively including submodules. Compile-time attributes are not included: when loading a compiled model they come from the artifact's rbln_config.json, and this dict is merged on top of it.
Returns:
| Type | Description |
|---|---|
dict[str, Any] | Dictionary of load overrides, keyed like the |