optimum-rbln 0.7.4a8__py3-none-any.whl → 0.7.5a0__py3-none-any.whl

optimum/rbln/__init__.py

@@ -337,6 +337,7 @@ else:
         globals()["__file__"],
         _import_structure,
         module_spec=__spec__,
+        extra_objects={"__version__": __version__},
     )
 
 

optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.7.4a8'
-__version_tuple__ = version_tuple = (0, 7, 4, 'a8')
+__version__ = version = '0.7.5a0'
+__version_tuple__ = version_tuple = (0, 7, 5, 'a0')

optimum/rbln/configuration_utils.py

@@ -19,7 +19,6 @@ from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
-import rebel
 import torch
 
 from .__version__ import __version__
@@ -452,6 +451,7 @@ class RBLNModelConfig:
         "activate_profiler",
     ]
     submodules: List[str] = []
+    subclass_non_save_attributes = []
 
     def init_submodule_config(
         self,
@@ -480,7 +480,11 @@ class RBLNModelConfig:
         return submodule_config
 
     def __setattr__(self, key, value):
-        if key != "_attributes_map" and key not in self.non_save_attributes:
+        if (
+            key != "_attributes_map"
+            and key not in self.non_save_attributes
+            and key not in self.subclass_non_save_attributes
+        ):
             self._attributes_map[key] = value
 
         if hasattr(self, "_frozen") and self._frozen:
@@ -528,7 +532,7 @@ class RBLNModelConfig:
 
         Args:
             cls_name (Optional[str]): The class name of the configuration. Defaults to the current class name.
-            create_runtimes (Optional[bool]): Whether to create RBLN runtimes. Defaults to True if an NPU is available.
+            create_runtimes (Optional[bool]): Whether to create RBLN runtimes. Defaults to True.
             optimize_host_memory (Optional[bool]): Whether to optimize host memory usage. Defaults to True.
             device (Optional[Union[int, List[int]]]): The device(s) to load the model onto. Can be a single device ID or a list.
             device_map (Optional[Dict[str, Union[int, List[int]]]]): Mapping from compiled model names to device IDs.
@@ -723,13 +727,35 @@ class RBLNModelConfig:
 
         return rbln_config, kwargs
 
+    def get_default_values_for_original_cls(self, func_name: str, keys: List[str]) -> Dict[str, Any]:
+        """
+        Get default values for original class attributes from RBLNModelConfig.
+
+        Args:
+            func_name (str): The name of the function to get the default values for.
+            keys (List[str]): The keys of the attributes to get.
+
+        Returns:
+            Dict[str, Any]: The default values for the attributes.
+        """
+        model_cls = self.rbln_model_cls.get_hf_class()
+        func = getattr(model_cls, func_name)
+        func_signature = inspect.signature(func)
+        default_values = {}
+        for key in keys:
+            if key in func_signature.parameters:
+                default_values[key] = func_signature.parameters[key].default
+            else:
+                raise ValueError(f"Default value for `{key}` is not set for the model class.")
+        return default_values
+
     @property
     def create_runtimes(self):
         context = ContextRblnConfig.get_current_context()["create_runtimes"]
         if context is not None:
             return context
         elif self._runtime_options["create_runtimes"] is None:
-            return rebel.npu_is_available()
+            return True
         return self._runtime_options["create_runtimes"]
 
     @create_runtimes.setter

optimum/rbln/diffusers/configurations/models/configuration_controlnet.py

@@ -18,6 +18,8 @@ from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNControlNetModelConfig(RBLNModelConfig):
+    subclass_non_save_attributes = ["_batch_size_is_specified"]
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
@@ -44,6 +46,8 @@ class RBLNControlNetModelConfig(RBLNModelConfig):
             ValueError: If batch_size is not a positive integer.
         """
         super().__init__(**kwargs)
+        self._batch_size_is_specified = batch_size is not None
+
         self.batch_size = batch_size or 1
         if not isinstance(self.batch_size, int) or self.batch_size < 0:
             raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
@@ -52,3 +56,7 @@ class RBLNControlNetModelConfig(RBLNModelConfig):
         self.unet_sample_size = unet_sample_size
         self.vae_sample_size = vae_sample_size
         self.text_model_hidden_size = text_model_hidden_size
+
+    @property
+    def batch_size_is_specified(self):
+        return self._batch_size_is_specified

optimum/rbln/diffusers/configurations/models/configuration_prior_transformer.py

@@ -18,6 +18,8 @@ from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNPriorTransformerConfig(RBLNModelConfig):
+    subclass_non_save_attributes = ["_batch_size_is_specified"]
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
@@ -36,9 +38,15 @@ class RBLNPriorTransformerConfig(RBLNModelConfig):
             ValueError: If batch_size is not a positive integer.
         """
         super().__init__(**kwargs)
+        self._batch_size_is_specified = batch_size is not None
+
         self.batch_size = batch_size or 1
         if not isinstance(self.batch_size, int) or self.batch_size < 0:
             raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
 
         self.embedding_dim = embedding_dim
         self.num_embeddings = num_embeddings
+
+    @property
+    def batch_size_is_specified(self):
+        return self._batch_size_is_specified

optimum/rbln/diffusers/configurations/models/configuration_transformer_sd3.py

@@ -18,6 +18,8 @@ from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNSD3Transformer2DModelConfig(RBLNModelConfig):
+    subclass_non_save_attributes = ["_batch_size_is_specified"]
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
@@ -38,6 +40,8 @@ class RBLNSD3Transformer2DModelConfig(RBLNModelConfig):
             ValueError: If batch_size is not a positive integer.
         """
         super().__init__(**kwargs)
+        self._batch_size_is_specified = batch_size is not None
+
         self.batch_size = batch_size or 1
         if not isinstance(self.batch_size, int) or self.batch_size < 0:
             raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
@@ -46,3 +50,7 @@ class RBLNSD3Transformer2DModelConfig(RBLNModelConfig):
         self.sample_size = sample_size
         if isinstance(self.sample_size, int):
             self.sample_size = (self.sample_size, self.sample_size)
+
+    @property
+    def batch_size_is_specified(self):
+        return self._batch_size_is_specified

optimum/rbln/diffusers/configurations/models/configuration_unet_2d_condition.py

@@ -18,6 +18,8 @@ from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNUNet2DConditionModelConfig(RBLNModelConfig):
+    subclass_non_save_attributes = ["_batch_size_is_specified"]
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
@@ -49,6 +51,8 @@ class RBLNUNet2DConditionModelConfig(RBLNModelConfig):
             ValueError: If batch_size is not a positive integer.
         """
         super().__init__(**kwargs)
+        self._batch_size_is_specified = batch_size is not None
+
         self.batch_size = batch_size or 1
         if not isinstance(self.batch_size, int) or self.batch_size < 0:
             raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
@@ -64,3 +68,7 @@ class RBLNUNet2DConditionModelConfig(RBLNModelConfig):
         self.sample_size = sample_size
         if isinstance(sample_size, int):
             self.sample_size = (sample_size, sample_size)
+
+    @property
+    def batch_size_is_specified(self):
+        return self._batch_size_is_specified

optimum/rbln/diffusers/configurations/pipelines/configuration_controlnet.py

@@ -16,13 +16,9 @@ from typing import Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 from ....transformers import RBLNCLIPTextModelConfig, RBLNCLIPTextModelWithProjectionConfig
-from ....utils.logging import get_logger
 from ..models import RBLNAutoencoderKLConfig, RBLNControlNetModelConfig, RBLNUNet2DConditionModelConfig
 
 
-logger = get_logger(__name__)
-
-
 class _RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
     submodules = ["text_encoder", "unet", "vae", "controlnet"]
     _vae_uses_encoder = False
@@ -58,7 +54,7 @@ class _RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
             sample_size (Optional[Tuple[int, int]]): Spatial dimensions for the UNet model.
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
-            guidance_scale (Optional[float]): Scale for classifier-free guidance. Deprecated parameter.
+            guidance_scale (Optional[float]): Scale for classifier-free guidance.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
@@ -79,7 +75,6 @@ class _RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
         self.unet = self.init_submodule_config(
             RBLNUNet2DConditionModelConfig,
             unet,
-            batch_size=batch_size,
             sample_size=sample_size,
         )
         self.vae = self.init_submodule_config(
@@ -89,14 +84,24 @@ class _RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
             uses_encoder=self.__class__._vae_uses_encoder,
             sample_size=image_size,  # image size is equal to sample size in vae
         )
-        self.controlnet = self.init_submodule_config(RBLNControlNetModelConfig, controlnet, batch_size=batch_size)
+        self.controlnet = self.init_submodule_config(RBLNControlNetModelConfig, controlnet)
+
+        # Get default guidance scale from original class to set UNet and ControlNet batch size
+        if guidance_scale is None:
+            guidance_scale = self.get_default_values_for_original_cls("__call__", ["guidance_scale"])["guidance_scale"]
 
         if guidance_scale is not None:
-            logger.warning("Specifying `guidance_scale` is deprecated. It will be removed in a future version.")
             do_classifier_free_guidance = guidance_scale > 1.0
             if do_classifier_free_guidance:
-                self.unet.batch_size = self.text_encoder.batch_size * 2
-                self.controlnet.batch_size = self.text_encoder.batch_size * 2
+                if not self.unet.batch_size_is_specified:
+                    self.unet.batch_size = self.text_encoder.batch_size * 2
+                if not self.controlnet.batch_size_is_specified:
+                    self.controlnet.batch_size = self.text_encoder.batch_size * 2
+            else:
+                if not self.unet.batch_size_is_specified:
+                    self.unet.batch_size = self.text_encoder.batch_size
+                if not self.controlnet.batch_size_is_specified:
+                    self.controlnet.batch_size = self.text_encoder.batch_size
 
     @property
     def batch_size(self):
@@ -157,7 +162,7 @@ class _RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
             sample_size (Optional[Tuple[int, int]]): Spatial dimensions for the UNet model.
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
-            guidance_scale (Optional[float]): Scale for classifier-free guidance. Deprecated parameter.
+            guidance_scale (Optional[float]): Scale for classifier-free guidance.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
@@ -181,7 +186,6 @@ class _RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
         self.unet = self.init_submodule_config(
             RBLNUNet2DConditionModelConfig,
             unet,
-            batch_size=batch_size,
             sample_size=sample_size,
         )
         self.vae = self.init_submodule_config(
@@ -191,14 +195,25 @@ class _RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
             uses_encoder=self.__class__._vae_uses_encoder,
             sample_size=image_size,  # image size is equal to sample size in vae
         )
-        self.controlnet = self.init_submodule_config(RBLNControlNetModelConfig, controlnet, batch_size=batch_size)
+        self.controlnet = self.init_submodule_config(RBLNControlNetModelConfig, controlnet)
 
-        if guidance_scale is not None:
-            logger.warning("Specifying `guidance_scale` is deprecated. It will be removed in a future version.")
-            do_classifier_free_guidance = guidance_scale > 1.0
-            if do_classifier_free_guidance:
+        # Get default guidance scale from original class to set UNet and ControlNet batch size
+        guidance_scale = (
+            guidance_scale
+            or self.get_default_values_for_original_cls("__call__", ["guidance_scale"])["guidance_scale"]
+        )
+
+        do_classifier_free_guidance = guidance_scale > 1.0
+        if do_classifier_free_guidance:
+            if not self.unet.batch_size_is_specified:
                 self.unet.batch_size = self.text_encoder.batch_size * 2
+            if not self.controlnet.batch_size_is_specified:
                 self.controlnet.batch_size = self.text_encoder.batch_size * 2
+        else:
+            if not self.unet.batch_size_is_specified:
+                self.unet.batch_size = self.text_encoder.batch_size
+            if not self.controlnet.batch_size_is_specified:
+                self.controlnet.batch_size = self.text_encoder.batch_size
 
     @property
     def batch_size(self):

optimum/rbln/diffusers/configurations/pipelines/configuration_kandinsky2_2.py

@@ -16,14 +16,10 @@ from typing import Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 from ....transformers import RBLNCLIPTextModelWithProjectionConfig, RBLNCLIPVisionModelWithProjectionConfig
-from ....utils.logging import get_logger
 from ..models import RBLNUNet2DConditionModelConfig, RBLNVQModelConfig
 from ..models.configuration_prior_transformer import RBLNPriorTransformerConfig
 
 
-logger = get_logger(__name__)
-
-
 class _RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
     submodules = ["unet", "movq"]
     _movq_uses_encoder = False
@@ -49,7 +45,7 @@ class _RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
                 Initialized as RBLNVQModelConfig if not provided.
             sample_size (Optional[Tuple[int, int]]): Spatial dimensions for the UNet model.
             batch_size (Optional[int]): Batch size for inference, applied to all submodules.
-            guidance_scale (Optional[float]): Scale for classifier-free guidance. Deprecated parameter.
+            guidance_scale (Optional[float]): Scale for classifier-free guidance.
             image_size (Optional[Tuple[int, int]]): Dimensions for the generated images.
                 Cannot be used together with img_height/img_width.
             img_height (Optional[int]): Height of the generated images.
@@ -70,9 +66,7 @@ class _RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
         if img_height is not None and img_width is not None:
             image_size = (img_height, img_width)
 
-        self.unet = self.init_submodule_config(
-            RBLNUNet2DConditionModelConfig, unet, batch_size=batch_size, sample_size=sample_size
-        )
+        self.unet = self.init_submodule_config(RBLNUNet2DConditionModelConfig, unet, sample_size=sample_size)
         self.movq = self.init_submodule_config(
             RBLNVQModelConfig,
             movq,
@@ -81,11 +75,16 @@ class _RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
             uses_encoder=self._movq_uses_encoder,
         )
 
-        if guidance_scale is not None:
-            logger.warning("Specifying `guidance_scale` is deprecated. It will be removed in a future version.")
+        # Get default guidance scale from original class to set UNet batch size
+        if guidance_scale is None:
+            guidance_scale = self.get_default_values_for_original_cls("__call__", ["guidance_scale"])["guidance_scale"]
+
+        if not self.unet.batch_size_is_specified:
             do_classifier_free_guidance = guidance_scale > 1.0
             if do_classifier_free_guidance:
                 self.unet.batch_size = self.movq.batch_size * 2
+            else:
+                self.unet.batch_size = self.movq.batch_size
 
     @property
     def batch_size(self):
@@ -136,7 +135,7 @@ class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
             prior (Optional[RBLNPriorTransformerConfig]): Configuration for the prior transformer component.
                 Initialized as RBLNPriorTransformerConfig if not provided.
             batch_size (Optional[int]): Batch size for inference, applied to all submodules.
-            guidance_scale (Optional[float]): Scale for classifier-free guidance. Deprecated parameter.
+            guidance_scale (Optional[float]): Scale for classifier-free guidance.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Note:
@@ -151,13 +150,18 @@ class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
             RBLNCLIPVisionModelWithProjectionConfig, image_encoder, batch_size=batch_size
         )
 
