optimum-rbln 0.1.9__py3-none-any.whl → 0.1.12__py3-none-any.whl

optimum/rbln/__init__.py

@@ -25,6 +25,7 @@ from typing import TYPE_CHECKING
 
 from transformers.utils import _LazyModule
 
+from .__version__ import __version__
 from .utils import check_version_compats
 
 
@@ -34,11 +35,10 @@ _import_structure = {
         "RBLNBertForQuestionAnswering",
         "RBLNDistilBertForQuestionAnswering",
         "RBLNResNetForImageClassification",
-        "RBLNT5ForConditionalGeneration",
-        "RBLNBartForConditionalGeneration",
         "RBLNXLMRobertaForSequenceClassification",
         "RBLNRobertaForSequenceClassification",
         "RBLNRobertaForMaskedLM",
+        "RBLNViTForImageClassification"
     ],
     "modeling_base": [
         "RBLNBaseModel",
@@ -49,18 +49,36 @@ _import_structure = {
         "RBLNModelForSequenceClassification",
         "RBLNModelForMaskedLM",
     ],
-    "modeling_seq2seq": [
-        "RBLNModelForSeq2SeqLM",
-    ],
     "transformers": [
         "BatchTextIteratorStreamer",
+        "RBLNAutoModel",
+        "RBLNAutoModelForAudioClassification",
+        "RBLNAutoModelForCausalLM",
+        "RBLNAutoModelForCTC",
+        "RBLNAutoModelForDepthEstimation",
+        "RBLNAutoModelForImageClassification",
+        "RBLNAutoModelForMaskedLM",
+        "RBLNAutoModelForQuestionAnswering",
+        "RBLNAutoModelForSeq2SeqLM",
+        "RBLNAutoModelForSequenceClassification",
+        "RBLNAutoModelForSpeechSeq2Seq",
+        "RBLNAutoModelForVision2Seq",
+        "RBLNBartForConditionalGeneration",
+        "RBLNBartModel",
+        "RBLNBertModel",
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelWithProjection",
+        "RBLNCLIPVisionModel",
         "RBLNDPTForDepthEstimation",
+        "RBLNExaoneForCausalLM",
         "RBLNGemmaForCausalLM",
         "RBLNGPT2LMHeadModel",
+        "RBLNQwen2ForCausalLM",
         "RBLNWav2Vec2ForCTC",
         "RBLNLlamaForCausalLM",
+        "RBLNT5ForConditionalGeneration",
+        "RBLNPhiForCausalLM",
+        "RBLNLlavaNextForConditionalGeneration",
         "RBLNMidmLMHeadModel",
         "RBLNMistralForCausalLM",
         "RBLNWhisperForConditionalGeneration",
@@ -80,7 +98,7 @@ _import_structure = {
         "RBLNStableDiffusionXLControlNetPipeline",
         "RBLNStableDiffusionXLControlNetImg2ImgPipeline",
     ],
-    "modeling_config": ["RBLNRuntimeConfig", "RBLNConfig"],
+    "modeling_config": ["RBLNCompileConfig", "RBLNConfig"],
 }
 
 if TYPE_CHECKING:
@@ -100,12 +118,12 @@ if TYPE_CHECKING:
     )
     from .modeling_alias import (
         RBLNASTForAudioClassification,
-        RBLNBartForConditionalGeneration,
         RBLNBertForQuestionAnswering,
         RBLNResNetForImageClassification,
         RBLNRobertaForMaskedLM,
         RBLNRobertaForSequenceClassification,
         RBLNT5ForConditionalGeneration,
+        RBLNViTForImageClassification,
         RBLNXLMRobertaForSequenceClassification,
     )
     from .modeling_base import (
@@ -117,18 +135,38 @@ if TYPE_CHECKING:
         RBLNModelForQuestionAnswering,
         RBLNModelForSequenceClassification,
     )
