Common

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

Key Classes

• RBLNDiffusionMixin: The base mixin class for all RBLN-compatible diffusion pipelines
• RBLNDiffusionMixinConfig: Configuration class for RBLN diffusion pipelines

API Reference

Classes

RBLNDiffusionMixinConfig

Bases: RBLNModelConfig

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, optimum_rbln_version=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 if an NPU is available.	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
optimum_rbln_version	Optional[str]	The optimum-rbln version used for this configuration.	None
_compile_cfgs	List[RBLNCompileConfig]	List of compilation configurations for the model.	[]
**kwargs	Dict[str, Any]	Additional keyword arguments.	{}

Raises:

Type	Description
ValueError	If unexpected keyword arguments are provided.

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 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"])

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=False, 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 Hugging Face 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.	False
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	Dict[str, 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 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.