-        self.prior = self.init_submodule_config(RBLNPriorTransformerConfig, prior, batch_size=batch_size)
+        self.prior = self.init_submodule_config(RBLNPriorTransformerConfig, prior)
+
+        # Get default guidance scale from original class to set UNet batch size
+        if guidance_scale is None:
+            guidance_scale = self.get_default_values_for_original_cls("__call__", ["guidance_scale"])["guidance_scale"]
 
-        if guidance_scale is not None:
-            logger.warning("Specifying `guidance_scale` is deprecated. It will be removed in a future version.")
+        if not self.prior.batch_size_is_specified:
             do_classifier_free_guidance = guidance_scale > 1.0
             if do_classifier_free_guidance:
                 self.prior.batch_size = self.text_encoder.batch_size * 2
+            else:
+                self.prior.batch_size = self.text_encoder.batch_size
 
     @property
     def batch_size(self):
@@ -208,7 +212,7 @@ class _RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
             batch_size (Optional[int]): Batch size for inference, applied to all submodules.
             img_height (Optional[int]): Height of the generated images.
             img_width (Optional[int]): Width of the generated images.
-            guidance_scale (Optional[float]): Scale for classifier-free guidance. Deprecated parameter.
+            guidance_scale (Optional[float]): Scale for classifier-free guidance.
             prior_prior (Optional[RBLNPriorTransformerConfig]): Direct configuration for the prior transformer.
                 Used if prior_pipe is not provided.
             prior_image_encoder (Optional[RBLNCLIPVisionModelWithProjectionConfig]): Direct configuration for the image encoder.

optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion.py

@@ -16,13 +16,9 @@ from typing import Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 from ....transformers import RBLNCLIPTextModelConfig
-from ....utils.logging import get_logger
 from ..models import RBLNAutoencoderKLConfig, RBLNUNet2DConditionModelConfig
 
 
-logger = get_logger(__name__)
-
-
 class _RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
     submodules = ["text_encoder", "unet", "vae"]
     _vae_uses_encoder = False
@@ -55,7 +51,7 @@ class _RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
             sample_size (Optional[Tuple[int, int]]): Spatial dimensions for the UNet model.
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
-            guidance_scale (Optional[float]): Scale for classifier-free guidance. Deprecated parameter.
+            guidance_scale (Optional[float]): Scale for classifier-free guidance.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
@@ -76,7 +72,6 @@ class _RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
         self.unet = self.init_submodule_config(
             RBLNUNet2DConditionModelConfig,
             unet,
-            batch_size=batch_size,
             sample_size=sample_size,
         )
         self.vae = self.init_submodule_config(
@@ -87,11 +82,16 @@ class _RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
             sample_size=image_size,  # image size is equal to sample size in vae
         )
 
-        if guidance_scale is not None:
-            logger.warning("Specifying `guidance_scale` is deprecated. It will be removed in a future version.")
+        # Get default guidance scale from original class to set UNet batch size
+        if guidance_scale is None:
+            guidance_scale = self.get_default_values_for_original_cls("__call__", ["guidance_scale"])["guidance_scale"]
+
+        if not self.unet.batch_size_is_specified:
             do_classifier_free_guidance = guidance_scale > 1.0
             if do_classifier_free_guidance:
                 self.unet.batch_size = self.text_encoder.batch_size * 2