-    from .modeling_config import RBLNConfig, RBLNRuntimeConfig
-    from .modeling_seq2seq import RBLNModelForSeq2SeqLM
+    from .modeling_config import RBLNCompileConfig, RBLNConfig
     from .transformers import (
         BatchTextIteratorStreamer,
+        RBLNAutoModel,
+        RBLNAutoModelForAudioClassification,
+        RBLNAutoModelForCausalLM,
+        RBLNAutoModelForCTC,
+        RBLNAutoModelForDepthEstimation,
+        RBLNAutoModelForImageClassification,
+        RBLNAutoModelForMaskedLM,
+        RBLNAutoModelForQuestionAnswering,
+        RBLNAutoModelForSeq2SeqLM,
+        RBLNAutoModelForSequenceClassification,
+        RBLNAutoModelForSpeechSeq2Seq,
+        RBLNAutoModelForVision2Seq,
+        RBLNBartForConditionalGeneration,
+        RBLNBartModel,
+        RBLNBertModel,
         RBLNCLIPTextModel,
         RBLNCLIPTextModelWithProjection,
+        RBLNCLIPVisionModel,
         RBLNDPTForDepthEstimation,
+        RBLNExaoneForCausalLM,
         RBLNGemmaForCausalLM,
         RBLNGPT2LMHeadModel,
         RBLNLlamaForCausalLM,
+        RBLNLlavaNextForConditionalGeneration,
         RBLNMidmLMHeadModel,
         RBLNMistralForCausalLM,
+        RBLNPhiForCausalLM,
+        RBLNQwen2ForCausalLM,
+        RBLNT5ForConditionalGeneration,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,
         RBLNXLMRobertaModel,

optimum/rbln/__version__.py

@@ -1 +1 @@
-__version__ = '0.1.9'
+__version__ = '0.1.12'

optimum/rbln/diffusers/models/autoencoder_kl.py

@@ -23,7 +23,7 @@
 
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Union
 
 import rebel
 import torch  # noqa: I001
@@ -34,7 +34,7 @@ from optimum.exporters import TasksManager
 from transformers import AutoConfig, AutoModel, PretrainedConfig
 
 from ...modeling_base import RBLNModel
-from ...modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
+from ...modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
 from ...utils.runtime_utils import RBLNPytorchRuntime
 
 
@@ -58,15 +58,12 @@ class RBLNRuntimeVAEDecoder(RBLNPytorchRuntime):
 
 
 class RBLNAutoencoderKL(RBLNModel):
-    model_type = "rbln_model"
     config_name = "config.json"
-    auto_model_class = AutoModel  # feature extraction
 
     def __post_init__(self, **kwargs):
-        self.dtype = torch.float32
-
-        self.rbln_use_encode = self.rbln_config.meta["rbln_use_encode"]
+        super().__post_init__(**kwargs)
 
+        self.rbln_use_encode = self.rbln_config.model_cfg["use_encode"]
         if self.rbln_use_encode:
             self.encoder = RBLNRuntimeVAEEncoder(runtime=self.model[0], main_input_name="x")
             self.decoder = RBLNRuntimeVAEDecoder(runtime=self.model[1], main_input_name="z")
@@ -81,20 +78,20 @@ class RBLNAutoencoderKL(RBLNModel):
             encoder_model.eval()
             decoder_model.eval()
 
-            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 = cls.compile(encoder_model, rbln_compile_config=rbln_config.compile_cfgs[0])
+            dec_compiled_model = cls.compile(decoder_model, rbln_compile_config=rbln_config.compile_cfgs[1])
 
-            return enc_compiled_model, dec_compiled_model
+            return {"encoder": enc_compiled_model, "decoder": dec_compiled_model}
 
         def compile_text2img():
             decoder_model = _VAEDecoder(model)
             decoder_model.eval()
 
-            dec_compiled_model = cls.compile(decoder_model, rbln_runtime_config=rbln_config["compiled_model"][0])
+            dec_compiled_model = cls.compile(decoder_model, rbln_compile_config=rbln_config.compile_cfgs[0])
 
             return dec_compiled_model
 
-        if rbln_config.meta.get("rbln_use_encode", False):
+        if rbln_config.model_cfg.get("use_encode", False):
             return compile_img2img()
         else:
             return compile_text2img()
@@ -133,23 +130,23 @@ class RBLNAutoencoderKL(RBLNModel):
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
         model_config: "PretrainedConfig",
-        rbln_unet_sample_size: Optional[int] = None,
-        rbln_img_width: Optional[int] = None,
-        rbln_img_height: Optional[int] = None,
-        rbln_batch_size: Optional[int] = None,
-        rbln_use_encode: Optional[bool] = None,
-        rbln_vae_scale_factor: Optional[int] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
-        meta = {}
+        rbln_unet_sample_size = rbln_kwargs.get("unet_sample_size", None)
+        rbln_img_width = rbln_kwargs.get("img_width", None)
+        rbln_img_height = rbln_kwargs.get("img_height", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+        rbln_use_encode = rbln_kwargs.get("use_encode", None)
+        rbln_vae_scale_factor = rbln_kwargs.get("vae_scale_factor", None)
+
         if rbln_batch_size is None:
             rbln_batch_size = 1
 
-        meta["rbln_use_encode"] = rbln_use_encode
-        meta["rbln_batch_size"] = rbln_batch_size
+        model_cfg = {}
 
         if rbln_use_encode:
-            meta["rbln_img_width"] = rbln_img_width
-            meta["rbln_img_height"] = rbln_img_height
+            model_cfg["img_width"] = rbln_img_width
+            model_cfg["img_height"] = rbln_img_height
 
             vae_enc_input_info = [
                 ("x", [rbln_batch_size, model_config.in_channels, rbln_img_height, rbln_img_width], "float32")
@@ -167,20 +164,23 @@ class RBLNAutoencoderKL(RBLNModel):
                 )
             ]
 
-            enc_rbln_runtime_config = RBLNRuntimeConfig(compiled_model_name="encoder", input_info=vae_enc_input_info)
-            dec_rbln_runtime_config = RBLNRuntimeConfig(compiled_model_name="decoder", input_info=vae_dec_input_info)
+            enc_rbln_compile_config = RBLNCompileConfig(compiled_model_name="encoder", input_info=vae_enc_input_info)
+            dec_rbln_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=vae_dec_input_info)
 
