optimum-rbln 0.8.2rc0__py3-none-any.whl → 0.8.3__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.

Potentially problematic release.


This version of optimum-rbln might be problematic. Click here for more details.

Files changed (105) hide show
  1. optimum/rbln/__init__.py +32 -9
  2. optimum/rbln/__version__.py +16 -3
  3. optimum/rbln/configuration_utils.py +20 -4
  4. optimum/rbln/diffusers/__init__.py +7 -0
  5. optimum/rbln/diffusers/configurations/models/configuration_autoencoder_kl.py +2 -2
  6. optimum/rbln/diffusers/configurations/models/configuration_autoencoder_kl_cosmos.py +2 -2
  7. optimum/rbln/diffusers/configurations/models/configuration_controlnet.py +2 -2
  8. optimum/rbln/diffusers/configurations/models/configuration_prior_transformer.py +2 -2
  9. optimum/rbln/diffusers/configurations/models/configuration_transformer_cosmos.py +2 -2
  10. optimum/rbln/diffusers/configurations/models/configuration_transformer_sd3.py +2 -2
  11. optimum/rbln/diffusers/configurations/models/configuration_unet_2d_condition.py +2 -2
  12. optimum/rbln/diffusers/configurations/models/configuration_vq_model.py +2 -2
  13. optimum/rbln/diffusers/configurations/pipelines/configuration_controlnet.py +3 -3
  14. optimum/rbln/diffusers/configurations/pipelines/configuration_cosmos.py +2 -2
  15. optimum/rbln/diffusers/configurations/pipelines/configuration_kandinsky2_2.py +4 -4
  16. optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion.py +2 -2
  17. optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.py +2 -2
  18. optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_xl.py +2 -2
  19. optimum/rbln/diffusers/modeling_diffusers.py +1 -1
  20. optimum/rbln/diffusers/models/__init__.py +3 -13
  21. optimum/rbln/diffusers/pipelines/__init__.py +11 -5
  22. optimum/rbln/diffusers/pipelines/auto_pipeline.py +237 -0
  23. optimum/rbln/diffusers/pipelines/cosmos/configuration_cosmos_guardrail.py +11 -6
  24. optimum/rbln/diffusers/pipelines/cosmos/cosmos_guardrail.py +14 -18
  25. optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py +1 -1
  26. optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py +1 -1
  27. optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py +1 -6
  28. optimum/rbln/modeling.py +3 -2
  29. optimum/rbln/modeling_base.py +29 -4
  30. optimum/rbln/ops/attn.py +158 -0
  31. optimum/rbln/ops/flash_attn.py +166 -0
  32. optimum/rbln/transformers/__init__.py +24 -0
  33. optimum/rbln/transformers/configuration_generic.py +6 -4
  34. optimum/rbln/transformers/modeling_generic.py +13 -8
  35. optimum/rbln/transformers/modeling_outputs.py +37 -0
  36. optimum/rbln/transformers/models/__init__.py +31 -16
  37. optimum/rbln/transformers/models/auto/__init__.py +2 -0
  38. optimum/rbln/transformers/models/auto/modeling_auto.py +14 -0
  39. optimum/rbln/transformers/models/bart/bart_architecture.py +1 -3
  40. optimum/rbln/transformers/models/bart/configuration_bart.py +2 -0
  41. optimum/rbln/transformers/models/bert/bert_architecture.py +16 -0
  42. optimum/rbln/transformers/models/bert/modeling_bert.py +8 -4
  43. optimum/rbln/transformers/models/blip_2/configuration_blip_2.py +2 -2
  44. optimum/rbln/transformers/models/blip_2/modeling_blip_2.py +7 -6
  45. optimum/rbln/transformers/models/clip/configuration_clip.py +3 -3
  46. optimum/rbln/transformers/models/colpali/colpali_architecture.py +1 -4
  47. optimum/rbln/transformers/models/colpali/configuration_colpali.py +2 -2
  48. optimum/rbln/transformers/models/colpali/modeling_colpali.py +2 -10
  49. optimum/rbln/transformers/models/decoderonly/configuration_decoderonly.py +43 -174
  50. optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py +101 -91
  51. optimum/rbln/transformers/models/decoderonly/decoderonly_runtime_utils.py +450 -0
  52. optimum/rbln/transformers/models/decoderonly/generation_decoderonly.py +88 -0
  53. optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py +296 -986
  54. optimum/rbln/transformers/models/depth_anything/__init__.py +16 -0
  55. optimum/rbln/transformers/models/depth_anything/configuration_depth_anything.py +24 -0
  56. optimum/rbln/transformers/models/depth_anything/modeling_depth_anything.py +25 -0
  57. optimum/rbln/transformers/models/gemma/gemma_architecture.py +1 -4
  58. optimum/rbln/transformers/models/gemma/modeling_gemma.py +9 -0
  59. optimum/rbln/transformers/models/gemma3/configuration_gemma3.py +3 -3
  60. optimum/rbln/transformers/models/gemma3/gemma3_runtime_utils.py +217 -0
  61. optimum/rbln/transformers/models/gemma3/modeling_gemma3.py +25 -251
  62. optimum/rbln/transformers/models/gpt2/modeling_gpt2.py +2 -0
  63. optimum/rbln/transformers/models/grounding_dino/__init__.py +10 -0
  64. optimum/rbln/transformers/models/grounding_dino/configuration_grounding_dino.py +86 -0
  65. optimum/rbln/transformers/models/grounding_dino/grounding_dino_architecture.py +507 -0
  66. optimum/rbln/transformers/models/grounding_dino/modeling_grounding_dino.py +1032 -0
  67. optimum/rbln/transformers/models/idefics3/configuration_idefics3.py +2 -2
  68. optimum/rbln/transformers/models/idefics3/modeling_idefics3.py +3 -9
  69. optimum/rbln/transformers/models/llama/modeling_llama.py +12 -3
  70. optimum/rbln/transformers/models/llava/configuration_llava.py +2 -2
  71. optimum/rbln/transformers/models/llava/modeling_llava.py +53 -14
  72. optimum/rbln/transformers/models/llava_next/configuration_llava_next.py +2 -2
  73. optimum/rbln/transformers/models/llava_next/modeling_llava_next.py +6 -16
  74. optimum/rbln/transformers/models/opt/modeling_opt.py +2 -30
  75. optimum/rbln/transformers/models/pegasus/configuration_pegasus.py +4 -0
  76. optimum/rbln/transformers/models/pegasus/modeling_pegasus.py +2 -0
  77. optimum/rbln/transformers/models/pegasus/pegasus_architecture.py +1 -3
  78. optimum/rbln/transformers/models/pixtral/configuration_pixtral.py +2 -2
  79. optimum/rbln/transformers/models/pixtral/modeling_pixtral.py +1 -4
  80. optimum/rbln/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py +3 -3
  81. optimum/rbln/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py +6 -15
  82. optimum/rbln/transformers/models/qwen2_5_vl/qwen2_5_vl_architecture.py +4 -7
  83. optimum/rbln/transformers/models/qwen3/modeling_qwen3.py +77 -3
  84. optimum/rbln/transformers/models/qwen3/qwen3_architecture.py +1 -4
  85. optimum/rbln/transformers/models/seq2seq/configuration_seq2seq.py +19 -2
  86. optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py +20 -1
  87. optimum/rbln/transformers/models/siglip/__init__.py +2 -6
  88. optimum/rbln/transformers/models/siglip/modeling_siglip.py +2 -2
  89. optimum/rbln/transformers/models/swin/__init__.py +16 -0
  90. optimum/rbln/transformers/models/swin/configuration_swin.py +42 -0
  91. optimum/rbln/transformers/models/swin/modeling_swin.py +341 -0
  92. optimum/rbln/transformers/models/t5/configuration_t5.py +2 -0
  93. optimum/rbln/transformers/models/t5/t5_architecture.py +8 -1
  94. optimum/rbln/transformers/models/time_series_transformer/configuration_time_series_transformer.py +2 -2
  95. optimum/rbln/transformers/models/time_series_transformer/modeling_time_series_transformer.py +4 -14
  96. optimum/rbln/transformers/models/whisper/configuration_whisper.py +10 -2
  97. optimum/rbln/transformers/models/whisper/modeling_whisper.py +20 -1
  98. optimum/rbln/transformers/models/xlm_roberta/__init__.py +2 -8
  99. optimum/rbln/transformers/utils/rbln_quantization.py +365 -65
  100. optimum/rbln/utils/runtime_utils.py +3 -3
  101. optimum/rbln/utils/submodule.py +10 -4
  102. {optimum_rbln-0.8.2rc0.dist-info → optimum_rbln-0.8.3.dist-info}/METADATA +1 -1
  103. {optimum_rbln-0.8.2rc0.dist-info → optimum_rbln-0.8.3.dist-info}/RECORD +105 -89
  104. {optimum_rbln-0.8.2rc0.dist-info → optimum_rbln-0.8.3.dist-info}/WHEEL +0 -0
  105. {optimum_rbln-0.8.2rc0.dist-info → optimum_rbln-0.8.3.dist-info}/licenses/LICENSE +0 -0