+            else:
+                self.unet.batch_size = self.text_encoder.batch_size
 
     @property
     def batch_size(self):

optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.py

@@ -16,13 +16,9 @@ from typing import Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 from ....transformers import RBLNCLIPTextModelWithProjectionConfig, RBLNT5EncoderModelConfig
-from ....utils.logging import get_logger
 from ..models import RBLNAutoencoderKLConfig, RBLNSD3Transformer2DModelConfig
 
 
-logger = get_logger(__name__)
-
-
 class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
     submodules = ["transformer", "text_encoder", "text_encoder_2", "text_encoder_3", "vae"]
     _vae_uses_encoder = False
@@ -63,7 +59,7 @@ class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
             batch_size (Optional[int]): Batch size for inference, applied to all submodules.
             img_height (Optional[int]): Height of the generated images.
             img_width (Optional[int]): Width of the generated images.
-            guidance_scale (Optional[float]): Scale for classifier-free guidance. Deprecated parameter.
+            guidance_scale (Optional[float]): Scale for classifier-free guidance.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
@@ -97,7 +93,6 @@ class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
         self.transformer = self.init_submodule_config(
             RBLNSD3Transformer2DModelConfig,
             transformer,
-            batch_size=batch_size,
             sample_size=sample_size,
         )
         self.vae = self.init_submodule_config(
@@ -108,11 +103,16 @@ class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
             sample_size=image_size,
         )
 
