optimum-rbln 0.1.1__py3-none-any.whl → 0.1.7__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (37) hide show
  1. optimum/rbln/__init__.py +9 -0
  2. optimum/rbln/__version__.py +1 -1
  3. optimum/rbln/diffusers/models/autoencoder_kl.py +16 -98
  4. optimum/rbln/diffusers/models/unet_2d_condition.py +1 -1
  5. optimum/rbln/diffusers/pipelines/controlnet/multicontrolnet.py +9 -11
  6. optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet.py +8 -0
  7. optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +9 -0
  8. optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +3 -0
  9. optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +8 -0
  10. optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +8 -0
  11. optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +9 -0
  12. optimum/rbln/modeling_base.py +175 -103
  13. optimum/rbln/modeling_seq2seq.py +58 -132
  14. optimum/rbln/transformers/__init__.py +4 -0
  15. optimum/rbln/transformers/models/__init__.py +2 -0
  16. optimum/rbln/transformers/models/clip/modeling_clip.py +0 -1
  17. optimum/rbln/transformers/models/dpt/__init__.py +24 -0
  18. optimum/rbln/transformers/models/dpt/modeling_dpt.py +89 -0
  19. optimum/rbln/transformers/models/gpt2/gpt2_architecture.py +24 -33
  20. optimum/rbln/transformers/models/gpt2/modeling_gpt2.py +52 -124
  21. optimum/rbln/transformers/models/llama/llama_architecture.py +62 -33
  22. optimum/rbln/transformers/models/llama/llama_architecture_cb.py +764 -0
  23. optimum/rbln/transformers/models/llama/modeling_llama.py +208 -140
  24. optimum/rbln/transformers/models/midm/__init__.py +32 -0
  25. optimum/rbln/transformers/models/midm/hf_hub_cached/configuration_midm.py +22 -0
  26. optimum/rbln/transformers/models/midm/hf_hub_cached/midm_bitext_tokenization.py +303 -0
  27. optimum/rbln/transformers/models/midm/hf_hub_cached/modeling_midm.py +1473 -0
  28. optimum/rbln/transformers/models/midm/hf_hub_cached/rotary_position_embedding.py +98 -0
  29. optimum/rbln/transformers/models/midm/midm_architecture.py +506 -0
  30. optimum/rbln/transformers/models/midm/modeling_midm.py +390 -0
  31. optimum/rbln/transformers/models/whisper/modeling_whisper.py +53 -123
  32. optimum/rbln/utils/__init__.py +1 -1
  33. optimum/rbln/utils/import_utils.py +46 -0
  34. {optimum_rbln-0.1.1.dist-info → optimum_rbln-0.1.7.dist-info}/METADATA +17 -50
  35. {optimum_rbln-0.1.1.dist-info → optimum_rbln-0.1.7.dist-info}/RECORD +37 -27
  36. {optimum_rbln-0.1.1.dist-info → optimum_rbln-0.1.7.dist-info}/WHEEL +1 -1
  37. {optimum_rbln-0.1.1.dist-info → optimum_rbln-0.1.7.dist-info}/licenses/LICENSE +0 -0
optimum/rbln/__init__.py CHANGED
@@ -25,6 +25,8 @@ from typing import TYPE_CHECKING
25
25
 
26
26
  from transformers.utils import _LazyModule
27
27
 
28
+ from .utils import check_version_compats
29
+
28
30
 
