optimum-rbln 0.1.4__py3-none-any.whl → 0.1.8__py3-none-any.whl

optimum/rbln/__init__.py

@@ -25,6 +25,8 @@ from typing import TYPE_CHECKING
 
 from transformers.utils import _LazyModule
 
+from .utils import check_version_compats
+
 
 _import_structure = {
     "modeling_alias": [
@@ -33,6 +35,9 @@ _import_structure = {
         "RBLNResNetForImageClassification",
         "RBLNT5ForConditionalGeneration",
         "RBLNBartForConditionalGeneration",
+        "RBLNXLMRobertaForSequenceClassification",
+        "RBLNRobertaForSequenceClassification",
+        "RBLNRobertaForMaskedLM",
     ],
     "modeling_base": [
         "RBLNBaseModel",
@@ -40,6 +45,8 @@ _import_structure = {
         "RBLNModelForQuestionAnswering",
         "RBLNModelForAudioClassification",
         "RBLNModelForImageClassification",
+        "RBLNModelForSequenceClassification",
+        "RBLNModelForMaskedLM",
     ],
     "modeling_seq2seq": [
         "RBLNModelForSeq2SeqLM",
@@ -48,11 +55,14 @@ _import_structure = {
         "BatchTextIteratorStreamer",
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelWithProjection",
+        "RBLNDPTForDepthEstimation",
+        "RBLNGemmaForCausalLM",
         "RBLNGPT2LMHeadModel",
         "RBLNWav2Vec2ForCTC",
         "RBLNLlamaForCausalLM",
         "RBLNMidmLMHeadModel",
         "RBLNWhisperForConditionalGeneration",
+        "RBLNXLMRobertaModel",
     ],
     "diffusers": [
         "RBLNStableDiffusionPipeline",
@@ -91,14 +101,19 @@ if TYPE_CHECKING:
         RBLNBartForConditionalGeneration,
         RBLNBertForQuestionAnswering,
         RBLNResNetForImageClassification,
+        RBLNRobertaForMaskedLM,
+        RBLNRobertaForSequenceClassification,
         RBLNT5ForConditionalGeneration,
+        RBLNXLMRobertaForSequenceClassification,
     )
     from .modeling_base import (
         RBLNBaseModel,
         RBLNModel,
         RBLNModelForAudioClassification,
         RBLNModelForImageClassification,
+        RBLNModelForMaskedLM,
         RBLNModelForQuestionAnswering,
+        RBLNModelForSequenceClassification,
     )
     from .modeling_config import RBLNConfig, RBLNRuntimeConfig
     from .modeling_seq2seq import RBLNModelForSeq2SeqLM
@@ -106,11 +121,14 @@ if TYPE_CHECKING:
         BatchTextIteratorStreamer,
         RBLNCLIPTextModel,
         RBLNCLIPTextModelWithProjection,
+        RBLNDPTForDepthEstimation,
+        RBLNGemmaForCausalLM,
         RBLNGPT2LMHeadModel,
         RBLNLlamaForCausalLM,
         RBLNMidmLMHeadModel,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,
+        RBLNXLMRobertaModel,
     )
 else:
     import sys
@@ -121,3 +139,6 @@ else:
         _import_structure,
         module_spec=__spec__,
     )
+
+
+check_version_compats()

optimum/rbln/__version__.py

@@ -1 +1 @@
-__version__ = '0.1.4'
+__version__ = '0.1.8'

optimum/rbln/diffusers/__init__.py

@@ -47,7 +47,6 @@ _import_structure = {
 }
 
 if TYPE_CHECKING:
-
     from .models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
     from .pipelines import (
         RBLNMultiControlNetModel,

optimum/rbln/diffusers/models/autoencoder_kl.py

@@ -23,7 +23,6 @@
 
 import logging
 from pathlib import Path
-from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 import rebel
@@ -37,7 +36,6 @@ from transformers import AutoConfig, AutoModel, PretrainedConfig
 from ...modeling_base import RBLNModel
 from ...modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
 from ...utils.runtime_utils import RBLNPytorchRuntime
-from ...utils.save_utils import maybe_save_preprocessors
 
 
 logger = logging.getLogger(__name__)
@@ -70,73 +68,13 @@ class RBLNAutoencoderKL(RBLNModel):
         self.rbln_use_encode = self.rbln_config.meta["rbln_use_encode"]
 
         if self.rbln_use_encode:
-            self.encoder = RBLNRuntimeVAEEncoder(runtime=self.runtimes[0], main_input_name="x")
-            self.decoder = RBLNRuntimeVAEDecoder(runtime=self.runtimes[1], main_input_name="z")
+            self.encoder = RBLNRuntimeVAEEncoder(runtime=self.model[0], main_input_name="x")
+            self.decoder = RBLNRuntimeVAEDecoder(runtime=self.model[1], main_input_name="z")
         else:
-            self.decoder = RBLNRuntimeVAEDecoder(runtime=self.runtimes[0], main_input_name="z")
+            self.decoder = RBLNRuntimeVAEDecoder(runtime=self.model[0], main_input_name="z")
 
     @classmethod
-    @torch.no_grad()
-    def _export(
-        cls,
-        model_id: str,
-        config: "PretrainedConfig",
-        use_auth_token: Optional[Union[bool, str]] = None,
-        revision: Optional[str] = None,
-        force_download: bool = False,
-        cache_dir: Optional[str] = None,
-        subfolder: str = "",
-        local_files_only: bool = False,
-        trust_remote_code: bool = False,
-        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
-        **kwargs,
-    ) -> "RBLNAutoencoderKL":
-        task = kwargs.pop("task", None)
-        if task is None:
-            task = TasksManager.infer_task_from_model(cls.auto_model_class)
-
-        if model_save_dir is None:
-            save_dir = TemporaryDirectory()
-            save_dir_path = Path(save_dir.name)
-        else:
-            save_dir = model_save_dir
-            if isinstance(save_dir, TemporaryDirectory):
-                save_dir_path = Path(model_save_dir.name)
-            else:
-                save_dir_path = Path(model_save_dir)
-                save_dir_path.mkdir(exist_ok=True)
-
-        rbln_config_kwargs, rbln_constructor_kwargs = cls.pop_rbln_kwargs_from_kwargs(kwargs)
-
-        model: AutoencoderKL = TasksManager.get_model_from_task(
-            task=None,
-            model_name_or_path=model_id,
-            subfolder=subfolder,
-            revision=revision,
-            framework="pt",
-            cache_dir=cache_dir,
-            use_auth_token=use_auth_token,
-            local_files_only=local_files_only,
-            force_download=force_download,
-            trust_remote_code=trust_remote_code,
-            **kwargs,
-        )
-
-        if config is None:
-            config = model.config
-
-        if not isinstance(config, PretrainedConfig):  # diffusers config
-            config = PretrainedConfig(**config)
-
-        config.save_pretrained(save_dir_path / subfolder)
-        preprocessors = maybe_save_preprocessors(model_id, save_dir_path, src_subfolder=subfolder)
-
-        # Get compilation arguments
-        if rbln_config_kwargs.get("rbln_config", None) is None:
-            rbln_config = cls.get_rbln_config(
-                preprocessors=preprocessors, model_config=model.config, **rbln_config_kwargs
-            )
-
+    def get_compiled_model(cls, model, rbln_config: RBLNConfig):
         def compile_img2img():
             encoder_model = _VAEEncoder(model)
             decoder_model = _VAEDecoder(model)
@@ -146,12 +84,7 @@ class RBLNAutoencoderKL(RBLNModel):
             enc_compiled_model = cls.compile(encoder_model, rbln_runtime_config=rbln_config["encoder"][0])
             dec_compiled_model = cls.compile(decoder_model, rbln_runtime_config=rbln_config["decoder"][0])
 
-            enc_compiled_model.save(
-                save_dir_path / subfolder / f"{rbln_config['encoder'][0].compiled_model_name}.rbln"
-            )
-            dec_compiled_model.save(
-                save_dir_path / subfolder / f"{rbln_config['decoder'][0].compiled_model_name}.rbln"
-            )
+            return enc_compiled_model, dec_compiled_model
 
         def compile_text2img():
             decoder_model = _VAEDecoder(model)
@@ -159,30 +92,12 @@ class RBLNAutoencoderKL(RBLNModel):
 
             dec_compiled_model = cls.compile(decoder_model, rbln_runtime_config=rbln_config["compiled_model"][0])
 
-            dec_compiled_model.save(
-                save_dir_path / subfolder / f"{rbln_config['compiled_model'][0].compiled_model_name}.rbln"
-            )
+            return dec_compiled_model
 
-        if rbln_config_kwargs.get("rbln_use_encode"):
-            compile_img2img()
+        if rbln_config.meta.get("rbln_use_encode", False):
+            return compile_img2img()
         else:
-            compile_text2img()
-
-        rbln_config.save(save_dir_path / subfolder)
-
-        return cls._from_pretrained(
-            model_id=save_dir_path,
-            config=config,
-            model_save_dir=save_dir,
-            use_auth_token=use_auth_token,
-            revision=revision,
-            force_download=force_download,
-            cache_dir=cache_dir,
-            subfolder=subfolder,
-            local_files_only=local_files_only,
-            **rbln_constructor_kwargs,
-            **kwargs,
-        )
+            return compile_text2img()
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
@@ -282,15 +197,18 @@ class RBLNAutoencoderKL(RBLNModel):
         rbln_config = RBLNConfig.from_rbln_runtime_configs([vae_config], _rbln_meta=meta)
         return rbln_config
 
-    def _create_runtimes(self, rbln_device_map: Dict[str, int]) -> List[rebel.Runtime]:
-        if len(self.compiled_models) == 1:
+    @classmethod
+    def _create_runtimes(
+        cls, compiled_models: List[rebel.RBLNCompiledModel], rbln_device_map: Dict[str, int]
+    ) -> List[rebel.Runtime]:
+        if len(compiled_models) == 1:
             device_val = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
-            return [self.compiled_models[0].create_runtime(tensor_type="pt", device=device_val)]
+            return [compiled_models[0].create_runtime(tensor_type="pt", device=device_val)]
 
         device_vals = [rbln_device_map["encoder"], rbln_device_map["decoder"]]
         return [
             compiled_model.create_runtime(tensor_type="pt", device=device_val)
-            for compiled_model, device_val in zip(self.compiled_models, device_vals)
+            for compiled_model, device_val in zip(compiled_models, device_vals)
         ]
 
     def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:

optimum/rbln/diffusers/models/controlnet.py

@@ -120,6 +120,9 @@ class RBLNControlNetModel(RBLNModel):
             model_name_or_path: Union[str, Path],
             **kwargs,
         ):
+            if "subfolder" in kwargs:
+                del kwargs["subfolder"]
+
             return ControlNetModel.from_pretrained(pretrained_model_name_or_path=model_name_or_path, **kwargs)
 
         tasktmp = TasksManager.get_model_from_task

optimum/rbln/diffusers/models/unet_2d_condition.py

@@ -244,6 +244,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
                         for i in range(3)
                     ]
                 )