-        if guidance_scale is not None:
-            logger.warning("Specifying `guidance_scale` is deprecated. It will be removed in a future version.")
+        # Get default guidance scale from original class to set Transformer batch size
+        if guidance_scale is None:
+            guidance_scale = self.get_default_values_for_original_cls("__call__", ["guidance_scale"])["guidance_scale"]
+
+        if not self.transformer.batch_size_is_specified:
             do_classifier_free_guidance = guidance_scale > 1.0
             if do_classifier_free_guidance:
                 self.transformer.batch_size = self.text_encoder.batch_size * 2
+            else:
+                self.transformer.batch_size = self.text_encoder.batch_size
 
     @property
     def max_seq_len(self):

optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_xl.py

@@ -16,13 +16,9 @@ from typing import Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 from ....transformers import RBLNCLIPTextModelConfig, RBLNCLIPTextModelWithProjectionConfig
-from ....utils.logging import get_logger
 from ..models import RBLNAutoencoderKLConfig, RBLNUNet2DConditionModelConfig
 
 
-logger = get_logger(__name__)
-
-
 class _RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
     submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]
     _vae_uses_encoder = False
@@ -58,7 +54,7 @@ class _RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
             sample_size (Optional[Tuple[int, int]]): Spatial dimensions for the UNet model.
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
-            guidance_scale (Optional[float]): Scale for classifier-free guidance. Deprecated parameter.
+            guidance_scale (Optional[float]): Scale for classifier-free guidance.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
@@ -82,7 +78,6 @@ class _RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
         self.unet = self.init_submodule_config(
             RBLNUNet2DConditionModelConfig,
             unet,
-            batch_size=batch_size,
             sample_size=sample_size,
         )
         self.vae = self.init_submodule_config(
@@ -93,11 +88,16 @@ class _RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
             sample_size=image_size,  # image size is equal to sample size in vae
         )
 