-            rbln_config = RBLNConfig.from_rbln_runtime_configs(
-                [enc_rbln_runtime_config, dec_rbln_runtime_config],
-                _rbln_meta=meta,
+            compile_cfgs = [enc_rbln_compile_config, dec_rbln_compile_config]
+            rbln_config = RBLNConfig(
+                rbln_cls=cls.__name__,
+                compile_cfgs=compile_cfgs,
+                rbln_kwargs=rbln_kwargs,
             )
+            rbln_config.model_cfg.update(model_cfg)
             return rbln_config
 
         if rbln_unet_sample_size is None:
             rbln_unet_sample_size = 64
 
-        meta["rbln_unet_sample_size"] = rbln_unet_sample_size
-        vae_config = RBLNRuntimeConfig(
+        model_cfg["unet_sample_size"] = rbln_unet_sample_size
+        vae_config = RBLNCompileConfig(
             input_info=[
                 (
                     "z",
@@ -194,7 +194,12 @@ class RBLNAutoencoderKL(RBLNModel):
                 )
             ],
         )
-        rbln_config = RBLNConfig.from_rbln_runtime_configs([vae_config], _rbln_meta=meta)
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[vae_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+        rbln_config.model_cfg.update(model_cfg)
         return rbln_config
 
     @classmethod

optimum/rbln/diffusers/models/controlnet.py

@@ -23,7 +23,7 @@
 
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import torch
 from diffusers import ControlNetModel
@@ -31,7 +31,7 @@ from optimum.exporters import TasksManager
 from transformers import AutoConfig, AutoModel, PretrainedConfig
 
 from ...modeling_base import RBLNModel
-from ...modeling_config import RBLNConfig, RBLNRuntimeConfig
+from ...modeling_config import RBLNCompileConfig, RBLNConfig
 
 
 if TYPE_CHECKING:
@@ -105,13 +105,10 @@ class _ControlNetModel_Cross_Attention(torch.nn.Module):
 
 
 class RBLNControlNetModel(RBLNModel):
-    model_type = "rbln_model"
-    auto_model_class = AutoModel  # feature extraction
-
     def __post_init__(self, **kwargs):
-        self.dtype = torch.float32
+        super().__post_init__(**kwargs)
         self.use_encoder_hidden_states = any(
-            item[0] == "encoder_hidden_states" for item in self.rbln_config["compiled_model"][0].input_info
+            item[0] == "encoder_hidden_states" for item in self.rbln_config.compile_cfgs[0].input_info
         )
 
     @classmethod
@@ -121,9 +118,6 @@ 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
@@ -155,14 +149,14 @@ class RBLNControlNetModel(RBLNModel):
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
         model_config: "PretrainedConfig",
-        rbln_max_seq_len: Optional[int] = None,
-        rbln_text_model_hidden_size: Optional[int] = None,
-        rbln_batch_size: Optional[int] = None,
-        rbln_img_width: Optional[int] = None,
-        rbln_img_height: Optional[int] = None,
-        rbln_vae_scale_factor: Optional[int] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
-        meta = {"type": "controlnet"}
+        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
+        rbln_text_model_hidden_size = rbln_kwargs.get("text_model_hidden_size", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+        rbln_img_width = rbln_kwargs.get("img_width", None)
+        rbln_img_height = rbln_kwargs.get("img_height", None)
+        rbln_vae_scale_factor = rbln_kwargs.get("vae_scale_factor", None)
 
         if rbln_batch_size is None:
             rbln_batch_size = 1
@@ -170,28 +164,29 @@ class RBLNControlNetModel(RBLNModel):
         if rbln_max_seq_len is None:
             rbln_max_seq_len = 77
 
+        if rbln_img_width is None or rbln_img_height is None or rbln_vae_scale_factor is None:
+            raise ValueError("rbln_img_width, rbln_img_height, and rbln_vae_scale_factor must be provided")
+
         input_width = rbln_img_width // rbln_vae_scale_factor
         input_height = rbln_img_height // rbln_vae_scale_factor
 
-        rbln_runtime_config = RBLNRuntimeConfig(
-            input_info=[
-                (
-                    "sample",
-                    [
-                        rbln_batch_size,
-                        model_config.in_channels,
-                        input_height,
-                        input_width,
-                    ],
-                    "float32",
-                ),
-                ("timestep", [], "float32"),
-            ],
-            batch_size=rbln_batch_size,
-        )
+        input_info = [
+            (
+                "sample",
+                [
+                    rbln_batch_size,
+                    model_config.in_channels,
+                    input_height,
+                    input_width,
+                ],
+                "float32",
+            ),
+            ("timestep", [], "float32"),
+        ]
+
         use_encoder_hidden_states = any(element != "DownBlock2D" for element in model_config.down_block_types)
         if use_encoder_hidden_states:
-            rbln_runtime_config.input_info.append(
+            input_info.append(
                 (
                     "encoder_hidden_states",
                     [
@@ -202,19 +197,34 @@ class RBLNControlNetModel(RBLNModel):
                     "float32",
                 )
             )
-        rbln_runtime_config.input_info.append(
-            ("controlnet_cond", [rbln_batch_size, 3, rbln_img_height, rbln_img_width], "float32")
-        )
-        rbln_runtime_config.input_info.append(("conditioning_scale", [], "float32"))
+
+        input_info.append(("controlnet_cond", [rbln_batch_size, 3, rbln_img_height, rbln_img_width], "float32"))
+        input_info.append(("conditioning_scale", [], "float32"))
+
         if hasattr(model_config, "addition_embed_type") and model_config.addition_embed_type == "text_time":
             if rbln_text_model_hidden_size is None:
                 rbln_text_model_hidden_size = 768
-            rbln_runtime_config.input_info.append(
-                ("text_embeds", [rbln_batch_size, rbln_text_model_hidden_size], "float32")
-            )
-            rbln_runtime_config.input_info.append(("time_ids", [rbln_batch_size, 6], "float32"))
+            input_info.append(("text_embeds", [rbln_batch_size, rbln_text_model_hidden_size], "float32"))
+            input_info.append(("time_ids", [rbln_batch_size, 6], "float32"))
+
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
+
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+
+        rbln_config.model_cfg.update(
+            {
+                "max_seq_len": rbln_max_seq_len,
+                "batch_size": rbln_batch_size,
+                "img_width": rbln_img_width,
+                "img_height": rbln_img_height,
+                "vae_scale_factor": rbln_vae_scale_factor,
+            }
+        )
 
-        rbln_config = RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
         return rbln_config
 
     def forward(

optimum/rbln/diffusers/models/unet_2d_condition.py

@@ -32,7 +32,7 @@ from optimum.exporters import TasksManager
 from transformers import AutoConfig, AutoModel, PretrainedConfig
 
 from ...modeling_base import RBLNModel
-from ...modeling_config import RBLNConfig, RBLNRuntimeConfig
+from ...modeling_config import RBLNCompileConfig, RBLNConfig
 
 
 if TYPE_CHECKING:
@@ -126,12 +126,9 @@ class _UNet_SDXL(torch.nn.Module):
 
 
 class RBLNUNet2DConditionModel(RBLNModel):
-    model_type = "rbln_model"
-    auto_model_class = AutoModel  # feature extraction
-
     def __post_init__(self, **kwargs):
-        self.dtype = torch.float32
-        self.in_features = self.rbln_config.meta.get("in_features", None)
+        super().__post_init__(**kwargs)
+        self.in_features = self.rbln_config.model_cfg.get("in_features", None)
         if self.in_features is not None:
 
             @dataclass
@@ -183,31 +180,31 @@ class RBLNUNet2DConditionModel(RBLNModel):
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
         model_config: "PretrainedConfig",
-        rbln_max_seq_len: Optional[int] = None,
-        rbln_text_model_hidden_size: Optional[int] = None,
-        rbln_batch_size: Optional[int] = None,
-        rbln_in_features: Optional[int] = None,
-        rbln_use_encode: Optional[bool] = None,
-        rbln_img_width: Optional[int] = None,
-        rbln_img_height: Optional[int] = None,
-        rbln_vae_scale_factor: Optional[int] = None,
-        rbln_is_controlnet: Optional[bool] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
-        meta = {"type": "unet"}
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
+        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
+        rbln_text_model_hidden_size = rbln_kwargs.get("text_model_hidden_size", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+        rbln_in_features = rbln_kwargs.get("in_features", None)
+        rbln_use_encode = rbln_kwargs.get("use_encode", None)
+        rbln_img_width = rbln_kwargs.get("img_width", None)
+        rbln_img_height = rbln_kwargs.get("img_height", None)
+        rbln_vae_scale_factor = rbln_kwargs.get("vae_scale_factor", None)
+        rbln_is_controlnet = rbln_kwargs.get("is_controlnet", None)
 
         if rbln_max_seq_len is None:
             rbln_max_seq_len = 77
-
-        meta["rbln_use_encode"] = rbln_use_encode
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
 
         if rbln_use_encode:
-            # FIXME :: robust img shape getter
+            if rbln_img_width is None or rbln_img_height is None or rbln_vae_scale_factor is None:
+                raise ValueError(
+                    "rbln_img_width, rbln_img_height, and rbln_vae_scale_factor must be provided when rbln_use_encode is True"
+                )
             input_width = rbln_img_width // rbln_vae_scale_factor
             input_height = rbln_img_height // rbln_vae_scale_factor
         else:
-            # FIXME :: model_config.sample_size can be tuple or list
             input_width, input_height = model_config.sample_size, model_config.sample_size
 
         input_info = [
@@ -232,6 +229,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
                 "float32",
             ),
         ]
+
         if rbln_is_controlnet:
             if len(model_config.block_out_channels) > 0:
                 input_info.extend(
@@ -304,24 +302,35 @@ class RBLNUNet2DConditionModel(RBLNModel):
                 )
             )
 
-        rbln_runtime_config = RBLNRuntimeConfig(
-            input_info=input_info,
-            batch_size=rbln_batch_size,
-        )
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
 
         if hasattr(model_config, "addition_embed_type") and model_config.addition_embed_type == "text_time":
-            # In case of sdxl
             if rbln_text_model_hidden_size is None:
                 rbln_text_model_hidden_size = 768
             if rbln_in_features is None:
                 rbln_in_features = model_config.projection_class_embeddings_input_dim
-            meta["in_features"] = rbln_in_features
-            rbln_runtime_config.input_info.append(
+            rbln_compile_config.input_info.append(
                 ("text_embeds", [rbln_batch_size, rbln_text_model_hidden_size], "float32")
             )
-            rbln_runtime_config.input_info.append(("time_ids", [rbln_batch_size, 6], "float32"))
+            rbln_compile_config.input_info.append(("time_ids", [rbln_batch_size, 6], "float32"))
+
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            rbln_kwargs=rbln_kwargs,
+        )
+
+        rbln_config.model_cfg.update(
+            {
+                "max_seq_len": rbln_max_seq_len,
+                "batch_size": rbln_batch_size,
+                "use_encode": rbln_use_encode,
+            }
+        )
+
+        if rbln_in_features is not None:
+            rbln_config.model_cfg["in_features"] = rbln_in_features
 
-        rbln_config = RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
         return rbln_config
 
     def forward(

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

@@ -102,6 +102,10 @@ class RBLNMultiControlNetModel(RBLNModel):
             real_save_path = save_directory + suffix
             model.save_pretrained(real_save_path)
 
+    @classmethod
+    def _get_rbln_config(cls, **rbln_config_kwargs):
+        pass
+
     def forward(
         self,
         sample: torch.FloatTensor,