Common

This page documents the common utilities and base classes used by all RBLN Diffusion model pipelines.

API Reference

Classes

RBLNDiffusionMixin

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.

To use this mixin:

1. Create a new pipeline class that inherits from both this mixin and the original StableDiffusionPipeline.
2. Define the required _submodules and _optional_submodules class variable listing the components to be compiled.

Example

class RBLNStableDiffusionPipeline(RBLNDiffusionMixin, StableDiffusionPipeline):
    _submodules = ["text_encoder", "unet", "vae"]

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

Functions

from_pretrained(model_id, *, export=None, model_save_dir=None, rbln_config={}, 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.

This method has two distinct operating modes

• When export=True: Takes a PyTorch-based diffusion model, compiles it for RBLN NPUs, and loads the compiled model
• When export=False: Loads an already compiled RBLN model from model_id without recompilation

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: A model ID from the HuggingFace Hub A local path to a saved model directory	required
export	bool	If True, takes a PyTorch model from model_id and compiles it for RBLN NPU execution. If False, loads an already compiled RBLN model from model_id without recompilation.	None
model_save_dir	Optional[PathLike]	Directory to save the compiled model artifacts. Only used when export=True. If not provided and export=True, a temporary directory is used.	None
rbln_config	Dict[str, Any]	Configuration options for RBLN compilation. Can include settings for specific submodules such as text_encoder, unet, and vae. Configuration can be tailored to the specific pipeline being compiled.	{}
lora_ids	Optional[Union[str, List[str]]]	LoRA adapter ID(s) to load and apply before compilation. LoRA weights are fused into the model weights during compilation. Only used when export=True.	None
lora_weights_names	Optional[Union[str, List[str]]]	Names of specific LoRA weight files to load, corresponding to lora_ids. Only used when export=True.	None
lora_scales	Optional[Union[float, List[float]]]	Scaling factor(s) to apply to the LoRA adapter(s). Only used when export=True.	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
RBLNDiffusionMixin	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.

Functions

__init__(cls_name=None, create_runtimes=None, optimize_host_memory=None, device=None, device_map=None, activate_profiler=None, npu=None, tensor_parallel_size=None, timeout=None, optimum_rbln_version=None, _torch_dtype=None, _compile_cfgs=[], **kwargs)

Initialize a RBLN model configuration with runtime options and compile configurations.

Parameters:

Name	Type	Description	Default
cls_name	Optional[str]	The class name of the configuration. Defaults to the current class name.	None
create_runtimes	Optional[bool]	Whether to create RBLN runtimes. Defaults to True.	None
optimize_host_memory	Optional[bool]	Whether to optimize host memory usage. Defaults to True.	None
device	Optional[Union[int, List[int]]]	The device(s) to load the model onto. Can be a single device ID or a list.	None
device_map	Optional[Dict[str, Union[int, List[int]]]]	Mapping from compiled model names to device IDs.	None
activate_profiler	Optional[bool]	Whether to activate the profiler for performance analysis.	None
npu	Optional[str]	The NPU device name to use for compilation.	None
tensor_parallel_size	Optional[int]	Size for tensor parallelism to distribute the model across devices.	None
timeout	Optional[int]	The timeout for the runtime in seconds. If it isn't provided, it will be set to 60 by default.	None
optimum_rbln_version	Optional[str]	The optimum-rbln version used for this configuration.	None
_torch_dtype	Optional[str]	The data type to use for the model.	None
_compile_cfgs	List[RBLNCompileConfig]	List of compilation configurations for the model.	[]
kwargs	Any	Additional keyword arguments.	{}

Raises:

Type	Description
ValueError	If unexpected keyword arguments are provided.

load(path, **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
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	RBLNModelConfig	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.