-        if guidance_scale is not None:
-            logger.warning("Specifying `guidance_scale` is deprecated. It will be removed in a future version.")
+        # Get default guidance scale from original class to set UNet batch size
+        if guidance_scale is None:
+            guidance_scale = self.get_default_values_for_original_cls("__call__", ["guidance_scale"])["guidance_scale"]
+
+        if not self.unet.batch_size_is_specified:
             do_classifier_free_guidance = guidance_scale > 1.0
             if do_classifier_free_guidance:
                 self.unet.batch_size = self.text_encoder.batch_size * 2
+            else:
+                self.unet.batch_size = self.text_encoder.batch_size
 
     @property
     def batch_size(self):

optimum/rbln/diffusers/modeling_diffusers.py

@@ -21,8 +21,6 @@ import torch
 
 from ..configuration_utils import ContextRblnConfig, RBLNModelConfig
 from ..modeling import RBLNModel
-
-# from ..transformers import RBLNCLIPTextModelConfig
 from ..utils.decorator_utils import remove_compile_time_kwargs
 from ..utils.logging import get_logger
 
@@ -70,6 +68,7 @@ class RBLNDiffusionMixin:
     _submodules = []
     _prefix = {}
     _rbln_config_class = None
+    _hf_class = None
 
     @staticmethod
     def _maybe_apply_and_fuse_lora(
@@ -125,6 +124,14 @@ class RBLNDiffusionMixin:
                 )
         return cls._rbln_config_class
 