optimum/rbln/__init__.py CHANGED
@@ -46,6 +46,8 @@ _import_structure = {
46
46
  "RBLNAutoModelForSequenceClassification",
47
47
  "RBLNAutoModelForSpeechSeq2Seq",
48
48
  "RBLNAutoModelForVision2Seq",
49
+ "RBLNAutoModelForTextEncoding",
50
+ "RBLNAutoModelForZeroShotObjectDetection",
49
51
  "RBLNBartForConditionalGeneration",
50
52
  "RBLNBartForConditionalGenerationConfig",
51
53
  "RBLNBartModel",
@@ -80,6 +82,8 @@ _import_structure = {
80
82
  "RBLNDistilBertForQuestionAnsweringConfig",
81
83
  "RBLNDPTForDepthEstimation",
82
84
  "RBLNDPTForDepthEstimationConfig",
85
+ "RBLNDepthAnythingForDepthEstimationConfig",
86
+ "RBLNDepthAnythingForDepthEstimation",
83
87
  "RBLNExaoneForCausalLM",
84
88
  "RBLNExaoneForCausalLMConfig",
85
89
  "RBLNGemmaModel",
@@ -94,6 +98,12 @@ _import_structure = {
94
98
  "RBLNGPT2ModelConfig",
95
99
  "RBLNGPT2LMHeadModel",
96
100
  "RBLNGPT2LMHeadModelConfig",
101
+ "RBLNGroundingDinoDecoder",
102
+ "RBLNGroundingDinoDecoderConfig",
103
+ "RBLNGroundingDinoForObjectDetection",
104
+ "RBLNGroundingDinoForObjectDetectionConfig",
105
+ "RBLNGroundingDinoEncoder",
106
+ "RBLNGroundingDinoEncoderConfig",
97
107
  "RBLNIdefics3VisionTransformer",
98
108
  "RBLNIdefics3ForConditionalGeneration",
99
109
  "RBLNIdefics3ForConditionalGenerationConfig",
@@ -146,6 +156,8 @@ _import_structure = {
146
156
  "RBLNRobertaForSequenceClassificationConfig",
147
157
  "RBLNSiglipVisionModel",
148
158
  "RBLNSiglipVisionModelConfig",
159
+ "RBLNSwinBackbone",
160
+ "RBLNSwinBackboneConfig",
149
161
  "RBLNT5EncoderModel",
150
162
  "RBLNT5EncoderModelConfig",
151
163
  "RBLNT5ForConditionalGeneration",
@@ -168,6 +180,9 @@ _import_structure = {
168
180
  "RBLNAutoencoderKLConfig",
169
181
  "RBLNAutoencoderKLCosmos",
170
182
  "RBLNAutoencoderKLCosmosConfig",
183
+ "RBLNAutoPipelineForImage2Image",
184
+ "RBLNAutoPipelineForInpainting",
185
+ "RBLNAutoPipelineForText2Image",
171
186
  "RBLNControlNetModel",
172
187
  "RBLNControlNetModelConfig",
173
188
  "RBLNCosmosTextToWorldPipeline",
@@ -231,16 +246,15 @@ _import_structure = {
231
246
  }
232
247
 
233
248
  if TYPE_CHECKING:
234
- from .configuration_utils import (
235
- RBLNAutoConfig,
236
- RBLNCompileConfig,
237
- RBLNModelConfig,
238
- )
249
+ from .configuration_utils import RBLNAutoConfig, RBLNCompileConfig, RBLNModelConfig
239
250
  from .diffusers import (
240
251
  RBLNAutoencoderKL,
241
252
  RBLNAutoencoderKLConfig,
242
253
  RBLNAutoencoderKLCosmos,
243
254
  RBLNAutoencoderKLCosmosConfig,
255
+ RBLNAutoPipelineForImage2Image,
256
+ RBLNAutoPipelineForInpainting,
257
+ RBLNAutoPipelineForText2Image,
244
258
  RBLNControlNetModel,
245
259
  RBLNControlNetModelConfig,
246
260
  RBLNCosmosSafetyChecker,
@@ -301,10 +315,7 @@ if TYPE_CHECKING:
301
315
  RBLNVQModel,
302
316
  RBLNVQModelConfig,
303
317
  )
304
- from .modeling import (
305
- RBLNBaseModel,
306
- RBLNModel,
307
- )
318
+ from .modeling import RBLNBaseModel, RBLNModel
308
319
  from .transformers import (
309
320
  RBLNASTForAudioClassification,
310
321
  RBLNASTForAudioClassificationConfig,
@@ -320,7 +331,9 @@ if TYPE_CHECKING:
320
331
  RBLNAutoModelForSeq2SeqLM,
321
332
  RBLNAutoModelForSequenceClassification,
322
333
  RBLNAutoModelForSpeechSeq2Seq,
334
+ RBLNAutoModelForTextEncoding,
323
335
  RBLNAutoModelForVision2Seq,
336
+ RBLNAutoModelForZeroShotObjectDetection,
324
337
  RBLNBartForConditionalGeneration,
325
338
  RBLNBartForConditionalGenerationConfig,
326
339
  RBLNBartModel,
@@ -351,6 +364,8 @@ if TYPE_CHECKING:
351
364
  RBLNDecoderOnlyModelConfig,
352
365
  RBLNDecoderOnlyModelForCausalLM,
353
366
  RBLNDecoderOnlyModelForCausalLMConfig,
367
+ RBLNDepthAnythingForDepthEstimation,
368
+ RBLNDepthAnythingForDepthEstimationConfig,
354
369
  RBLNDistilBertForQuestionAnswering,
355
370
  RBLNDistilBertForQuestionAnsweringConfig,
356
371
  RBLNDPTForDepthEstimation,
@@ -369,6 +384,12 @@ if TYPE_CHECKING:
369
384
  RBLNGPT2LMHeadModelConfig,
370
385
  RBLNGPT2Model,
371
386
  RBLNGPT2ModelConfig,
387
+ RBLNGroundingDinoDecoder,
388
+ RBLNGroundingDinoDecoderConfig,
389
+ RBLNGroundingDinoEncoder,
390
+ RBLNGroundingDinoEncoderConfig,
391
+ RBLNGroundingDinoForObjectDetection,
392
+ RBLNGroundingDinoForObjectDetectionConfig,
372
393
  RBLNIdefics3ForConditionalGeneration,
373
394
  RBLNIdefics3ForConditionalGenerationConfig,
374
395
  RBLNIdefics3VisionTransformer,
@@ -421,6 +442,8 @@ if TYPE_CHECKING:
421
442
  RBLNRobertaForSequenceClassificationConfig,
422
443
  RBLNSiglipVisionModel,
423
444
  RBLNSiglipVisionModelConfig,
445
+ RBLNSwinBackbone,
446
+ RBLNSwinBackboneConfig,
424
447
  RBLNT5EncoderModel,
425
448
  RBLNT5EncoderModelConfig,
426
449
  RBLNT5ForConditionalGeneration,
@@ -1,7 +1,14 @@
1
1
  # file generated by setuptools-scm
2
2
  # don't change, don't track in version control
3
3
 
4
- __all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
4
+ __all__ = [
5
+ "__version__",
6
+ "__version_tuple__",
7
+ "version",
8
+ "version_tuple",
9
+ "__commit_id__",
10
+ "commit_id",
11
+ ]
5
12
 
6
13
  TYPE_CHECKING = False
7
14
  if TYPE_CHECKING:
@@ -9,13 +16,19 @@ if TYPE_CHECKING:
9
16
  from typing import Union
10
17
 
11
18
  VERSION_TUPLE = Tuple[Union[int, str], ...]
19
+ COMMIT_ID = Union[str, None]
12
20
  else:
13
21
  VERSION_TUPLE = object
22
+ COMMIT_ID = object
14
23
 
15
24
  version: str
16
25
  __version__: str
17
26
  __version_tuple__: VERSION_TUPLE
18
27
  version_tuple: VERSION_TUPLE
28
+ commit_id: COMMIT_ID
29
+ __commit_id__: COMMIT_ID
19
30
 
20
- __version__ = version = '0.8.2rc0'
21
- __version_tuple__ = version_tuple = (0, 8, 2, 'rc0')
31
+ __version__ = version = '0.8.3'
32
+ __version_tuple__ = version_tuple = (0, 8, 3)
33
+
34
+ __commit_id__ = commit_id = None
@@ -21,6 +21,7 @@ from typing import Any, Dict, List, Optional, Protocol, Tuple, Type, Union, runt
21
21
 
22
22
  import numpy as np
23
23
  import torch
24
+ from packaging.version import Version
24
25
 
25
26
  from .__version__ import __version__
26
27
  from .utils.depreacate_utils import warn_deprecated_npu
@@ -491,7 +492,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
491
492
  self,
492
493
  submodule_config_cls: Type["RBLNModelConfig"],
493
494
  submodule_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
494
- **kwargs: Dict[str, Any],
495
+ **kwargs: Any,
495
496
  ) -> "RBLNModelConfig":
496
497
  # Initialize a submodule config from a dict or a RBLNModelConfig.
497
498
  # kwargs is specified from the predecessor config.
@@ -566,7 +567,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
566
567
  timeout: Optional[int] = None,
567
568
  optimum_rbln_version: Optional[str] = None,
568
569
  _compile_cfgs: List[RBLNCompileConfig] = [],
569
- **kwargs: Dict[str, Any],
570
+ **kwargs: Any,
570
571
  ):
571
572
  """
572
573
  Initialize a RBLN model configuration with runtime options and compile configurations.
@@ -621,6 +622,21 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
621
622
  self.set_compile_cfgs([RBLNCompileConfig(**cfg) for cfg in self._compile_cfgs])
622
623
 
623
624
  if len(kwargs) > 0:
625
+ if optimum_rbln_version is not None: # loaded from file
626
+ if Version(__version__) < Version(optimum_rbln_version):
627
+ diff = "newer"
628
+ elif Version(__version__) > Version(optimum_rbln_version):
629
+ diff = "older"
630
+ else:
631
+ diff = None
632
+ if diff is not None:
633
+ raise ValueError(
634
+ f"Unexpected arguments: {kwargs.keys()}\n"
635
+ f"Maybe you are trying to load a model compiled with {diff} version of optimum-rbln. "
636
+ "It is recommended to use the same version to compile and load the model.\n"
637
+ f"Current version: {__version__}, Loaded version: {optimum_rbln_version}"
638
+ )
639
+
624
640
  raise ValueError(f"Unexpected arguments: {kwargs.keys()}")
625
641
 
626
642
  @property
@@ -717,7 +733,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
717
733
  json.dump(serializable_data, jsonf, indent=2)
718
734
 
719
735
  @classmethod
720
- def load(cls, path: str, **kwargs: Dict[str, Any]) -> "RBLNModelConfig":
736
+ def load(cls, path: str, **kwargs: Any) -> "RBLNModelConfig":
721
737
  """
722
738
  Load a RBLNModelConfig from a path.
723
739
 
@@ -750,7 +766,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
750
766
  def initialize_from_kwargs(
751
767
  cls: Type["RBLNModelConfig"],
752
768
  rbln_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
753
- **kwargs: Dict[str, Any],
769
+ **kwargs: Any,
754
770
  ) -> Tuple["RBLNModelConfig", Dict[str, Any]]:
755
771
  # Initialize RBLNModelConfig from kwargs.
756
772
  kwargs_keys = list(kwargs.keys())
@@ -59,6 +59,9 @@ _import_structure = {
59
59
  "RBLNVQModelConfig",
60
60
  ],
61
61
  "pipelines": [
62
+ "RBLNAutoPipelineForImage2Image",
63
+ "RBLNAutoPipelineForInpainting",
64
+ "RBLNAutoPipelineForText2Image",
62
65
  "RBLNCosmosTextToWorldPipeline",
63
66
  "RBLNCosmosVideoToWorldPipeline",
64
67
  "RBLNCosmosSafetyChecker",
@@ -135,6 +138,7 @@ if TYPE_CHECKING:
135
138
  from .modeling_diffusers import RBLNDiffusionMixin
136
139
  from .models import (
137
140
  RBLNAutoencoderKL,
141
+ RBLNAutoencoderKLCosmos,
138
142
  RBLNControlNetModel,
139
143
  RBLNCosmosTransformer3DModel,
140
144
  RBLNPriorTransformer,
@@ -143,6 +147,9 @@ if TYPE_CHECKING:
143
147
  RBLNVQModel,
144
148
  )
145
149
  from .pipelines import (
150
+ RBLNAutoPipelineForImage2Image,
151
+ RBLNAutoPipelineForInpainting,
152
+ RBLNAutoPipelineForText2Image,
146
153
  RBLNCosmosSafetyChecker,
147
154
  RBLNCosmosTextToWorldPipeline,
148
155
  RBLNCosmosVideoToWorldPipeline,
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional, Tuple
15
+ from typing import Any, Optional, Tuple
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
 
@@ -33,7 +33,7 @@ class RBLNAutoencoderKLConfig(RBLNModelConfig):
33
33
  vae_scale_factor: Optional[float] = None, # TODO: rename to scaling_factor
34
34
  in_channels: Optional[int] = None,
35
35
  latent_channels: Optional[int] = None,
36
- **kwargs: Dict[str, Any],
36
+ **kwargs: Any,
37
37
  ):
38
38
  """
39
39
  Args:
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional
15
+ from typing import Any, Optional
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
  from ....utils.logging import get_logger
@@ -35,7 +35,7 @@ class RBLNAutoencoderKLCosmosConfig(RBLNModelConfig):
35
35
  vae_scale_factor_temporal: Optional[int] = None,
36
36
  vae_scale_factor_spatial: Optional[int] = None,
37
37
  use_slicing: Optional[bool] = None,
38
- **kwargs: Dict[str, Any],
38
+ **kwargs: Any,
39
39
  ):
40
40
  """
41
41
  Args:
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional, Tuple
15
+ from typing import Any, Optional, Tuple
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
 
@@ -29,7 +29,7 @@ class RBLNControlNetModelConfig(RBLNModelConfig):
29
29
  unet_sample_size: Optional[Tuple[int, int]] = None,
30
30
  vae_sample_size: Optional[Tuple[int, int]] = None,
31
31
  text_model_hidden_size: Optional[int] = None,
32
- **kwargs: Dict[str, Any],
32
+ **kwargs: Any,
33
33
  ):
34
34
  """
35
35
  Args:
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional
15
+ from typing import Any, Optional
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
 
@@ -32,7 +32,7 @@ class RBLNPriorTransformerConfig(RBLNModelConfig):
32
32
  batch_size: Optional[int] = None,
33
33
  embedding_dim: Optional[int] = None,
34
34
  num_embeddings: Optional[int] = None,
35
- **kwargs: Dict[str, Any],
35
+ **kwargs: Any,
36
36
  ):
37
37
  """
38
38
  Args:
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional
15
+ from typing import Any, Optional
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
 
@@ -33,7 +33,7 @@ class RBLNCosmosTransformer3DModelConfig(RBLNModelConfig):
33
33
  num_latent_frames: Optional[int] = None,
34
34
  latent_height: Optional[int] = None,
35
35
  latent_width: Optional[int] = None,
36
- **kwargs: Dict[str, Any],
36
+ **kwargs: Any,
37
37
  ):
38
38
  """
39
39
  Args:
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional, Tuple, Union
15
+ from typing import Any, Optional, Tuple, Union
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
 
@@ -27,7 +27,7 @@ class RBLNSD3Transformer2DModelConfig(RBLNModelConfig):
27
27
  batch_size: Optional[int] = None,
28
28
  sample_size: Optional[Union[int, Tuple[int, int]]] = None,
29
29
  prompt_embed_length: Optional[int] = None,
30
- **kwargs: Dict[str, Any],
30
+ **kwargs: Any,
31
31
  ):
32
32
  """
33
33
  Args:
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional, Tuple
15
+ from typing import Any, Optional, Tuple
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
 
@@ -38,7 +38,7 @@ class RBLNUNet2DConditionModelConfig(RBLNModelConfig):
38
38
  in_features: Optional[int] = None,
39
39
  text_model_hidden_size: Optional[int] = None,
40
40
  image_model_hidden_size: Optional[int] = None,
41
- **kwargs: Dict[str, Any],
41
+ **kwargs: Any,
42
42
  ):
43
43
  """
44
44
  Args:
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional, Tuple
15
+ from typing import Any, Optional, Tuple
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
 
@@ -33,7 +33,7 @@ class RBLNVQModelConfig(RBLNModelConfig):
33
33
  vqmodel_scale_factor: Optional[float] = None, # TODO: rename to scaling_factor
34
34
  in_channels: Optional[int] = None,
35
35
  latent_channels: Optional[int] = None,
36
- **kwargs: Dict[str, Any],
36
+ **kwargs: Any,
37
37
  ):
38
38
  """
39
39
  Args:
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional, Tuple
15
+ from typing import Any, Optional, Tuple
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
  from ....transformers import RBLNCLIPTextModelConfig, RBLNCLIPTextModelWithProjectionConfig
@@ -38,7 +38,7 @@ class RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
38
38
  sample_size: Optional[Tuple[int, int]] = None,
39
39
  image_size: Optional[Tuple[int, int]] = None,
40
40
  guidance_scale: Optional[float] = None,
41
- **kwargs: Dict[str, Any],
41
+ **kwargs: Any,
42
42
  ):
43
43
  """
44
44
  Args:
@@ -178,7 +178,7 @@ class RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
178
178
  sample_size: Optional[Tuple[int, int]] = None,
179
179
  image_size: Optional[Tuple[int, int]] = None,
180
180
  guidance_scale: Optional[float] = None,
181
- **kwargs: Dict[str, Any],
181
+ **kwargs: Any,
182
182
  ):
183
183
  """
184
184
  Args:
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional
15
+ from typing import Any, Optional
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
  from ....transformers import RBLNT5EncoderModelConfig
@@ -41,7 +41,7 @@ class RBLNCosmosPipelineBaseConfig(RBLNModelConfig):
41
41
  num_frames: Optional[int] = None,
42
42
  fps: Optional[int] = None,
43
43
  max_seq_len: Optional[int] = None,
44
- **kwargs: Dict[str, Any],
44
+ **kwargs: Any,
45
45
  ):
