optimum-rbln 0.7.4a4__py3-none-any.whl → 0.7.4a6__py3-none-any.whl

optimum/rbln/diffusers/__init__.py

@@ -23,6 +23,34 @@ ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES["optimum.rbln"])
 
 
 _import_structure = {
+    "configurations": [
+        "RBLNAutoencoderKLConfig",
+        "RBLNControlNetModelConfig",
+        "RBLNKandinskyV22CombinedPipelineConfig",
+        "RBLNKandinskyV22Img2ImgCombinedPipelineConfig",
+        "RBLNKandinskyV22Img2ImgPipelineConfig",
+        "RBLNKandinskyV22InpaintCombinedPipelineConfig",
+        "RBLNKandinskyV22InpaintPipelineConfig",
+        "RBLNKandinskyV22PipelineConfig",
+        "RBLNKandinskyV22PriorPipelineConfig",
+        "RBLNPriorTransformerConfig",
+        "RBLNStableDiffusionControlNetPipelineConfig",
+        "RBLNStableDiffusionControlNetImg2ImgPipelineConfig",
+        "RBLNStableDiffusionImg2ImgPipelineConfig",
+        "RBLNStableDiffusionInpaintPipelineConfig",
+        "RBLNStableDiffusionPipelineConfig",
+        "RBLNStableDiffusionXLControlNetPipelineConfig",
+        "RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig",
+        "RBLNStableDiffusionXLImg2ImgPipelineConfig",
+        "RBLNStableDiffusionXLInpaintPipelineConfig",
+        "RBLNStableDiffusionXLPipelineConfig",
+        "RBLNStableDiffusion3PipelineConfig",
+        "RBLNStableDiffusion3Img2ImgPipelineConfig",
+        "RBLNStableDiffusion3InpaintPipelineConfig",
+        "RBLNSD3Transformer2DModelConfig",
+        "RBLNUNet2DConditionModelConfig",
+        "RBLNVQModelConfig",
+    ],
     "pipelines": [
         "RBLNKandinskyV22CombinedPipeline",
         "RBLNKandinskyV22Img2ImgCombinedPipeline",
@@ -60,6 +88,34 @@ _import_structure = {
 }
 
 if TYPE_CHECKING:
+    from .configurations import (
+        RBLNAutoencoderKLConfig,
+        RBLNControlNetModelConfig,
+        RBLNKandinskyV22CombinedPipelineConfig,
+        RBLNKandinskyV22Img2ImgCombinedPipelineConfig,
+        RBLNKandinskyV22Img2ImgPipelineConfig,
+        RBLNKandinskyV22InpaintCombinedPipelineConfig,
+        RBLNKandinskyV22InpaintPipelineConfig,
+        RBLNKandinskyV22PipelineConfig,
+        RBLNKandinskyV22PriorPipelineConfig,
+        RBLNPriorTransformerConfig,
+        RBLNSD3Transformer2DModelConfig,
+        RBLNStableDiffusion3Img2ImgPipelineConfig,
+        RBLNStableDiffusion3InpaintPipelineConfig,
+        RBLNStableDiffusion3PipelineConfig,
+        RBLNStableDiffusionControlNetImg2ImgPipelineConfig,
+        RBLNStableDiffusionControlNetPipelineConfig,
+        RBLNStableDiffusionImg2ImgPipelineConfig,
+        RBLNStableDiffusionInpaintPipelineConfig,
+        RBLNStableDiffusionPipelineConfig,
+        RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig,
+        RBLNStableDiffusionXLControlNetPipelineConfig,
+        RBLNStableDiffusionXLImg2ImgPipelineConfig,
+        RBLNStableDiffusionXLInpaintPipelineConfig,
+        RBLNStableDiffusionXLPipelineConfig,
+        RBLNUNet2DConditionModelConfig,
+        RBLNVQModelConfig,
+    )
     from .modeling_diffusers import RBLNDiffusionMixin
     from .models import (
         RBLNAutoencoderKL,

optimum/rbln/diffusers/configurations/__init__.py

@@ -0,0 +1,30 @@
+from .models import (
+    RBLNAutoencoderKLConfig,
+    RBLNControlNetModelConfig,
+    RBLNPriorTransformerConfig,
+    RBLNSD3Transformer2DModelConfig,
+    RBLNUNet2DConditionModelConfig,
+    RBLNVQModelConfig,
+)
+from .pipelines import (
+    RBLNKandinskyV22CombinedPipelineConfig,
+    RBLNKandinskyV22Img2ImgCombinedPipelineConfig,
+    RBLNKandinskyV22Img2ImgPipelineConfig,
+    RBLNKandinskyV22InpaintCombinedPipelineConfig,
+    RBLNKandinskyV22InpaintPipelineConfig,
+    RBLNKandinskyV22PipelineConfig,
+    RBLNKandinskyV22PriorPipelineConfig,
+    RBLNStableDiffusion3Img2ImgPipelineConfig,
+    RBLNStableDiffusion3InpaintPipelineConfig,
+    RBLNStableDiffusion3PipelineConfig,
+    RBLNStableDiffusionControlNetImg2ImgPipelineConfig,
+    RBLNStableDiffusionControlNetPipelineConfig,
+    RBLNStableDiffusionImg2ImgPipelineConfig,
+    RBLNStableDiffusionInpaintPipelineConfig,
+    RBLNStableDiffusionPipelineConfig,
+    RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig,
+    RBLNStableDiffusionXLControlNetPipelineConfig,
+    RBLNStableDiffusionXLImg2ImgPipelineConfig,
+    RBLNStableDiffusionXLInpaintPipelineConfig,
+    RBLNStableDiffusionXLPipelineConfig,
+)

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

@@ -0,0 +1,6 @@
+from .configuration_autoencoder_kl import RBLNAutoencoderKLConfig
+from .configuration_controlnet import RBLNControlNetModelConfig
+from .configuration_prior_transformer import RBLNPriorTransformerConfig
+from .configuration_transformer_sd3 import RBLNSD3Transformer2DModelConfig
+from .configuration_unet_2d_condition import RBLNUNet2DConditionModelConfig
+from .configuration_vq_model import RBLNVQModelConfig

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

@@ -0,0 +1,66 @@
+# 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 typing import Optional, Tuple
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNAutoencoderKLConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        sample_size: Optional[Tuple[int, int]] = None,
+        uses_encoder: Optional[bool] = None,
+        vae_scale_factor: Optional[float] = None,  # TODO: rename to scaling_factor
+        in_channels: Optional[int] = None,
+        latent_channels: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            sample_size (Optional[Tuple[int, int]]): The spatial dimensions (height, width) of the input/output images.
+                If an integer is provided, it's used for both height and width.
+            uses_encoder (Optional[bool]): Whether to include the encoder part of the VAE in the model.
+                When False, only the decoder is used (for latent-to-image conversion).
+            vae_scale_factor (Optional[float]): The scaling factor between pixel space and latent space.
+                Determines how much smaller the latent representations are compared to the original images.
+            in_channels (Optional[int]): Number of input channels for the model.
+            latent_channels (Optional[int]): Number of channels in the latent space.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        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.uses_encoder = uses_encoder
+        self.vae_scale_factor = vae_scale_factor
+        self.in_channels = in_channels
+        self.latent_channels = latent_channels
+        self.sample_size = sample_size
+        if isinstance(sample_size, int):
+            self.sample_size = (sample_size, sample_size)
+
+    @property
+    def image_size(self):
+        return self.sample_size
+
+    @property
+    def latent_sample_size(self):
+        return (self.image_size[0] // self.vae_scale_factor, self.image_size[1] // self.vae_scale_factor)

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

@@ -0,0 +1,54 @@
+# 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 typing import Optional, Tuple
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNControlNetModelConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        max_seq_len: Optional[int] = None,
+        unet_sample_size: Optional[Tuple[int, int]] = None,
+        vae_sample_size: Optional[Tuple[int, int]] = None,
+        text_model_hidden_size: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            max_seq_len (Optional[int]): Maximum sequence length for text inputs when used
+                with cross-attention.
+            unet_sample_size (Optional[Tuple[int, int]]): The spatial dimensions (height, width)
+                of the UNet output samples.
+            vae_sample_size (Optional[Tuple[int, int]]): The spatial dimensions (height, width)
+                of the VAE input/output images.
+            text_model_hidden_size (Optional[int]): Hidden size of the text encoder model used
+                for conditioning.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        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.max_seq_len = max_seq_len
+        self.unet_sample_size = unet_sample_size
+        self.vae_sample_size = vae_sample_size
+        self.text_model_hidden_size = text_model_hidden_size

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

@@ -0,0 +1,44 @@
+# 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 typing import Optional
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNPriorTransformerConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        embedding_dim: Optional[int] = None,
+        num_embeddings: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            embedding_dim (Optional[int]): Dimension of the embedding vectors in the model.
+            num_embeddings (Optional[int]): Number of discrete embeddings in the codebook.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        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

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

@@ -0,0 +1,48 @@
+# 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 typing import Optional, Tuple, Union
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNSD3Transformer2DModelConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        sample_size: Optional[Union[int, Tuple[int, int]]] = None,
+        prompt_embed_length: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            sample_size (Optional[Union[int, Tuple[int, int]]]): The spatial dimensions (height, width)
+                of the generated samples. If an integer is provided, it's used for both height and width.
+            prompt_embed_length (Optional[int]): The length of the embedded prompt vectors that
+                will be used to condition the transformer model.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        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.prompt_embed_length = prompt_embed_length
+        self.sample_size = sample_size
+        if isinstance(self.sample_size, int):
+            self.sample_size = (self.sample_size, self.sample_size)

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

@@ -0,0 +1,66 @@
+# 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 typing import Optional, Tuple
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNUNet2DConditionModelConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        sample_size: Optional[Tuple[int, int]] = None,
+        in_channels: Optional[int] = None,
+        cross_attention_dim: Optional[int] = None,
+        use_additional_residuals: Optional[bool] = None,
+        max_seq_len: Optional[int] = None,
+        in_features: Optional[int] = None,
+        text_model_hidden_size: Optional[int] = None,
+        image_model_hidden_size: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            sample_size (Optional[Tuple[int, int]]): The spatial dimensions (height, width) of the generated samples.
+                If an integer is provided, it's used for both height and width.
+            in_channels (Optional[int]): Number of input channels for the UNet.
+            cross_attention_dim (Optional[int]): Dimension of the cross-attention features.
+            use_additional_residuals (Optional[bool]): Whether to use additional residual connections in the model.
+            max_seq_len (Optional[int]): Maximum sequence length for text inputs when used with cross-attention.
+            in_features (Optional[int]): Number of input features for the model.
+            text_model_hidden_size (Optional[int]): Hidden size of the text encoder model.
+            image_model_hidden_size (Optional[int]): Hidden size of the image encoder model.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        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.in_channels = in_channels
+        self.cross_attention_dim = cross_attention_dim
+        self.use_additional_residuals = use_additional_residuals
+        self.max_seq_len = max_seq_len
+        self.in_features = in_features
+        self.text_model_hidden_size = text_model_hidden_size
+        self.image_model_hidden_size = image_model_hidden_size
+
+        self.sample_size = sample_size
+        if isinstance(sample_size, int):
+            self.sample_size = (sample_size, sample_size)

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

@@ -0,0 +1,67 @@
+# 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 typing import Optional, Tuple
+
+from ....configuration_utils import RBLNModelConfig
+
+
+class RBLNVQModelConfig(RBLNModelConfig):
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        sample_size: Optional[Tuple[int, int]] = None,
+        uses_encoder: Optional[bool] = None,
+        vqmodel_scale_factor: Optional[float] = None,  # TODO: rename to scaling_factor
+        in_channels: Optional[int] = None,
+        latent_channels: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
+            sample_size (Optional[Tuple[int, int]]): The spatial dimensions (height, width) of the input/output images.
+                If an integer is provided, it's used for both height and width.
+            uses_encoder (Optional[bool]): Whether to include the encoder part of the VAE in the model.
+                When False, only the decoder is used (for latent-to-image conversion).
+            vqmodel_scale_factor (Optional[float]): The scaling factor between pixel space and latent space.
+                Determines the downsampling ratio between original images and latent representations.
+            in_channels (Optional[int]): Number of input channels for the model.
+            latent_channels (Optional[int]): Number of channels in the latent space.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If batch_size is not a positive integer.
+        """
+        super().__init__(**kwargs)
+        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.uses_encoder = uses_encoder
+        self.sample_size = sample_size
+        if isinstance(self.sample_size, int):
+            self.sample_size = (self.sample_size, self.sample_size)
+
+        self.vqmodel_scale_factor = vqmodel_scale_factor
+        self.in_channels = in_channels
+        self.latent_channels = latent_channels
+
+    @property
+    def image_size(self):
+        return self.sample_size
+
+    @property
+    def latent_sample_size(self):
+        return (self.image_size[0] // self.vqmodel_scale_factor, self.image_size[1] // self.vqmodel_scale_factor)

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

@@ -0,0 +1,30 @@
+from .configuration_controlnet import (
+    RBLNStableDiffusionControlNetImg2ImgPipelineConfig,
+    RBLNStableDiffusionControlNetPipelineConfig,
+    RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig,
+    RBLNStableDiffusionXLControlNetPipelineConfig,
+)
+from .configuration_kandinsky2_2 import (
+    RBLNKandinskyV22CombinedPipelineConfig,
+    RBLNKandinskyV22Img2ImgCombinedPipelineConfig,
+    RBLNKandinskyV22Img2ImgPipelineConfig,
+    RBLNKandinskyV22InpaintCombinedPipelineConfig,
+    RBLNKandinskyV22InpaintPipelineConfig,
+    RBLNKandinskyV22PipelineConfig,
+    RBLNKandinskyV22PriorPipelineConfig,
+)
+from .configuration_stable_diffusion import (
+    RBLNStableDiffusionImg2ImgPipelineConfig,
+    RBLNStableDiffusionInpaintPipelineConfig,
+    RBLNStableDiffusionPipelineConfig,
+)
+from .configuration_stable_diffusion_3 import (
+    RBLNStableDiffusion3Img2ImgPipelineConfig,
+    RBLNStableDiffusion3InpaintPipelineConfig,
+    RBLNStableDiffusion3PipelineConfig,
+)
+from .configuration_stable_diffusion_xl import (
+    RBLNStableDiffusionXLImg2ImgPipelineConfig,
+    RBLNStableDiffusionXLInpaintPipelineConfig,
+    RBLNStableDiffusionXLPipelineConfig,
+)