29
31
  _import_structure = {
30
32
  "modeling_alias": [
@@ -48,9 +50,11 @@ _import_structure = {
48
50
  "BatchTextIteratorStreamer",
49
51
  "RBLNCLIPTextModel",
50
52
  "RBLNCLIPTextModelWithProjection",
53
+ "RBLNDPTForDepthEstimation",
51
54
  "RBLNGPT2LMHeadModel",
52
55
  "RBLNWav2Vec2ForCTC",
53
56
  "RBLNLlamaForCausalLM",
57
+ "RBLNMidmLMHeadModel",
54
58
  "RBLNWhisperForConditionalGeneration",
55
59
  ],
56
60
  "diffusers": [
@@ -105,8 +109,10 @@ if TYPE_CHECKING:
105
109
  BatchTextIteratorStreamer,
106
110
  RBLNCLIPTextModel,
107
111
  RBLNCLIPTextModelWithProjection,
112
+ RBLNDPTForDepthEstimation,
108
113
  RBLNGPT2LMHeadModel,
109
114
  RBLNLlamaForCausalLM,
115
+ RBLNMidmLMHeadModel,
110
116
  RBLNWav2Vec2ForCTC,
111
117
  RBLNWhisperForConditionalGeneration,
112
118
  )
@@ -119,3 +125,6 @@ else:
119
125
  _import_structure,
120
126
  module_spec=__spec__,
121
127
  )
128
+
129
+
130
+ check_version_compats()
@@ -1 +1 @@
1
- __version__ = '0.1.1'
1
+ __version__ = '0.1.7'
@@ -23,7 +23,6 @@
23
23
 
24
24
  import logging
25
25
  from pathlib import Path
26
- from tempfile import TemporaryDirectory
27
26
  from typing import TYPE_CHECKING, Dict, List, Optional, Union
28
27
 
29
28
  import rebel
@@ -37,7 +36,6 @@ from transformers import AutoConfig, AutoModel, PretrainedConfig
37
36
  from ...modeling_base import RBLNModel
38
37
  from ...modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
39
38
  from ...utils.runtime_utils import RBLNPytorchRuntime
40
- from ...utils.save_utils import maybe_save_preprocessors
41
39
 
42
40
 
43
41
  logger = logging.getLogger(__name__)
@@ -70,73 +68,13 @@ class RBLNAutoencoderKL(RBLNModel):
70
68
  self.rbln_use_encode = self.rbln_config.meta["rbln_use_encode"]
71
69
 
72
70
  if self.rbln_use_encode:
73
- self.encoder = RBLNRuntimeVAEEncoder(runtime=self.runtimes[0], main_input_name="x")
74
- self.decoder = RBLNRuntimeVAEDecoder(runtime=self.runtimes[1], main_input_name="z")
71
+ self.encoder = RBLNRuntimeVAEEncoder(runtime=self.model[0], main_input_name="x")
72
+ self.decoder = RBLNRuntimeVAEDecoder(runtime=self.model[1], main_input_name="z")
75
73
  else:
76
- self.decoder = RBLNRuntimeVAEDecoder(runtime=self.runtimes[0], main_input_name="z")
74
+ self.decoder = RBLNRuntimeVAEDecoder(runtime=self.model[0], main_input_name="z")
77
75
 
78
76
  @classmethod
79
- @torch.no_grad()
80
- def _export(
81
- cls,
82
- model_id: str,
83
- config: "PretrainedConfig",
84
- use_auth_token: Optional[Union[bool, str]] = None,
85
- revision: Optional[str] = None,
86
- force_download: bool = False,
87
- cache_dir: Optional[str] = None,
88
- subfolder: str = "",
89
- local_files_only: bool = False,
90
- trust_remote_code: bool = False,
91
- model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
92
- **kwargs,
93
- ) -> "RBLNAutoencoderKL":
94
- task = kwargs.pop("task", None)
95
- if task is None:
96
- task = TasksManager.infer_task_from_model(cls.auto_model_class)
97
-
98
- if model_save_dir is None:
99
- save_dir = TemporaryDirectory()
100
- save_dir_path = Path(save_dir.name)
101
- else:
102
- save_dir = model_save_dir
103
- if isinstance(save_dir, TemporaryDirectory):
104
- save_dir_path = Path(model_save_dir.name)
105
- else:
106
- save_dir_path = Path(model_save_dir)
107
- save_dir_path.mkdir(exist_ok=True)
108
-
109
- rbln_config_kwargs, rbln_constructor_kwargs = cls.pop_rbln_kwargs_from_kwargs(kwargs)
110
-
111
- model: AutoencoderKL = TasksManager.get_model_from_task(
112
- task=None,
113
- model_name_or_path=model_id,
114
- subfolder=subfolder,
115
- revision=revision,
116
- framework="pt",
117
- cache_dir=cache_dir,
118
- use_auth_token=use_auth_token,
119
- local_files_only=local_files_only,
120
- force_download=force_download,
121
- trust_remote_code=trust_remote_code,
122
- **kwargs,
123
- )
124
-
125
- if config is None:
126
- config = model.config
127
-
128
- if not isinstance(config, PretrainedConfig): # diffusers config
129
- config = PretrainedConfig(**config)
130
-
131
- config.save_pretrained(save_dir_path / subfolder)
132
- preprocessors = maybe_save_preprocessors(model_id, save_dir_path, src_subfolder=subfolder)
133
-
134
- # Get compilation arguments
135
- if rbln_config_kwargs.get("rbln_config", None) is None:
136
- rbln_config = cls.get_rbln_config(
137
- preprocessors=preprocessors, model_config=model.config, **rbln_config_kwargs
138
- )
139
-
77
+ def get_compiled_model(cls, model, rbln_config: RBLNConfig):
140
78
  def compile_img2img():
141
79
  encoder_model = _VAEEncoder(model)
142
80
  decoder_model = _VAEDecoder(model)
@@ -146,12 +84,7 @@ class RBLNAutoencoderKL(RBLNModel):
146
84
  enc_compiled_model = cls.compile(encoder_model, rbln_runtime_config=rbln_config["encoder"][0])
147
85
  dec_compiled_model = cls.compile(decoder_model, rbln_runtime_config=rbln_config["decoder"][0])
148
86
 
149
- enc_compiled_model.save(
150
- save_dir_path / subfolder / f"{rbln_config['encoder'][0].compiled_model_name}.rbln"
151
- )
152
- dec_compiled_model.save(
153
- save_dir_path / subfolder / f"{rbln_config['decoder'][0].compiled_model_name}.rbln"
154
- )
87
+ return enc_compiled_model, dec_compiled_model
155
88
 
156
89
  def compile_text2img():
157
90
  decoder_model = _VAEDecoder(model)
@@ -159,30 +92,12 @@ class RBLNAutoencoderKL(RBLNModel):
159
92
 
160
93
  dec_compiled_model = cls.compile(decoder_model, rbln_runtime_config=rbln_config["compiled_model"][0])
161
94
 
162
- dec_compiled_model.save(
163
- save_dir_path / subfolder / f"{rbln_config['compiled_model'][0].compiled_model_name}.rbln"
164
- )
95
+ return dec_compiled_model
165
96
 
166
- if rbln_config_kwargs.get("rbln_use_encode"):
167
- compile_img2img()
97
+ if rbln_config.meta.get("rbln_use_encode", False):
98
+ return compile_img2img()
168
99
  else:
169
- compile_text2img()
170
-
171
- rbln_config.save(save_dir_path / subfolder)
172
-
173
- return cls._from_pretrained(
174
- model_id=save_dir_path,
175
- config=config,
176
- model_save_dir=save_dir,
177
- use_auth_token=use_auth_token,
178
- revision=revision,
179
- force_download=force_download,
180
- cache_dir=cache_dir,
181
- subfolder=subfolder,
182
- local_files_only=local_files_only,
183
- **rbln_constructor_kwargs,
184
- **kwargs,
185
- )
100
+ return compile_text2img()
186
101
 
187
102
  @classmethod
188
103
  def from_pretrained(cls, *args, **kwargs):
@@ -282,15 +197,18 @@ class RBLNAutoencoderKL(RBLNModel):
282
197
  rbln_config = RBLNConfig.from_rbln_runtime_configs([vae_config], _rbln_meta=meta)
283
198
  return rbln_config
284
199
 
285
- def _create_runtimes(self, rbln_device_map: Dict[str, int]) -> List[rebel.Runtime]:
286
- if len(self.compiled_models) == 1:
200
+ @classmethod
201
+ def _create_runtimes(
202
+ cls, compiled_models: List[rebel.RBLNCompiledModel], rbln_device_map: Dict[str, int]
203
+ ) -> List[rebel.Runtime]:
204
+ if len(compiled_models) == 1:
287
205
  device_val = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
288
- return [self.compiled_models[0].create_runtime(tensor_type="pt", device=device_val)]
206
+ return [compiled_models[0].create_runtime(tensor_type="pt", device=device_val)]
289
207
 
290
208
  device_vals = [rbln_device_map["encoder"], rbln_device_map["decoder"]]
291
209
  return [
292
210
  compiled_model.create_runtime(tensor_type="pt", device=device_val)
293
- for compiled_model, device_val in zip(self.compiled_models, device_vals)
211
+ for compiled_model, device_val in zip(compiled_models, device_vals)
294
212
  ]
295
213
 
296
214
  def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
@@ -314,7 +314,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
314
314
  if rbln_text_model_hidden_size is None:
315
315
  rbln_text_model_hidden_size = 768
316
316
  if rbln_in_features is None:
317
- rbln_in_features = 2816
317
+ rbln_in_features = model_config.projection_class_embeddings_input_dim
318
318
  meta["in_features"] = rbln_in_features
319
319
  rbln_runtime_config.input_info.append(
320
320
  ("text_embeds", [rbln_batch_size, rbln_text_model_hidden_size], "float32")
@@ -35,7 +35,7 @@ from optimum.exporters import TasksManager
35
35
  from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
36
36
 
37
37
  from ....modeling_base import RBLNBaseModel
38
- from ....modeling_config import RBLNConfig
38
+ from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig
39
39
  from ...models.controlnet import RBLNControlNetModel
40
40
 
41
41
 
@@ -111,7 +111,6 @@ class RBLNMultiControlNetModel(RBLNBaseModel):
111
111
  local_files_only: bool = False,
112
112
  **kwargs,
113
113
  ) -> RBLNBaseModel:
114
-
115
114
  if isinstance(model_id, str):
116
115
  model_path = Path(model_id)
117
116
  else:
@@ -150,6 +149,7 @@ class RBLNMultiControlNetModel(RBLNBaseModel):
150
149
  )