46
46
  """
47
47
  Args:
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional, Tuple
15
+ from typing import Any, Optional, Tuple
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
  from ....transformers import RBLNCLIPTextModelWithProjectionConfig, RBLNCLIPVisionModelWithProjectionConfig
@@ -37,7 +37,7 @@ class RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
37
37
  img_width: Optional[int] = None,
38
38
  height: Optional[int] = None,
39
39
  width: Optional[int] = None,
40
- **kwargs: Dict[str, Any],
40
+ **kwargs: Any,
41
41
  ):
42
42
  """
43
43
  Args:
@@ -148,7 +148,7 @@ class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
148
148
  *,
149
149
  batch_size: Optional[int] = None,
150
150
  guidance_scale: Optional[float] = None,
151
- **kwargs: Dict[str, Any],
151
+ **kwargs: Any,
152
152
  ):
153
153
  """
154
154
  Initialize a configuration for Kandinsky 2.2 prior pipeline optimized for RBLN NPU.
@@ -226,7 +226,7 @@ class RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
226
226
  prior_text_encoder: Optional[RBLNCLIPTextModelWithProjectionConfig] = None,
227
227
  unet: Optional[RBLNUNet2DConditionModelConfig] = None,
228
228
  movq: Optional[RBLNVQModelConfig] = None,
229
- **kwargs: Dict[str, Any],
229
+ **kwargs: Any,
230
230
  ):
231
231
  """
