optimum-rbln 0.8.3a1__tar.gz → 0.8.3a3__tar.gz

{optimum_rbln-0.8.3a1 → optimum_rbln-0.8.3a3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.3a1
+Version: 0.8.3a3
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai

{optimum_rbln-0.8.3a1 → optimum_rbln-0.8.3a3}/src/optimum/rbln/__init__.py

@@ -81,6 +81,8 @@ _import_structure = {
         "RBLNDistilBertForQuestionAnsweringConfig",
         "RBLNDPTForDepthEstimation",
         "RBLNDPTForDepthEstimationConfig",
+        "RBLNDepthAnythingForDepthEstimationConfig",
+        "RBLNDepthAnythingForDepthEstimation",
         "RBLNExaoneForCausalLM",
         "RBLNExaoneForCausalLMConfig",
         "RBLNGemmaModel",
@@ -169,6 +171,9 @@ _import_structure = {
         "RBLNAutoencoderKLConfig",
         "RBLNAutoencoderKLCosmos",
         "RBLNAutoencoderKLCosmosConfig",
+        "RBLNAutoPipelineForImage2Image",
+        "RBLNAutoPipelineForInpainting",
+        "RBLNAutoPipelineForText2Image",
         "RBLNControlNetModel",
         "RBLNControlNetModelConfig",
         "RBLNCosmosTextToWorldPipeline",
@@ -238,6 +243,9 @@ if TYPE_CHECKING:
         RBLNAutoencoderKLConfig,
         RBLNAutoencoderKLCosmos,
         RBLNAutoencoderKLCosmosConfig,
+        RBLNAutoPipelineForImage2Image,
+        RBLNAutoPipelineForInpainting,
+        RBLNAutoPipelineForText2Image,
         RBLNControlNetModel,
         RBLNControlNetModelConfig,
         RBLNCosmosSafetyChecker,
@@ -346,6 +354,8 @@ if TYPE_CHECKING:
         RBLNDecoderOnlyModelConfig,
         RBLNDecoderOnlyModelForCausalLM,
         RBLNDecoderOnlyModelForCausalLMConfig,
+        RBLNDepthAnythingForDepthEstimation,
+        RBLNDepthAnythingForDepthEstimationConfig,
         RBLNDistilBertForQuestionAnswering,
         RBLNDistilBertForQuestionAnsweringConfig,
         RBLNDPTForDepthEstimation,

{optimum_rbln-0.8.3a1 → optimum_rbln-0.8.3a3}/src/optimum/rbln/__version__.py

@@ -1,7 +1,14 @@
 # file generated by setuptools-scm
 # don't change, don't track in version control
 
-__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
 
 TYPE_CHECKING = False
 if TYPE_CHECKING:
@@ -9,13 +16,19 @@ if TYPE_CHECKING:
     from typing import Union
 
     VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
 else:
     VERSION_TUPLE = object
+    COMMIT_ID = object
 
 version: str
 __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
 
-__version__ = version = '0.8.3a1'
-__version_tuple__ = version_tuple = (0, 8, 3, 'a1')
+__version__ = version = '0.8.3a3'
+__version_tuple__ = version_tuple = (0, 8, 3, 'a3')
+
+__commit_id__ = commit_id = None

{optimum_rbln-0.8.3a1 → optimum_rbln-0.8.3a3}/src/optimum/rbln/diffusers/__init__.py

@@ -59,6 +59,9 @@ _import_structure = {
         "RBLNVQModelConfig",
     ],
     "pipelines": [
+        "RBLNAutoPipelineForImage2Image",
+        "RBLNAutoPipelineForInpainting",
+        "RBLNAutoPipelineForText2Image",
         "RBLNCosmosTextToWorldPipeline",
         "RBLNCosmosVideoToWorldPipeline",
         "RBLNCosmosSafetyChecker",
@@ -144,6 +147,9 @@ if TYPE_CHECKING:
         RBLNVQModel,
     )
     from .pipelines import (
+        RBLNAutoPipelineForImage2Image,
+        RBLNAutoPipelineForInpainting,
+        RBLNAutoPipelineForText2Image,
         RBLNCosmosSafetyChecker,
         RBLNCosmosTextToWorldPipeline,
         RBLNCosmosVideoToWorldPipeline,

{optimum_rbln-0.8.3a1 → optimum_rbln-0.8.3a3}/src/optimum/rbln/diffusers/pipelines/__init__.py

@@ -18,6 +18,11 @@ from transformers.utils import _LazyModule
 
 
 _import_structure = {
+    "auto_pipeline": [
+        "RBLNAutoPipelineForImage2Image",
+        "RBLNAutoPipelineForInpainting",
+        "RBLNAutoPipelineForText2Image",
+    ],
     "controlnet": [
         "RBLNMultiControlNetModel",
         "RBLNStableDiffusionControlNetImg2ImgPipeline",
@@ -56,6 +61,11 @@ _import_structure = {
     ],
 }
 if TYPE_CHECKING:
+    from .auto_pipeline import (
+        RBLNAutoPipelineForImage2Image,
+        RBLNAutoPipelineForInpainting,
+        RBLNAutoPipelineForText2Image,
+    )
     from .controlnet import (
         RBLNMultiControlNetModel,
         RBLNStableDiffusionControlNetImg2ImgPipeline,

optimum_rbln-0.8.3a3/src/optimum/rbln/diffusers/pipelines/auto_pipeline.py

@@ -0,0 +1,237 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import importlib
+from typing import Type
+
+from diffusers.models.controlnets import ControlNetUnionModel
+from diffusers.pipelines.auto_pipeline import (
+    AUTO_IMAGE2IMAGE_PIPELINES_MAPPING,
+    AUTO_INPAINT_PIPELINES_MAPPING,
+    AUTO_TEXT2IMAGE_PIPELINES_MAPPING,
+    AutoPipelineForImage2Image,
+    AutoPipelineForInpainting,
+    AutoPipelineForText2Image,
+    _get_task_class,
+)
+from huggingface_hub.utils import validate_hf_hub_args
+
+from optimum.rbln.modeling_base import RBLNBaseModel
+from optimum.rbln.utils.model_utils import (
+    MODEL_MAPPING,
+    convert_hf_to_rbln_model_name,
+    convert_rbln_to_hf_model_name,
+    get_rbln_model_cls,
+)
+
+
+class RBLNAutoPipelineBase:
+    _model_mapping = None
+    _model_mapping_names = None
+
+    @classmethod
+    def get_rbln_cls(cls, pretrained_model_name_or_path, export=True, **kwargs):
+        if export:
+            hf_model_class = cls.infer_hf_model_class(pretrained_model_name_or_path, **kwargs)
+            rbln_class_name = convert_hf_to_rbln_model_name(hf_model_class.__name__)
+        else:
+            rbln_class_name = cls.get_rbln_model_cls_name(pretrained_model_name_or_path, **kwargs)
+            if convert_rbln_to_hf_model_name(rbln_class_name) not in cls._model_mapping_names:
+                raise ValueError(
+                    f"The architecture '{rbln_class_name}' is not supported by the `{cls.__name__}.from_pretrained()` method. "
+                    "Please use the `from_pretrained()` method of the appropriate class to load this model, "
+                    f"or directly use '{rbln_class_name}.from_pretrained()`."
+                )
+
+        try:
+            rbln_cls = get_rbln_model_cls(rbln_class_name)
+        except AttributeError as e:
+            raise AttributeError(
+                f"Class '{rbln_class_name}' not found in 'optimum.rbln' module for model ID '{pretrained_model_name_or_path}'. "
+                "Ensure that the class name is correctly mapped and available in the 'optimum.rbln' module."
+            ) from e
+
+        return rbln_cls
+
+    @classmethod
+    def get_rbln_model_cls_name(cls, pretrained_model_name_or_path, **kwargs):
+        """
+        Retrieve the path to the compiled model directory for a given RBLN model.
+
+        Args:
+            pretrained_model_name_or_path (str): Identifier of the model.
+
+        Returns:
+            str: Path to the compiled model directory.
+        """
+        model_index_config = cls.load_config(pretrained_model_name_or_path)
+
+        if "_class_name" not in model_index_config:
+            raise ValueError(
+                "The `_class_name` field is missing from model_index_config. This is unexpected and should be reported as an issue. "
+                "Please use the `from_pretrained()` method of the appropriate class to load this model."
+            )
+
+        return model_index_config["_class_name"]
+
+    @classmethod
+    def infer_hf_model_class(
+        cls,
+        pretrained_model_or_path,
+        cache_dir=None,
+        force_download=False,
+        proxies=None,
+        token=None,
+        local_files_only=False,
+        revision=None,
+        **kwargs,
+    ):
+        config = cls.load_config(
+            pretrained_model_or_path,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            proxies=proxies,
+            token=token,
+            local_files_only=local_files_only,
+            revision=revision,
+        )
+        pipeline_key_name = cls.get_pipeline_key_name(config, **kwargs)
+
+        pipeline_cls = _get_task_class(cls._model_mapping, pipeline_key_name)
+
+        return pipeline_cls
+
+    @classmethod
+    def get_pipeline_key_name(cls, config, **kwargs):
+        orig_class_name = config["_class_name"]
+        if "ControlPipeline" in orig_class_name:
+            to_replace = "ControlPipeline"
+        else:
+            to_replace = "Pipeline"
+
+        if "controlnet" in kwargs:
+            if isinstance(kwargs["controlnet"], ControlNetUnionModel):
+                orig_class_name = config["_class_name"].replace(to_replace, "ControlNetUnionPipeline")
+            else:
+                orig_class_name = config["_class_name"].replace(to_replace, "ControlNetPipeline")
+        if "enable_pag" in kwargs:
+            enable_pag = kwargs.pop("enable_pag")
+            if enable_pag:
+                orig_class_name = orig_class_name.replace(to_replace, "PAGPipeline")
+
+        return orig_class_name
+
+    @classmethod
+    @validate_hf_hub_args
+    def from_pretrained(cls, model_id, **kwargs):
+        rbln_cls = cls.get_rbln_cls(model_id, **kwargs)
+        return rbln_cls.from_pretrained(model_id, **kwargs)
+
+    @classmethod
+    def from_model(cls, model, **kwargs):
+        rbln_cls = get_rbln_model_cls(f"RBLN{model.__class__.__name__}")
+        return rbln_cls.from_model(model, **kwargs)
+
+    @staticmethod
+    def register(rbln_cls: Type[RBLNBaseModel], exist_ok=False):
+        """
+        Register a new RBLN model class.
+
+        Args:
+            rbln_cls (Type[RBLNBaseModel]): The RBLN model class to register.
+            exist_ok (bool): Whether to allow registering an already registered model.
+        """
+        if not issubclass(rbln_cls, RBLNBaseModel):
+            raise ValueError("`rbln_cls` must be a subclass of RBLNBaseModel.")
+
+        native_cls = getattr(importlib.import_module("optimum.rbln"), rbln_cls.__name__, None)
+        if rbln_cls.__name__ in MODEL_MAPPING or native_cls is not None:
+            if not exist_ok:
+                raise ValueError(f"Model for {rbln_cls.__name__} already registered.")
+
+        MODEL_MAPPING[rbln_cls.__name__] = rbln_cls
+
+
+class RBLNAutoPipelineForText2Image(RBLNAutoPipelineBase, AutoPipelineForText2Image):
+    _model_mapping = AUTO_TEXT2IMAGE_PIPELINES_MAPPING
+    _model_mapping_names = {x[0]: x[1].__name__ for x in AUTO_TEXT2IMAGE_PIPELINES_MAPPING.items()}
+
+
+class RBLNAutoPipelineForImage2Image(RBLNAutoPipelineBase, AutoPipelineForImage2Image):
+    _model_mapping = AUTO_IMAGE2IMAGE_PIPELINES_MAPPING
+    _model_mapping_names = {x[0]: x[1].__name__ for x in AUTO_IMAGE2IMAGE_PIPELINES_MAPPING.items()}
+
+    @classmethod
+    def get_pipeline_key_name(cls, config, **kwargs):
+        orig_class_name = config["_class_name"]
+        # the `orig_class_name` can be:
+        # `- *Pipeline` (for regular text-to-image checkpoint)
+        #  - `*ControlPipeline` (for Flux tools specific checkpoint)
+        # `- *Img2ImgPipeline` (for refiner checkpoint)
+        if "Img2Img" in orig_class_name:
+            to_replace = "Img2ImgPipeline"
+        elif "ControlPipeline" in orig_class_name:
+            to_replace = "ControlPipeline"
+        else:
+            to_replace = "Pipeline"
+
+        if "controlnet" in kwargs:
+            if isinstance(kwargs["controlnet"], ControlNetUnionModel):
+                orig_class_name = orig_class_name.replace(to_replace, "ControlNetUnion" + to_replace)
+            else:
+                orig_class_name = orig_class_name.replace(to_replace, "ControlNet" + to_replace)
+        if "enable_pag" in kwargs:
+            enable_pag = kwargs.pop("enable_pag")
+            if enable_pag:
+                orig_class_name = orig_class_name.replace(to_replace, "PAG" + to_replace)
+
+        if to_replace == "ControlPipeline":
+            orig_class_name = orig_class_name.replace(to_replace, "ControlImg2ImgPipeline")
+
+        return orig_class_name
+
+
+class RBLNAutoPipelineForInpainting(RBLNAutoPipelineBase, AutoPipelineForInpainting):
+    _model_mapping = AUTO_INPAINT_PIPELINES_MAPPING
+    _model_mapping_names = {x[0]: x[1].__name__ for x in AUTO_INPAINT_PIPELINES_MAPPING.items()}
+
+    @classmethod
+    def get_pipeline_key_name(cls, config, **kwargs):
+        orig_class_name = config["_class_name"]
+
+        # The `orig_class_name`` can be:
+        # `- *InpaintPipeline` (for inpaint-specific checkpoint)
+        #  - `*ControlPipeline` (for Flux tools specific checkpoint)
+        #  - or *Pipeline (for regular text-to-image checkpoint)
+        if "Inpaint" in orig_class_name:
+            to_replace = "InpaintPipeline"
+        elif "ControlPipeline" in orig_class_name:
+            to_replace = "ControlPipeline"
+        else:
+            to_replace = "Pipeline"
+
+        if "controlnet" in kwargs:
+            if isinstance(kwargs["controlnet"], ControlNetUnionModel):
+                orig_class_name = orig_class_name.replace(to_replace, "ControlNetUnion" + to_replace)
+            else:
+                orig_class_name = orig_class_name.replace(to_replace, "ControlNet" + to_replace)
+        if "enable_pag" in kwargs:
+            enable_pag = kwargs.pop("enable_pag")
+            if enable_pag:
+                orig_class_name = orig_class_name.replace(to_replace, "PAG" + to_replace)
+        if to_replace == "ControlPipeline":
+            orig_class_name = orig_class_name.replace(to_replace, "ControlInpaintPipeline")
+
+        return orig_class_name

{optimum_rbln-0.8.3a1 → optimum_rbln-0.8.3a3}/src/optimum/rbln/transformers/__init__.py

@@ -69,6 +69,8 @@ _import_structure = {
         "RBLNDistilBertForQuestionAnsweringConfig",
         "RBLNDPTForDepthEstimation",
         "RBLNDPTForDepthEstimationConfig",
+        "RBLNDepthAnythingForDepthEstimation",
+        "RBLNDepthAnythingForDepthEstimationConfig",
         "RBLNExaoneForCausalLM",
         "RBLNExaoneForCausalLMConfig",
         "RBLNGemmaModel",
@@ -204,6 +206,8 @@ if TYPE_CHECKING:
         RBLNDecoderOnlyModelConfig,
         RBLNDecoderOnlyModelForCausalLM,
         RBLNDecoderOnlyModelForCausalLMConfig,
+        RBLNDepthAnythingForDepthEstimation,
+        RBLNDepthAnythingForDepthEstimationConfig,
         RBLNDistilBertForQuestionAnswering,
         RBLNDistilBertForQuestionAnsweringConfig,
         RBLNDPTForDepthEstimation,

{optimum_rbln-0.8.3a1 → optimum_rbln-0.8.3a3}/src/optimum/rbln/transformers/models/__init__.py

@@ -90,6 +90,7 @@ _import_structure = {
         "RBLNDecoderOnlyModelForCausalLM",
         "RBLNDecoderOnlyModelForCausalLMConfig",
     ],
+    "depth_anything": ["RBLNDepthAnythingForDepthEstimationConfig", "RBLNDepthAnythingForDepthEstimation"],
     "dpt": [
         "RBLNDPTForDepthEstimation",
         "RBLNDPTForDepthEstimationConfig",
@@ -219,6 +220,7 @@ if TYPE_CHECKING:
         RBLNDecoderOnlyModelForCausalLM,
         RBLNDecoderOnlyModelForCausalLMConfig,
     )
+    from .depth_anything import RBLNDepthAnythingForDepthEstimation, RBLNDepthAnythingForDepthEstimationConfig
     from .distilbert import RBLNDistilBertForQuestionAnswering, RBLNDistilBertForQuestionAnsweringConfig
     from .dpt import RBLNDPTForDepthEstimation, RBLNDPTForDepthEstimationConfig
     from .exaone import RBLNExaoneForCausalLM, RBLNExaoneForCausalLMConfig

optimum_rbln-0.8.3a3/src/optimum/rbln/transformers/models/depth_anything/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .configuration_depth_anything import RBLNDepthAnythingForDepthEstimationConfig
+from .modeling_depth_anything import RBLNDepthAnythingForDepthEstimation

optimum_rbln-0.8.3a3/src/optimum/rbln/transformers/models/depth_anything/configuration_depth_anything.py

@@ -0,0 +1,24 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...configuration_generic import RBLNModelForDepthEstimationConfig
+
+
+class RBLNDepthAnythingForDepthEstimationConfig(RBLNModelForDepthEstimationConfig):
+    """
+    Configuration class for DepthAnythingForDepthEstimation.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Depth Anything V2 Small models for depth estimation tasks.
+    """

optimum_rbln-0.8.3a3/src/optimum/rbln/transformers/models/depth_anything/modeling_depth_anything.py

@@ -0,0 +1,25 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...modeling_generic import RBLNModelForDepthEstimation
+
+
+class RBLNDepthAnythingForDepthEstimation(RBLNModelForDepthEstimation):
+    """
+    RBLN optimized DepthAnythingForDepthEstimation model for depth estimation tasks.
+
+    This class provides hardware-accelerated inference for Depth Anything V2 Small
+    models on RBLN devices, providing the most capable monocular depth estimation (MDE) model.
+    """

{optimum_rbln-0.8.3a1 → optimum_rbln-0.8.3a3}/tests/test_diffusers.py

@@ -4,6 +4,9 @@ import torch
 from diffusers import ControlNetModel
 
 from optimum.rbln import (
+    RBLNAutoPipelineForImage2Image,
+    # RBLNAutoPipelineForInpainting, FIXME: add inpainting tests
+    RBLNAutoPipelineForText2Image,
     RBLNKandinskyV22CombinedPipeline,
     RBLNKandinskyV22Img2ImgCombinedPipeline,
     RBLNStableDiffusion3Img2ImgPipeline,
@@ -19,6 +22,7 @@ from .test_base import BaseHubTest, BaseTest
 
 
 class TestSDModel(BaseTest.TestModel, BaseHubTest.TestHub):
+    RBLN_AUTO_CLASS = RBLNAutoPipelineForText2Image
     RBLN_CLASS = RBLNStableDiffusionPipeline
     HF_MODEL_ID = "hf-internal-testing/tiny-sd-pipe"
     GENERATION_KWARGS = {
@@ -40,6 +44,7 @@ class TestSDModel(BaseTest.TestModel, BaseHubTest.TestHub):
 
 
 class TestSDModelBatch(BaseTest.TestModel):
+    RBLN_AUTO_CLASS = RBLNAutoPipelineForText2Image
     RBLN_CLASS = RBLNStableDiffusionPipeline
     HF_MODEL_ID = "hf-internal-testing/tiny-sd-pipe"
     GENERATION_KWARGS = {
@@ -61,6 +66,7 @@ class TestSDModelBatch(BaseTest.TestModel):
 
 
 class TestSDXLModel(BaseTest.TestModel):
+    RBLN_AUTO_CLASS = RBLNAutoPipelineForText2Image
     RBLN_CLASS = RBLNStableDiffusionXLPipeline
     HF_MODEL_ID = "hf-internal-testing/tiny-sdxl-pipe"
     GENERATION_KWARGS = {
@@ -80,6 +86,7 @@ class TestSDXLModel(BaseTest.TestModel):
 
 
 class TestSDImg2ImgModel(BaseTest.TestModel):
+    RBLN_AUTO_CLASS = RBLNAutoPipelineForImage2Image
     RBLN_CLASS = RBLNStableDiffusionImg2ImgPipeline
     HF_MODEL_ID = "hf-internal-testing/tiny-sd-pipe"
     GENERATION_KWARGS = {
@@ -105,6 +112,7 @@ class TestSDImg2ImgModel(BaseTest.TestModel):
 
 
 class TestSDControlNetModel(BaseTest.TestModel):
+    RBLN_AUTO_CLASS = RBLNAutoPipelineForText2Image
     RBLN_CLASS = RBLNStableDiffusionControlNetPipeline
     HF_MODEL_ID = "hf-internal-testing/tiny-stable-diffusion-torch"
     CONTROLNET_ID = "hf-internal-testing/tiny-controlnet"
@@ -134,6 +142,7 @@ class TestSDControlNetModel(BaseTest.TestModel):
 
 
 class TestSDXLControlNetModel(BaseTest.TestModel):
+    RBLN_AUTO_CLASS = RBLNAutoPipelineForText2Image
     RBLN_CLASS = RBLNStableDiffusionXLControlNetPipeline
     HF_MODEL_ID = "hf-internal-testing/tiny-sdxl-pipe"
     CONTROLNET_ID = "hf-internal-testing/tiny-controlnet-sdxl"
@@ -164,6 +173,7 @@ class TestSDXLControlNetModel(BaseTest.TestModel):
 
 
 class TestSD3Model(BaseTest.TestModel):
+    RBLN_AUTO_CLASS = RBLNAutoPipelineForText2Image
     RBLN_CLASS = RBLNStableDiffusion3Pipeline
     HF_MODEL_ID = "hf-internal-testing/tiny-sd3-pipe"
     GENERATION_KWARGS = {
@@ -181,6 +191,7 @@ class TestSD3Model(BaseTest.TestModel):
 
 
 class TestSD3Img2ImgModel(BaseTest.TestModel):
+    RBLN_AUTO_CLASS = RBLNAutoPipelineForImage2Image
     RBLN_CLASS = RBLNStableDiffusion3Img2ImgPipeline
     HF_MODEL_ID = "hf-internal-testing/tiny-sd3-pipe"
     GENERATION_KWARGS = {
@@ -198,6 +209,7 @@ class TestSD3Img2ImgModel(BaseTest.TestModel):
 
 
 class TestSDMultiControlNetModel(BaseTest.TestModel):
+    RBLN_AUTO_CLASS = RBLNAutoPipelineForText2Image
     RBLN_CLASS = RBLNStableDiffusionControlNetPipeline
     HF_MODEL_ID = "hf-internal-testing/tiny-stable-diffusion-torch"
     CONTROLNET_ID = "hf-internal-testing/tiny-controlnet"
@@ -234,6 +246,7 @@ class TestSDMultiControlNetModel(BaseTest.TestModel):
 
 
 class TestKandinskyV22Model(BaseTest.TestModel):
+    RBLN_AUTO_CLASS = RBLNAutoPipelineForText2Image
     RBLN_CLASS = RBLNKandinskyV22CombinedPipeline
     HF_MODEL_ID = "hf-internal-testing/tiny-random-kandinsky-v22-decoder"
     GENERATION_KWARGS = {
@@ -280,6 +293,7 @@ class TestKandinskyV22Model(BaseTest.TestModel):
 
 
 class TestKandinskyV22Img2ImgModel(BaseTest.TestModel):
+    RBLN_AUTO_CLASS = RBLNAutoPipelineForImage2Image
     RBLN_CLASS = RBLNKandinskyV22Img2ImgCombinedPipeline
     HF_MODEL_ID = "hf-internal-testing/tiny-random-kandinsky-v22-decoder"