151
150
 
152
151
  def _save_pretrained(self, save_directory: Union[str, Path]):
152
+ # TODO(kblee) : 확인 부탁드립니다
153
153
  idx = 0
154
154
  real_save_dir_path = save_directory
155
155
  for compiled_model in self.compiled_models:
@@ -176,7 +176,6 @@ class RBLNMultiControlNetModel(RBLNBaseModel):
176
176
  trust_remote_code: bool = False,
177
177
  **kwargs,
178
178
  ) -> "RBLNMultiControlNetModel":
179
-
180
179
  task = kwargs.pop("task", None)
181
180
  if task is None:
182
181
  task = TasksManager.infer_task_from_model(cls.auto_model_class)
@@ -219,13 +218,12 @@ class RBLNMultiControlNetModel(RBLNBaseModel):
219
218
  **kwargs,
220
219
  )
221
220
 
222
- def _create_runtimes(self, rbln_device_map: Dict[str, int]) -> List[rebel.Runtime]:
223
- device_val = rbln_device_map["compiled_model"]
224
-
225
- return [
226
- compiled_model.create_runtime(tensor_type="pt", device=device_val)
227
- for compiled_model in self.compiled_models
228
- ]
221
+ @classmethod
222
+ def _create_runtimes(
223
+ cls, compiled_models: List[rebel.RBLNCompiledModel], rbln_device_map: Dict[str, int]
224
+ ) -> List[rebel.Runtime]:
225
+ device = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
226
+ return [compiled_model.create_runtime(tensor_type="pt", device=device) for compiled_model in compiled_models]
229
227
 