232
232
  Initialize a configuration for combined Kandinsky 2.2 pipelines optimized for RBLN NPU.
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional, Tuple
15
+ from typing import Any, Optional, Tuple
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
  from ....transformers import RBLNCLIPTextModelConfig
@@ -37,7 +37,7 @@ class RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
37
37
  sample_size: Optional[Tuple[int, int]] = None,
38
38
  image_size: Optional[Tuple[int, int]] = None,
39
39
  guidance_scale: Optional[float] = None,
40
- **kwargs: Dict[str, Any],
40
+ **kwargs: Any,
41
41
  ):
42
42
  """
43
43
  Args:
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional, Tuple
15
+ from typing import Any, Optional, Tuple
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
  from ....transformers import RBLNCLIPTextModelWithProjectionConfig, RBLNT5EncoderModelConfig
@@ -40,7 +40,7 @@ class RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
40
40
  height: Optional[int] = None,
41
41
  width: Optional[int] = None,
42
42
  guidance_scale: Optional[float] = None,
43
- **kwargs: Dict[str, Any],
43
+ **kwargs: Any,
44
44
  ):
45
45
  """
46
46
  Args:
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Any, Dict, Optional, Tuple
15
+ from typing import Any, Optional, Tuple
16
16
 
17
17
  from ....configuration_utils import RBLNModelConfig