+            if len(model_config.block_out_channels) > 1:
                 input_info.append(
                     (
                         "down_block_additional_residuals_3",
@@ -251,7 +252,6 @@ class RBLNUNet2DConditionModel(RBLNModel):
                         "float32",
                     )
                 )
-            if len(model_config.block_out_channels) > 1:
                 input_info.extend(
                     [
                         (
@@ -262,6 +262,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
                         for i in range(4, 6)
                     ]
                 )
+            if len(model_config.block_out_channels) > 2:
                 input_info.append(
                     (
                         f"down_block_additional_residuals_{6}",
@@ -269,7 +270,6 @@ class RBLNUNet2DConditionModel(RBLNModel):
                         "float32",
                     )
                 )
-            if len(model_config.block_out_channels) > 2:
                 input_info.extend(
                     [
                         (
@@ -314,7 +314,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
             if rbln_text_model_hidden_size is None:
                 rbln_text_model_hidden_size = 768
             if rbln_in_features is None:
-                rbln_in_features = 2816
+                rbln_in_features = model_config.projection_class_embeddings_input_dim
             meta["in_features"] = rbln_in_features
             rbln_runtime_config.input_info.append(
                 ("text_embeds", [rbln_batch_size, rbln_text_model_hidden_size], "float32")

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

@@ -24,17 +24,15 @@
 import logging
 import os
 from pathlib import Path
-from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
-import rebel
 import torch
 from diffusers import ControlNetModel
 from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
 from optimum.exporters import TasksManager
-from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
+from transformers import AutoConfig, AutoModel
 
-from ....modeling_base import RBLNBaseModel
+from ....modeling_base import RBLNModel
 from ....modeling_config import RBLNConfig
 from ...models.controlnet import RBLNControlNetModel
 
@@ -42,38 +40,16 @@ from ...models.controlnet import RBLNControlNetModel
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    from transformers import (
-        PretrainedConfig,
-        PreTrainedModel,
-    )
+    pass
 
 
-class RBLNMultiControlNetModel(RBLNBaseModel):
-    model_type = "rbln_model"
-    auto_model_class = AutoModel
-
+class RBLNMultiControlNetModel(RBLNModel):
     def __init__(
         self,
-        models: List[Union[PreTrainedModel, rebel.RBLNCompiledModel]],
-        config: PretrainedConfig = None,
-        preprocessors: Optional[List] = None,
-        rbln_config: Optional[RBLNConfig] = None,
+        models: List[RBLNControlNetModel],
         **kwargs,
     ):
-        super().__init__(
-            models,
-            config,
-            preprocessors,
-            rbln_config,
-            **kwargs,
-        )
-
-        if not isinstance(config, PretrainedConfig):
-            config = PretrainedConfig(**config)
-
-        for i in range(len(models)):
-            self.runtimes[i].config = config
-        self.nets = self.runtimes
+        self.nets = models
         self.dtype = torch.float32
 
     @classmethod
@@ -83,7 +59,7 @@ class RBLNMultiControlNetModel(RBLNBaseModel):
             model_name_or_path: Union[str, Path],
             **kwargs,
         ):
-            return MultiControlNetModel.from_pretrained(pretrained_model_path=model_name_or_path, **kwargs)
+            return MultiControlNetModel.from_pretrained(pretrained_model_name_or_path=model_name_or_path, **kwargs)
 
         tasktmp = TasksManager.get_model_from_task
         configtmp = AutoConfig.from_pretrained
@@ -101,131 +77,31 @@ class RBLNMultiControlNetModel(RBLNBaseModel):
     def _from_pretrained(
         cls,
         model_id: Union[str, Path],
-        config: "PretrainedConfig",
-        use_auth_token: Optional[Union[bool, str]] = None,
-        revision: Optional[str] = None,
-        force_download: bool = False,
-        cache_dir: Optional[str] = None,
-        file_name: Optional[str] = None,
-        subfolder: str = "",
-        local_files_only: bool = False,
         **kwargs,
-    ) -> RBLNBaseModel:
-
-        if isinstance(model_id, str):
-            model_path = Path(model_id)
-        else:
-            model_path = model_id / "controlnet"
+    ) -> RBLNModel:
 
-        rbln_files = []
-        rbln_config_filenames = []
         idx = 0
-        model_load_path = model_path
+        controlnets = []
+        model_path_to_load = model_id
 
-        while model_load_path.is_dir():
-            rbln_files.append(list(model_load_path.glob("**/*.rbln"))[0])
-            rbln_config_filenames.append(model_load_path)
+        while os.path.isdir(model_path_to_load):
+            controlnet = RBLNControlNetModel.from_pretrained(model_path_to_load, export=False, **kwargs)
+            controlnets.append(controlnet)
+            rbln_config = RBLNConfig.load(model_path_to_load)
             idx += 1
-            model_load_path = Path(str(model_path) + f"_{idx}")
-
-        if len(rbln_files) == 0:
-            raise FileNotFoundError(f"Could not find any rbln model file in {model_path}")
-
-        if len(rbln_config_filenames) == 0:
-            raise FileNotFoundError(f"Could not find `rbln_config.json` file in {model_path}")
-
-        models = []
-        for rconf, rfiles in zip(rbln_config_filenames, rbln_files):
-            rbln_config = RBLNConfig.load(str(rconf))
-            models.append(rebel.RBLNCompiledModel(rfiles))
-
-        preprocessors = []
+            model_path_to_load = model_id + f"_{idx}"
 
         return cls(
-            models,
-            config,
-            preprocessors,
+            controlnets,
             rbln_config=rbln_config,
             **kwargs,
         )
 
-    def _save_pretrained(self, save_directory: Union[str, Path]):
-        idx = 0
-        real_save_dir_path = save_directory
-        for compiled_model in self.compiled_models:
-            dst_path = Path(real_save_dir_path) / "compiled_model.rbln"
-            if not os.path.exists(real_save_dir_path):
-                os.makedirs(real_save_dir_path)
-            compiled_model.save(dst_path)
-            self.rbln_config.save(real_save_dir_path)
-            idx += 1
-            real_save_dir_path = save_directory + f"_{idx}"
-
-    @classmethod
-    @torch.no_grad()
-    def _export(
-        cls,
-        model_id: str,
-        config: "PretrainedConfig",
-        use_auth_token: Optional[Union[bool, str]] = None,
-        revision: Optional[str] = None,
-        force_download: bool = False,
-        cache_dir: Optional[str] = None,
-        subfolder: str = "",
-        local_files_only: bool = False,
-        trust_remote_code: bool = False,
-        **kwargs,
-    ) -> "RBLNMultiControlNetModel":
-
-        task = kwargs.pop("task", None)
-        if task is None:
-            task = TasksManager.infer_task_from_model(cls.auto_model_class)
-
-        save_dir = TemporaryDirectory()
-        save_dir_path = Path(save_dir.name)
-
-        rbln_config_kwargs, rbln_constructor_kwargs = cls.pop_rbln_kwargs_from_kwargs(kwargs)
-        img_width = rbln_config_kwargs.pop("rbln_img_width", None)
-        img_height = rbln_config_kwargs.pop("rbln_img_height", None)
-        vae_scale_factor = rbln_config_kwargs.pop("rbln_vae_scale_factor", None)
-        batch_size = rbln_config_kwargs.pop("rbln_batch_size", None)
-
-        model: MultiControlNetModel = TasksManager.get_model_from_task(
-            task=task,
-            model_name_or_path=model_id,
-        )
-
-        model_path_to_load = model_id
-        real_save_dir_path = save_dir_path / "controlnet"
-
-        for idx in range(len(model.nets)):
+    def save_pretrained(self, save_directory: Union[str, Path], **kwargs):
+        for idx, model in enumerate(self.nets):
             suffix = "" if idx == 0 else f"_{idx}"
-            controlnet = RBLNControlNetModel.from_pretrained(
-                model_path_to_load + suffix,
-                export=True,
-                rbln_batch_size=batch_size,
-                rbln_img_width=img_width,
-                rbln_img_height=img_height,
-                rbln_vae_scale_factor=vae_scale_factor,
-            )
-            controlnet.save_pretrained(real_save_dir_path)
-            real_save_dir_path = save_dir_path / f"controlnet_{idx+1}"
-
-        return cls._from_pretrained(
-            model_id=save_dir_path,
-            config=config,
-            model_save_dir=save_dir,
-            **rbln_constructor_kwargs,
-            **kwargs,
-        )
-
-    def _create_runtimes(self, rbln_device_map: Dict[str, int]) -> List[rebel.Runtime]:
-        device_val = rbln_device_map["compiled_model"]
-
-        return [
-            compiled_model.create_runtime(tensor_type="pt", device=device_val)
-            for compiled_model in self.compiled_models
-        ]
+            real_save_path = save_directory + suffix
+            model.save_pretrained(real_save_path)
 
     def forward(
         self,
@@ -243,9 +119,9 @@ class RBLNMultiControlNetModel(RBLNBaseModel):
         return_dict: bool = True,
     ):
         for i, (image, scale, controlnet) in enumerate(zip(controlnet_cond, conditioning_scale, self.nets)):
-            output = controlnet(
+            output = controlnet.model[0](
                 sample=sample.contiguous(),
-                timestep=timestep,
+                timestep=timestep.float(),
                 encoder_hidden_states=encoder_hidden_states,
                 controlnet_cond=image,
                 conditioning_scale=torch.tensor(scale),