230
228
  def forward(
231
229
  self,
@@ -245,7 +243,7 @@ class RBLNMultiControlNetModel(RBLNBaseModel):
245
243
  for i, (image, scale, controlnet) in enumerate(zip(controlnet_cond, conditioning_scale, self.nets)):
246
244
  output = controlnet(
247
245
  sample=sample.contiguous(),
248
- timestep=timestep,
246
+ timestep=timestep.float(),
249
247
  encoder_hidden_states=encoder_hidden_states,
250
248
  controlnet_cond=image,
251
249
  conditioning_scale=torch.tensor(scale),
@@ -161,6 +161,14 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
161
161
 
162
162
  model.models = [vae.model[0], text_encoder.model[0], unet.model[0], controlnet.model[0]]
163
163
 
164
+ if rbln_constructor_kwargs.pop("rbln_optimize_host_memory", None) is False:
165
+ model.compiled_models = [
166
+ vae.compiled_models[0],
167
+ text_encoder.compiled_models[0],
168
+ unet.compiled_models[0],
169
+ controlnet.compiled_models[0],
170
+ ]
171
+
164
172
  return model
165
173
 
166
174
  def check_inputs(
@@ -160,6 +160,15 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
160
160
 
161
161
  model.models = [vae.model[0], vae.model[1], text_encoder.model[0], unet.model[0], controlnet.model[0]]
162
162
 
163
+ if rbln_constructor_kwargs.pop("rbln_optimize_host_memory", None) is False:
164
+ model.compiled_models = [
165
+ vae.compiled_models[0],
166
+ vae.compiled_models[1],
167
+ text_encoder.compiled_models[0],
168
+ unet.compiled_models[0],
169
+ controlnet.compiled_models[0],
170
+ ]
171
+
163
172
  return model
164
173
 
165
174
  def check_inputs(
@@ -122,4 +122,7 @@ class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
122
122
 
123
123
  model.models = [vae.model[0], text_encoder.model[0], unet.model[0]]
124
124
 
125
+ if rbln_constructor_kwargs.pop("rbln_optimize_host_memory", None) is False:
126
+ model.compiled_models = [vae.compiled_models[0], text_encoder.compiled_models[0], unet.compiled_models[0]]
127
+
125
128
  return model
@@ -126,4 +126,12 @@ class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
126
126
  # vae encoder, vae decoder, text_encoder, unet
127
127
  model.models = [vae.model[0], vae.model[1], text_encoder.model[0], unet.model[0]]
128
128
 
129
+ if rbln_constructor_kwargs.pop("rbln_optimize_host_memory", None) is False:
130
+ model.compiled_models = [
131
+ vae.compiled_models[0],
132
+ vae.compiled_models[1],
133
+ text_encoder.compiled_models[0],
134
+ unet.compiled_models[0],
135
+ ]
136
+
129
137
  return model
@@ -122,4 +122,12 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
122
122
 
123
123
  model.models = [vae.model[0], unet.model[0], text_encoder.model[0], text_encoder_2.model[0]]
124
124
 
125
+ if rbln_constructor_kwargs.pop("rbln_optimize_host_memory", None) is False:
126
+ model.compiled_models = [
127
+ vae.compiled_models[0],
128
+ unet.compiled_models[0],
129
+ text_encoder.compiled_models[0],
130
+ text_encoder_2.compiled_models[0],
131
+ ]
132
+
125
133
  return model
@@ -124,4 +124,13 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
124
124
 
125
125
  model.models = [vae.model[0], vae.model[1], unet.model[0], text_encoder.model[0], text_encoder_2.model[0]]
126
126
 
127
+ if rbln_constructor_kwargs.pop("rbln_optimize_host_memory", None) is False:
128
+ model.compiled_models = [
129
+ vae.compiled_models[0],
130
+ vae.compiled_models[1],
131
+ unet.compiled_models[0],
132
+ text_encoder.compiled_models[0],
133
+ text_encoder_2.compiled_models[0],
134
+ ]
135
+
127
136
  return model