18
18
  from ....transformers import RBLNCLIPTextModelConfig, RBLNCLIPTextModelWithProjectionConfig
@@ -38,7 +38,7 @@ class RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
38
38
  sample_size: Optional[Tuple[int, int]] = None,
39
39
  image_size: Optional[Tuple[int, int]] = None,
40
40
  guidance_scale: Optional[float] = None,
41
- **kwargs: Dict[str, Any],
41
+ **kwargs: Any,
42
42
  ):
43
43
  """
44
44
  Args:
@@ -136,7 +136,7 @@ class RBLNDiffusionMixin:
136
136
  lora_ids: Optional[Union[str, List[str]]] = None,
137
137
  lora_weights_names: Optional[Union[str, List[str]]] = None,
138
138
  lora_scales: Optional[Union[float, List[float]]] = None,
139
- **kwargs: Dict[str, Any],
139
+ **kwargs: Any,
140
140
  ) -> "RBLNDiffusionMixin":
141
141
  """
142
142
  Load a pretrained diffusion pipeline from a model checkpoint, with optional compilation for RBLN NPUs.
@@ -35,20 +35,10 @@ _import_structure = {
35
35
  }
36
36
 
37
37
  if TYPE_CHECKING:
38
- from .autoencoders import (
39
- RBLNAutoencoderKL,
40
- RBLNAutoencoderKLCosmos,
41
- RBLNVQModel,
42
- )
38
+ from .autoencoders import RBLNAutoencoderKL, RBLNAutoencoderKLCosmos, RBLNVQModel
43
39
  from .controlnet import RBLNControlNetModel
