optimum-rbln 0.8.1a0__py3-none-any.whl → 0.8.1a2__py3-none-any.whl

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionImg2ImgPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion pipeline for image-to-image generation.
+
+    This pipeline compiles Stable Diffusion models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for transforming input images based on text prompts with controlled strength and guidance.
+    """
+
     original_class = StableDiffusionImg2ImgPipeline
     _rbln_config_class = RBLNStableDiffusionImg2ImgPipelineConfig
     _submodules = ["text_encoder", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionInpaintPipeline(RBLNDiffusionMixin, StableDiffusionInpaintPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion pipeline for image inpainting.
+
+    This pipeline compiles Stable Diffusion models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for filling masked regions of images based on text prompts with seamless integration.
+    """
+
     original_class = StableDiffusionInpaintPipeline
     _rbln_config_class = RBLNStableDiffusionInpaintPipelineConfig
     _submodules = ["text_encoder", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3Pipeline(RBLNDiffusionMixin, StableDiffusion3Pipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion 3 pipeline for advanced text-to-image generation.
+
+    This pipeline compiles Stable Diffusion 3 models to run efficiently on RBLN NPUs, enabling high-performance
+    inference with improved text understanding, enhanced image quality, and superior prompt adherence.
+    """
+
     original_class = StableDiffusion3Pipeline
     _rbln_config_class = RBLNStableDiffusion3PipelineConfig
     _submodules = ["transformer", "text_encoder_3", "text_encoder", "text_encoder_2", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3Img2ImgPipeline(RBLNDiffusionMixin, StableDiffusion3Img2ImgPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion 3 pipeline for advanced image-to-image generation.
+
+    This pipeline compiles Stable Diffusion 3 models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for transforming input images with superior text understanding and enhanced visual quality.
+    """
+
     original_class = StableDiffusion3Img2ImgPipeline
     _rbln_config_class = RBLNStableDiffusion3Img2ImgPipelineConfig
     _submodules = ["transformer", "text_encoder_3", "text_encoder", "text_encoder_2", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3InpaintPipeline(RBLNDiffusionMixin, StableDiffusion3InpaintPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion 3 pipeline for advanced image inpainting.
+
+    This pipeline compiles Stable Diffusion 3 models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for filling masked regions with superior text understanding and seamless content generation.
+    """
+
     original_class = StableDiffusion3InpaintPipeline
     _rbln_config_class = RBLNStableDiffusion3InpaintPipelineConfig
     _submodules = ["transformer", "text_encoder_3", "text_encoder", "text_encoder_2", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLPipeline(RBLNDiffusionMixin, StableDiffusionXLPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion XL pipeline for high-resolution text-to-image generation.
+
+    This pipeline compiles Stable Diffusion XL models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for generating high-quality images with enhanced detail and improved prompt adherence.
+    """
+
     original_class = StableDiffusionXLPipeline
     _rbln_config_class = RBLNStableDiffusionXLPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionXLImg2ImgPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion XL pipeline for high-resolution image-to-image generation.
+
+    This pipeline compiles Stable Diffusion XL models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for transforming input images with enhanced quality and detail preservation.
+    """
+
     original_class = StableDiffusionXLImg2ImgPipeline
     _rbln_config_class = RBLNStableDiffusionXLImg2ImgPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py

@@ -19,6 +19,13 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLInpaintPipeline(RBLNDiffusionMixin, StableDiffusionXLInpaintPipeline):
+    """
+    RBLN-accelerated implementation of Stable Diffusion XL pipeline for high-resolution image inpainting.
+
+    This pipeline compiles Stable Diffusion XL models to run efficiently on RBLN NPUs, enabling high-performance
+    inference for filling masked regions with enhanced quality and seamless blending capabilities.
+    """
+
     original_class = StableDiffusionXLInpaintPipeline
     _rbln_config_class = RBLNStableDiffusionXLInpaintPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]

optimum/rbln/modeling.py

@@ -14,7 +14,7 @@
 
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Dict, List, Optional, Union, get_args, get_origin, get_type_hints
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, get_args, get_origin, get_type_hints
 
 import rebel
 import torch
@@ -35,27 +35,12 @@ logger = get_logger(__name__)
 
 
 class RBLNModel(RBLNBaseModel):
-    """
-    A class that inherits from RBLNBaseModel for models consisting of a single `torch.nn.Module`.
-
-    This class supports all the functionality of RBLNBaseModel, including loading and saving models using
-    the `from_pretrained` and `save_pretrained` methods, compiling PyTorch models for execution on RBLN NPU
-    devices.
-
-    Example:
-        ```python
-        model = RBLNModel.from_pretrained("model_id", export=True, rbln_npu="npu_name")
-        outputs = model(**inputs)
-        ```
-    """
-
     _output_class = None
 
     @classmethod
     def update_kwargs(cls, kwargs):
-        """
-        Update user-given kwargs to get proper pytorch model.
-        """
+        # Update user-given kwargs to get proper pytorch model.
+
         return kwargs
 
     @classmethod
@@ -66,10 +51,9 @@ class RBLNModel(RBLNBaseModel):
         subfolder: str,
         rbln_config: RBLNModelConfig,
     ):
-        """
-        If you are unavoidably running on a CPU rather than an RBLN device,
-        store the torch tensor, weight, etc. in this function.
-        """
+        # If you are unavoidably running on a CPU rather than an RBLN device,
+        # store the torch tensor, weight, etc. in this function.
+        pass
 
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> torch.nn.Module:
@@ -88,11 +72,32 @@ class RBLNModel(RBLNBaseModel):
         cls,
         model: "PreTrainedModel",
         config: Optional[PretrainedConfig] = None,
-        rbln_config: Optional[RBLNModelConfig] = None,
+        rbln_config: Optional[Union[RBLNModelConfig, Dict]] = None,
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         subfolder: str = "",
-        **kwargs,
-    ):
+        **kwargs: Dict[str, Any],
+    ) -> "RBLNModel":
+        """
+        Converts and compiles a pre-trained HuggingFace library model into a RBLN model.
+        This method performs the actual model conversion and compilation process.
+
+        Args:
+            model: The PyTorch model to be compiled. The object must be an instance of the HuggingFace transformers PreTrainedModel class.
+            rbln_config: Configuration for RBLN model compilation and runtime. This can be provided as a dictionary or an instance of the model's configuration class (e.g., `RBLNLlamaForCausalLMConfig` for Llama models).
+                For detailed configuration options, see the specific model's configuration class documentation.
+
+            kwargs: Additional keyword arguments. Arguments with the prefix 'rbln_' are passed to rbln_config, while the remaining arguments are passed to the HuggingFace library.
+
+        The method performs the following steps:
+
+        1. Compiles the PyTorch model into an optimized RBLN graph
+        2. Configures the model for the specified NPU device
+        3. Creates the necessary runtime objects if requested
+        4. Saves the compiled model and configurations
+
+        Returns:
+            A RBLN model instance ready for inference on RBLN NPU devices.
+        """
         preprocessors = kwargs.pop("preprocessors", [])
         rbln_config, kwargs = cls.prepare_rbln_config(rbln_config=rbln_config, **kwargs)
 
@@ -246,12 +251,7 @@ class RBLNModel(RBLNBaseModel):
 
     @classmethod
     def get_hf_output_class(cls):
-        """
-        Dynamically gets the output class from the corresponding HuggingFace model class.
-
-        Returns:
-            type: The appropriate output class from transformers or diffusers
-        """
+        # Dynamically gets the output class from the corresponding HuggingFace model class.
         if cls._output_class:
             return cls._output_class
 
@@ -278,10 +278,8 @@ class RBLNModel(RBLNBaseModel):
         return BaseModelOutput
 
     def _prepare_output(self, output, return_dict):
-        """
-        Prepare model output based on return_dict flag.
-        This method can be overridden by subclasses to provide task-specific output handling.
-        """
+        # Prepare model output based on return_dict flag.
+        # This method can be overridden by subclasses to provide task-specific output handling.
         tuple_output = (output,) if not isinstance(output, (tuple, list)) else tuple(output)
         if not return_dict:
             return tuple_output

optimum/rbln/modeling_base.py

@@ -15,7 +15,7 @@
 import importlib
 import os
 import shutil
-from abc import ABC, abstractmethod
+from abc import ABC
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
@@ -24,7 +24,7 @@ import rebel
 import torch
 from transformers import AutoConfig, AutoModel, GenerationConfig, PretrainedConfig
 
-from .configuration_utils import RBLNAutoConfig, RBLNCompileConfig, RBLNModelConfig
+from .configuration_utils import RBLNAutoConfig, RBLNCompileConfig, RBLNModelConfig, get_rbln_config_class
 from .utils.hub import PushToHubMixin, pull_compiled_model_from_hub, validate_files
 from .utils.logging import get_logger
 from .utils.runtime_utils import UnavailableRuntime
@@ -47,40 +47,6 @@ class RBLNBaseModelConfig(RBLNModelConfig):
 
 
 class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
-    """
-    An abstract base class for compiling, loading, and saving neural network models from the huggingface
-    transformers and diffusers libraries to run on RBLN NPU devices.
-
-    This class supports loading and saving models using the `from_pretrained` and `save_pretrained` methods,
-    similar to the huggingface libraries.
-
-    The `from_pretrained` method loads a model corresponding to the given `model_id` from a local repository
-    or the huggingface hub onto the NPU. If the model is a PyTorch model and `export=True` is passed as a
-    kwarg, it compiles the PyTorch model corresponding to the given `model_id` before loading. If `model_id`
-    is an already rbln-compiled model, it can be directly loaded onto the NPU with `export=False`.
-
-    `rbln_npu` is a kwarg required for compilation, specifying the name of the NPU to be used. If this
-    keyword is not specified, the NPU installed on the host machine is used. If no NPU is installed on the
-    host machine, an error occurs.
-
-    `rbln_device` specifies the device to be used at runtime. If not specified, device 0 is used.
-
-    `rbln_create_runtimes` indicates whether to create runtime objects. If False, the runtime does not load
-    the model onto the NPU. This option is particularly useful when you want to perform compilation only on a
-    host machine without an NPU.
-
-    `RBLNModel`, `RBLNModelFor*`, etc. are all child classes of RBLNBaseModel.
-
-    Models compiled in this way can be saved to a local repository using `save_pretrained` or uploaded to
-    the huggingface hub.
-
-    It also supports generation through `generate` (for transformers models that support generation).
-
-    RBLNBaseModel is a class for models consisting of an arbitrary number of `torch.nn.Module`s, and
-    therefore is an abstract class without explicit implementations of `forward` or `export` functions.
-    To inherit from this class, `forward`, `export`, etc. must be implemented.
-    """
-
     model_type = "rbln_model"
     auto_model_class = AutoModel
     config_class = AutoConfig
@@ -156,7 +122,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         subfolder: str = "",
         local_files_only: bool = False,
     ) -> str:
-        """Load the directory containing the compiled model files."""
+        # Load the directory containing the compiled model files.
         model_path = Path(model_id)
 
         if model_path.is_dir():
@@ -372,19 +338,40 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
     def prepare_rbln_config(
         cls, rbln_config: Optional[Union[Dict[str, Any], RBLNModelConfig]] = None, **kwargs
     ) -> Tuple[RBLNModelConfig, Dict[str, Any]]:
-        """
-        Extract rbln-config from kwargs and convert it to RBLNModelConfig.
-        """
+        # Extract rbln-config from kwargs and convert it to RBLNModelConfig.
+
         config_cls = cls.get_rbln_config_class()
         rbln_config, kwargs = config_cls.initialize_from_kwargs(rbln_config, **kwargs)
         return rbln_config, kwargs
 
     @classmethod
-    def from_pretrained(cls, model_id: Union[str, Path], export: bool = False, **kwargs) -> "RBLNBaseModel":
+    def from_pretrained(
+        cls: Type["RBLNBaseModel"],
+        model_id: Union[str, Path],
+        export: bool = False,
+        rbln_config: Optional[Union[Dict, RBLNModelConfig]] = None,
+        **kwargs: Dict[str, Any],
+    ) -> "RBLNBaseModel":
+        """
+        The `from_pretrained()` function is utilized in its standard form as in the HuggingFace transformers library.
+        User can use this function to load a pre-trained model from the HuggingFace library and convert it to a RBLN model to be run on RBLN NPUs.
+
+        Args:
+            model_id: The model id of the pre-trained model to be loaded. It can be downloaded from the HuggingFace model hub or a local path, or a model id of a compiled model using the RBLN Compiler.
+            export: A boolean flag to indicate whether the model should be compiled.
+            rbln_config: Configuration for RBLN model compilation and runtime. This can be provided as a dictionary or an instance of the model's configuration class (e.g., `RBLNLlamaForCausalLMConfig` for Llama models).
+                For detailed configuration options, see the specific model's configuration class documentation.
+
+            kwargs: Additional keyword arguments. Arguments with the prefix 'rbln_' are passed to rbln_config, while the remaining arguments are passed to the HuggingFace library.
+
+        Returns:
+            A RBLN model instance ready for inference on RBLN NPU devices.
+        """
+
         if isinstance(model_id, Path):
             model_id = model_id.as_posix()
         from_pretrained_method = cls._export if export else cls._from_pretrained
-        return from_pretrained_method(model_id=model_id, **kwargs)
+        return from_pretrained_method(model_id=model_id, **kwargs, rbln_config=rbln_config)
 
     @classmethod
     def compile(cls, model, rbln_compile_config: Optional[RBLNCompileConfig] = None, **kwargs):
@@ -411,15 +398,13 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
 
     @classmethod
     def get_hf_class(cls):
-        """
-        Lazily loads and caches the corresponding HuggingFace model class.
-        Removes 'RBLN' prefix from the class name to get the original class name
-        (e.g., RBLNLlamaForCausalLM -> LlamaForCausalLM) and imports it from
-        the transformers/diffusers module.
+        # Lazily loads and caches the corresponding HuggingFace model class.
+        # Removes 'RBLN' prefix from the class name to get the original class name
+        # (e.g., RBLNLlamaForCausalLM -> LlamaForCausalLM) and imports it from
+        # the transformers/diffusers module.
 
-        Returns:
-            type: The original HuggingFace model class
-        """
+        # Returns:
+        #     type: The original HuggingFace model class
         if cls._hf_class is None:
             hf_cls_name = cls.__name__[4:]
             library = importlib.import_module(cls.hf_library_name)
@@ -428,18 +413,10 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
 
     @classmethod
     def get_rbln_config_class(cls) -> Type[RBLNModelConfig]:
-        """
-        Lazily loads and caches the corresponding RBLN model config class.
-        """
+        # Lazily loads and caches the corresponding RBLN model config class.
         if cls._rbln_config_class is None:
             rbln_config_class_name = cls.__name__ + "Config"
-            library = importlib.import_module("optimum.rbln")
-            cls._rbln_config_class = getattr(library, rbln_config_class_name, None)
-            if cls._rbln_config_class is None:
-                raise ValueError(
-                    f"RBLN config class {rbln_config_class_name} not found. This is an internal error. "
-                    "Please report it to the developers."
-                )
+            cls._rbln_config_class = get_rbln_config_class(rbln_config_class_name)
         return cls._rbln_config_class
 
     def can_generate(self):
@@ -449,17 +426,15 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         return self
 
     def parameters(self):
-        """
-        Provides a dummy parameter generator for compatibility.
+        # A dummy parameter generator for compatibility.
 
-        This method mimics the interface of torch.nn.Module.parameters()
-        specifically for code that uses `next(model.parameters())` to infer
-        the device or dtype. It yields a single dummy tensor on CPU with float32 dtype.
+        # This method mimics the interface of torch.nn.Module.parameters()
+        # specifically for code that uses `next(model.parameters())` to infer
+        # the device or dtype. It yields a single dummy tensor on CPU with float32 dtype.
 
-        Warning:
-            This does NOT yield the actual model parameters used by the RBLN runtime.
-            Code relying on iterating through all model parameters will not work as expected.
-        """
+        # Warning:
+        #     This does NOT yield the actual model parameters used by the RBLN runtime.
+        #     Code relying on iterating through all model parameters will not work as expected.
         yield torch.tensor([1.0], dtype=torch.float32, device=torch.device("cpu"))
 
     def __call__(self, *args, **kwargs):
@@ -547,7 +522,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
 
     @staticmethod
     def _raise_missing_compiled_file_error(missing_files: List[str]):
-        """Raises a KeyError with a message indicating missing compiled model files."""
+        # Raises a KeyError with a message indicating missing compiled model files.
 
         if len(missing_files) == 1:
             message = f"The rbln model folder is missing the required '{missing_files[0]}.rbln' file. "
@@ -563,40 +538,3 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             "and ensure the compilation completes successfully."
         )
         raise KeyError(message)
-
-    @classmethod
-    @abstractmethod
-    def _update_rbln_config(cls, **rbln_config_kwargs) -> RBLNModelConfig:
-        pass
-
-    @classmethod
-    @abstractmethod
-    def _create_runtimes(
-        cls,
-        compiled_models: List[rebel.RBLNCompiledModel],
-        rbln_config: RBLNModelConfig,
-    ) -> List[rebel.Runtime]:
-        # compiled_models -> runtimes
-        pass
-
-    @classmethod
-    @abstractmethod
-    def get_pytorch_model(cls, *args, **kwargs):
-        pass
-
-    @classmethod
-    @abstractmethod
-    def from_model(
-        cls,
-        model: "PreTrainedModel",
-        config: Optional[PretrainedConfig] = None,
-        rbln_config: Optional[RBLNModelConfig] = None,
-        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
-        subfolder: str = "",
-        **kwargs,
-    ):
-        pass
-
-    @abstractmethod
-    def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
-        pass

optimum/rbln/transformers/__init__.py

@@ -18,16 +18,9 @@ from transformers.utils import _LazyModule
 
 
 _import_structure = {
-    "configuration_alias": [
-        "RBLNASTForAudioClassificationConfig",
-        "RBLNDistilBertForQuestionAnsweringConfig",
-        "RBLNResNetForImageClassificationConfig",
-        "RBLNXLMRobertaForSequenceClassificationConfig",
-        "RBLNRobertaForSequenceClassificationConfig",
-        "RBLNRobertaForMaskedLMConfig",
-        "RBLNViTForImageClassificationConfig",
-    ],
     "models": [
+        "RBLNASTForAudioClassification",
+        "RBLNASTForAudioClassificationConfig",
         "RBLNAutoModel",
         "RBLNAutoModelForAudioClassification",
         "RBLNAutoModelForCausalLM",
@@ -51,12 +44,12 @@ _import_structure = {
         "RBLNBertForQuestionAnsweringConfig",
         "RBLNBertModel",
         "RBLNBertModelConfig",
-        "RBLNBlip2VisionModelConfig",
-        "RBLNBlip2VisionModel",
-        "RBLNBlip2QFormerModel",
-        "RBLNBlip2QFormerModelConfig",
         "RBLNBlip2ForConditionalGeneration",
         "RBLNBlip2ForConditionalGenerationConfig",
+        "RBLNBlip2QFormerModel",
+        "RBLNBlip2QFormerModelConfig",
+        "RBLNBlip2VisionModel",
+        "RBLNBlip2VisionModelConfig",
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelConfig",
         "RBLNCLIPTextModelWithProjection",
@@ -67,40 +60,48 @@ _import_structure = {
         "RBLNCLIPVisionModelWithProjectionConfig",
         "RBLNDecoderOnlyModelForCausalLM",
         "RBLNDecoderOnlyModelForCausalLMConfig",
+        "RBLNDistilBertForQuestionAnswering",
+        "RBLNDistilBertForQuestionAnsweringConfig",
         "RBLNDPTForDepthEstimation",
         "RBLNDPTForDepthEstimationConfig",
         "RBLNExaoneForCausalLM",
         "RBLNExaoneForCausalLMConfig",
-        "RBLNGemmaForCausalLM",
-        "RBLNGemmaForCausalLMConfig",
         "RBLNGemma3ForCausalLM",
         "RBLNGemma3ForCausalLMConfig",
         "RBLNGemma3ForConditionalGeneration",
         "RBLNGemma3ForConditionalGenerationConfig",
+        "RBLNGemmaForCausalLM",
+        "RBLNGemmaForCausalLMConfig",
         "RBLNGPT2LMHeadModel",
         "RBLNGPT2LMHeadModelConfig",
-        "RBLNIdefics3VisionTransformer",
         "RBLNIdefics3ForConditionalGeneration",
         "RBLNIdefics3ForConditionalGenerationConfig",
+        "RBLNIdefics3VisionTransformer",
         "RBLNIdefics3VisionTransformerConfig",
         "RBLNLlamaForCausalLM",
         "RBLNLlamaForCausalLMConfig",
-        "RBLNOPTForCausalLM",
-        "RBLNOPTForCausalLMConfig",
         "RBLNLlavaNextForConditionalGeneration",
         "RBLNLlavaNextForConditionalGenerationConfig",
         "RBLNMidmLMHeadModel",
         "RBLNMidmLMHeadModelConfig",
         "RBLNMistralForCausalLM",
         "RBLNMistralForCausalLMConfig",
+        "RBLNOPTForCausalLM",
+        "RBLNOPTForCausalLMConfig",
         "RBLNPhiForCausalLM",
         "RBLNPhiForCausalLMConfig",
-        "RBLNQwen2ForCausalLM",
-        "RBLNQwen2ForCausalLMConfig",
         "RBLNQwen2_5_VisionTransformerPretrainedModel",
         "RBLNQwen2_5_VisionTransformerPretrainedModelConfig",
         "RBLNQwen2_5_VLForConditionalGeneration",
         "RBLNQwen2_5_VLForConditionalGenerationConfig",
+        "RBLNQwen2ForCausalLM",
+        "RBLNQwen2ForCausalLMConfig",
+        "RBLNResNetForImageClassification",
+        "RBLNResNetForImageClassificationConfig",
+        "RBLNRobertaForMaskedLM",
+        "RBLNRobertaForMaskedLMConfig",
+        "RBLNRobertaForSequenceClassification",
+        "RBLNRobertaForSequenceClassificationConfig",
         "RBLNSiglipVisionModel",
         "RBLNSiglipVisionModelConfig",
         "RBLNT5EncoderModel",
@@ -109,44 +110,23 @@ _import_structure = {
         "RBLNT5ForConditionalGenerationConfig",
         "RBLNTimeSeriesTransformerForPrediction",
         "RBLNTimeSeriesTransformerForPredictionConfig",
+        "RBLNViTForImageClassification",
+        "RBLNViTForImageClassificationConfig",
         "RBLNWav2Vec2ForCTC",
         "RBLNWav2Vec2ForCTCConfig",
         "RBLNWhisperForConditionalGeneration",
         "RBLNWhisperForConditionalGenerationConfig",
+        "RBLNXLMRobertaForSequenceClassification",
+        "RBLNXLMRobertaForSequenceClassificationConfig",
         "RBLNXLMRobertaModel",
         "RBLNXLMRobertaModelConfig",
     ],
-    "modeling_alias": [
-        "RBLNASTForAudioClassification",
-        "RBLNDistilBertForQuestionAnswering",
-        "RBLNResNetForImageClassification",
-        "RBLNXLMRobertaForSequenceClassification",
-        "RBLNRobertaForSequenceClassification",
-        "RBLNRobertaForMaskedLM",
-        "RBLNViTForImageClassification",
-    ],
 }
 
 if TYPE_CHECKING:
-    from .configuration_alias import (
-        RBLNASTForAudioClassificationConfig,
-        RBLNDistilBertForQuestionAnsweringConfig,
-        RBLNResNetForImageClassificationConfig,
-        RBLNRobertaForMaskedLMConfig,
-        RBLNRobertaForSequenceClassificationConfig,
-        RBLNViTForImageClassificationConfig,
-        RBLNXLMRobertaForSequenceClassificationConfig,
-    )
-    from .modeling_alias import (
-        RBLNASTForAudioClassification,
-        RBLNDistilBertForQuestionAnswering,
-        RBLNResNetForImageClassification,
-        RBLNRobertaForMaskedLM,
-        RBLNRobertaForSequenceClassification,
-        RBLNViTForImageClassification,
-        RBLNXLMRobertaForSequenceClassification,
-    )
     from .models import (
+        RBLNASTForAudioClassification,
+        RBLNASTForAudioClassificationConfig,
         RBLNAutoModel,
         RBLNAutoModelForAudioClassification,
         RBLNAutoModelForCausalLM,
@@ -186,6 +166,8 @@ if TYPE_CHECKING:
         RBLNCLIPVisionModelWithProjectionConfig,
         RBLNDecoderOnlyModelForCausalLM,
         RBLNDecoderOnlyModelForCausalLMConfig,
+        RBLNDistilBertForQuestionAnswering,
+        RBLNDistilBertForQuestionAnsweringConfig,
         RBLNDPTForDepthEstimation,
         RBLNDPTForDepthEstimationConfig,
         RBLNExaoneForCausalLM,
@@ -220,6 +202,12 @@ if TYPE_CHECKING:
         RBLNQwen2_5_VLForConditionalGenerationConfig,
         RBLNQwen2ForCausalLM,
         RBLNQwen2ForCausalLMConfig,
+        RBLNResNetForImageClassification,
+        RBLNResNetForImageClassificationConfig,
+        RBLNRobertaForMaskedLM,
+        RBLNRobertaForMaskedLMConfig,
+        RBLNRobertaForSequenceClassification,
+        RBLNRobertaForSequenceClassificationConfig,
         RBLNSiglipVisionModel,
         RBLNSiglipVisionModelConfig,
         RBLNT5EncoderModel,
@@ -228,10 +216,14 @@ if TYPE_CHECKING:
         RBLNT5ForConditionalGenerationConfig,
         RBLNTimeSeriesTransformerForPrediction,
         RBLNTimeSeriesTransformerForPredictionConfig,
+        RBLNViTForImageClassification,
+        RBLNViTForImageClassificationConfig,
         RBLNWav2Vec2ForCTC,
         RBLNWav2Vec2ForCTCConfig,
         RBLNWhisperForConditionalGeneration,
         RBLNWhisperForConditionalGenerationConfig,
+        RBLNXLMRobertaForSequenceClassification,
+        RBLNXLMRobertaForSequenceClassificationConfig,
         RBLNXLMRobertaModel,
         RBLNXLMRobertaModelConfig,
     )