+    @classmethod
+    def get_hf_class(cls):
+        if cls._hf_class is None:
+            hf_cls_name = cls.__name__[4:]
+            library = importlib.import_module("diffusers")
+            cls._hf_class = getattr(library, hf_cls_name, None)
+        return cls._hf_class
+
     @classmethod
     def from_pretrained(
         cls,

optimum/rbln/diffusers/models/controlnet.py

@@ -222,7 +222,7 @@ class RBLNControlNetModel(RBLNModel):
                 f"Mismatch between ControlNet's runtime batch size ({sample_batch_size}) and compiled batch size ({compiled_batch_size}). "
                 "This may be caused by the 'guidance_scale' parameter, which doubles the runtime batch size of ControlNet in Stable Diffusion. "
                 "Adjust the batch size of ControlNet during compilation to match the runtime batch size.\n\n"
-                "For details, see: https://docs.rbln.ai/software/optimum/model_api.html#stable-diffusion"
+                "For details, see: https://docs.rbln.ai/software/optimum/model_api/diffusers/pipelines/controlnet.html#important-batch-size-configuration-for-guidance-scale"
             )
 
         added_cond_kwargs = {} if added_cond_kwargs is None else added_cond_kwargs

optimum/rbln/diffusers/models/transformers/transformer_sd3.py

@@ -161,7 +161,7 @@ class RBLNSD3Transformer2DModel(RBLNModel):
                 f"Mismatch between transformer's runtime batch size ({sample_batch_size}) and compiled batch size ({compiled_batch_size}). "
                 "This may be caused by the 'guidance scale' parameter, which doubles the runtime batch size in Stable Diffusion. "
                 "Adjust the batch size of transformer during compilation.\n\n"
-                "For details, see: https://docs.rbln.ai/software/optimum/model_api.html#stable-diffusion"
+                "For details, see: https://docs.rbln.ai/software/optimum/model_api/diffusers/pipelines/stable_diffusion_3.html#important-batch-size-configuration-for-guidance-scale"
             )
 
         return super().forward(

optimum/rbln/diffusers/models/unets/unet_2d_condition.py

@@ -346,7 +346,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
                 f"Mismatch between UNet's runtime batch size ({sample_batch_size}) and compiled batch size ({compiled_batch_size}). "
                 "This may be caused by the 'guidance scale' parameter, which doubles the runtime batch size of UNet in Stable Diffusion. "
                 "Adjust the batch size of UNet during compilation to match the runtime batch size.\n\n"
-                "For details, see: https://docs.rbln.ai/software/optimum/model_api.html#stable-diffusion"
+                "For details, see: https://docs.rbln.ai/software/optimum/model_api/diffusers/pipelines/stable_diffusion.html#important-batch-size-configuration-for-guidance-scale"
             )
 
         added_cond_kwargs = {} if added_cond_kwargs is None else added_cond_kwargs

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet.py

@@ -39,6 +39,7 @@ from diffusers.utils.torch_utils import is_compiled_module, is_torch_version
 
 from ....utils.decorator_utils import remove_compile_time_kwargs
 from ....utils.logging import get_logger
+from ...configurations import RBLNStableDiffusionControlNetPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
@@ -49,6 +50,7 @@ logger = get_logger(__name__)
 
 class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionControlNetPipeline):
     original_class = StableDiffusionControlNetPipeline
+    _rbln_config_class = RBLNStableDiffusionControlNetPipelineConfig
     _submodules = ["text_encoder", "unet", "vae", "controlnet"]
 
     # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet.py

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py

@@ -37,6 +37,7 @@ from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module
 
 from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...configurations import RBLNStableDiffusionControlNetImg2ImgPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
@@ -48,6 +49,7 @@ logger = logging.get_logger(__name__)
 class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionControlNetImg2ImgPipeline):
     original_class = StableDiffusionControlNetImg2ImgPipeline
     _submodules = ["text_encoder", "unet", "vae", "controlnet"]
+    _rbln_config_class = RBLNStableDiffusionControlNetImg2ImgPipelineConfig
 
     # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_img2img.py
     def check_inputs(