44
- from .transformers import (
45
- RBLNCosmosTransformer3DModel,
46
- RBLNPriorTransformer,
47
- RBLNSD3Transformer2DModel,
48
- )
49
- from .unets import (
50
- RBLNUNet2DConditionModel,
51
- )
40
+ from .transformers import RBLNCosmosTransformer3DModel, RBLNPriorTransformer, RBLNSD3Transformer2DModel
41
+ from .unets import RBLNUNet2DConditionModel
52
42
  else:
53
43
  import sys
54
44
 
@@ -18,6 +18,11 @@ from transformers.utils import _LazyModule
18
18
 
19
19
 
20
20
  _import_structure = {
21
+ "auto_pipeline": [
22
+ "RBLNAutoPipelineForImage2Image",
23
+ "RBLNAutoPipelineForInpainting",
24
+ "RBLNAutoPipelineForText2Image",
25
+ ],
21
26
  "controlnet": [
22
27
  "RBLNMultiControlNetModel",
23
28
  "RBLNStableDiffusionControlNetImg2ImgPipeline",
@@ -56,6 +61,11 @@ _import_structure = {
56
61
  ],
57
62
  }
58
63
  if TYPE_CHECKING:
64
+ from .auto_pipeline import (
65
+ RBLNAutoPipelineForImage2Image,
66
+ RBLNAutoPipelineForInpainting,
67
+ RBLNAutoPipelineForText2Image,
68
+ )
59
69
  from .controlnet import (
60
70
  RBLNMultiControlNetModel,
61
71
  RBLNStableDiffusionControlNetImg2ImgPipeline,
@@ -63,11 +73,7 @@ if TYPE_CHECKING:
63
73
  RBLNStableDiffusionXLControlNetImg2ImgPipeline,
64
74
  RBLNStableDiffusionXLControlNetPipeline,
65
75
  )
66
- from .cosmos import (
67
- RBLNCosmosSafetyChecker,
68
- RBLNCosmosTextToWorldPipeline,
69
- RBLNCosmosVideoToWorldPipeline,
70
- )
76
+ from .cosmos import RBLNCosmosSafetyChecker, RBLNCosmosTextToWorldPipeline, RBLNCosmosVideoToWorldPipeline
71
77
  from .kandinsky2_2 import (
72
78
  RBLNKandinskyV22CombinedPipeline,
73
79
  RBLNKandinskyV22Img2ImgCombinedPipeline,