optimum-rbln 0.1.13__py3-none-any.whl → 0.2.0__py3-none-any.whl

optimum/rbln/__init__.py

@@ -30,27 +30,15 @@ from .utils import check_version_compats
 
 
 _import_structure = {
-    "modeling_alias": [
-        "RBLNASTForAudioClassification",
-        "RBLNBertForQuestionAnswering",
-        "RBLNDistilBertForQuestionAnswering",
-        "RBLNResNetForImageClassification",
-        "RBLNXLMRobertaForSequenceClassification",
-        "RBLNRobertaForSequenceClassification",
-        "RBLNRobertaForMaskedLM",
-        "RBLNViTForImageClassification",
-    ],
-    "modeling_base": [
+    "modeling": [
         "RBLNBaseModel",
         "RBLNModel",
-        "RBLNModelForQuestionAnswering",
-        "RBLNModelForAudioClassification",
-        "RBLNModelForImageClassification",
-        "RBLNModelForSequenceClassification",
-        "RBLNModelForMaskedLM",
+    ],
+    "modeling_config": [
+        "RBLNCompileConfig",
+        "RBLNConfig",
     ],
     "transformers": [
-        "BatchTextIteratorStreamer",
         "RBLNAutoModel",
         "RBLNAutoModelForAudioClassification",
         "RBLNAutoModelForCausalLM",
@@ -84,6 +72,14 @@ _import_structure = {
         "RBLNMistralForCausalLM",
         "RBLNWhisperForConditionalGeneration",
         "RBLNXLMRobertaModel",
+        "RBLNASTForAudioClassification",
+        "RBLNBertForQuestionAnswering",
+        "RBLNDistilBertForQuestionAnswering",
+        "RBLNResNetForImageClassification",
+        "RBLNXLMRobertaForSequenceClassification",
+        "RBLNRobertaForSequenceClassification",
+        "RBLNRobertaForMaskedLM",
+        "RBLNViTForImageClassification",
     ],
     "diffusers": [
         "RBLNStableDiffusionPipeline",
@@ -92,55 +88,54 @@ _import_structure = {
         "RBLNUNet2DConditionModel",
         "RBLNControlNetModel",
         "RBLNStableDiffusionImg2ImgPipeline",
+        "RBLNStableDiffusionInpaintPipeline",
         "RBLNStableDiffusionControlNetImg2ImgPipeline",
         "RBLNMultiControlNetModel",
         "RBLNStableDiffusionXLImg2ImgPipeline",
+        "RBLNStableDiffusionXLInpaintPipeline",
         "RBLNStableDiffusionControlNetPipeline",
         "RBLNStableDiffusionXLControlNetPipeline",
         "RBLNStableDiffusionXLControlNetImg2ImgPipeline",
+        "RBLNSD3Transformer2DModel",
+        "RBLNStableDiffusion3Img2ImgPipeline",
+        "RBLNStableDiffusion3InpaintPipeline",
+        "RBLNStableDiffusion3Pipeline",
+        "RBLNDiffusionMixin",
     ],
-    "modeling_config": ["RBLNCompileConfig", "RBLNConfig"],
-    "modeling_diffusers": ["RBLNDiffusionMixin"],
 }
 
 if TYPE_CHECKING:
     from .diffusers import (
         RBLNAutoencoderKL,
         RBLNControlNetModel,
+        RBLNDiffusionMixin,
         RBLNMultiControlNetModel,
+        RBLNSD3Transformer2DModel,
+        RBLNStableDiffusion3Img2ImgPipeline,
+        RBLNStableDiffusion3InpaintPipeline,
+        RBLNStableDiffusion3Pipeline,
         RBLNStableDiffusionControlNetImg2ImgPipeline,
         RBLNStableDiffusionControlNetPipeline,
         RBLNStableDiffusionImg2ImgPipeline,
+        RBLNStableDiffusionInpaintPipeline,
         RBLNStableDiffusionPipeline,
         RBLNStableDiffusionXLControlNetImg2ImgPipeline,
         RBLNStableDiffusionXLControlNetPipeline,
         RBLNStableDiffusionXLImg2ImgPipeline,
+        RBLNStableDiffusionXLInpaintPipeline,
         RBLNStableDiffusionXLPipeline,
         RBLNUNet2DConditionModel,
     )
-    from .modeling_alias import (
-        RBLNASTForAudioClassification,
-        RBLNBertForQuestionAnswering,
-        RBLNResNetForImageClassification,
-        RBLNRobertaForMaskedLM,
-        RBLNRobertaForSequenceClassification,
-        RBLNT5ForConditionalGeneration,
-        RBLNViTForImageClassification,
-        RBLNXLMRobertaForSequenceClassification,
-    )
-    from .modeling_base import (
+    from .modeling import (
         RBLNBaseModel,
         RBLNModel,
-        RBLNModelForAudioClassification,
-        RBLNModelForImageClassification,
-        RBLNModelForMaskedLM,
-        RBLNModelForQuestionAnswering,
-        RBLNModelForSequenceClassification,
     )
-    from .modeling_config import RBLNCompileConfig, RBLNConfig
-    from .modeling_diffusers import RBLNDiffusionMixin
+    from .modeling_config import (
+        RBLNCompileConfig,
+        RBLNConfig,
+    )
     from .transformers import (
-        BatchTextIteratorStreamer,
+        RBLNASTForAudioClassification,
         RBLNAutoModel,
         RBLNAutoModelForAudioClassification,
         RBLNAutoModelForCausalLM,
@@ -155,10 +150,12 @@ if TYPE_CHECKING:
         RBLNAutoModelForVision2Seq,
         RBLNBartForConditionalGeneration,
         RBLNBartModel,
+        RBLNBertForQuestionAnswering,
         RBLNBertModel,
         RBLNCLIPTextModel,
         RBLNCLIPTextModelWithProjection,
         RBLNCLIPVisionModel,
+        RBLNDistilBertForQuestionAnswering,
         RBLNDPTForDepthEstimation,
         RBLNExaoneForCausalLM,
         RBLNGemmaForCausalLM,
@@ -169,12 +166,18 @@ if TYPE_CHECKING:
         RBLNMistralForCausalLM,
         RBLNPhiForCausalLM,
         RBLNQwen2ForCausalLM,
+        RBLNResNetForImageClassification,
+        RBLNRobertaForMaskedLM,
+        RBLNRobertaForSequenceClassification,
         RBLNT5EncoderModel,
         RBLNT5ForConditionalGeneration,
+        RBLNViTForImageClassification,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,
+        RBLNXLMRobertaForSequenceClassification,
         RBLNXLMRobertaModel,
     )
+
 else:
     import sys
 

optimum/rbln/__version__.py

@@ -1 +1,16 @@
-__version__ = '0.1.13'
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 0)

optimum/rbln/diffusers/__init__.py

@@ -36,27 +36,51 @@ _import_structure = {
         "RBLNStableDiffusionPipeline",
         "RBLNStableDiffusionXLPipeline",
         "RBLNStableDiffusionImg2ImgPipeline",
+        "RBLNStableDiffusionInpaintPipeline",
         "RBLNStableDiffusionControlNetImg2ImgPipeline",
         "RBLNMultiControlNetModel",
         "RBLNStableDiffusionXLImg2ImgPipeline",
+        "RBLNStableDiffusionXLInpaintPipeline",
         "RBLNStableDiffusionControlNetPipeline",
         "RBLNStableDiffusionXLControlNetPipeline",
         "RBLNStableDiffusionXLControlNetImg2ImgPipeline",
+        "RBLNStableDiffusion3Pipeline",
+        "RBLNStableDiffusion3Img2ImgPipeline",
+        "RBLNStableDiffusion3InpaintPipeline",
+    ],
+    "models": [
+        "RBLNAutoencoderKL",
+        "RBLNUNet2DConditionModel",
+        "RBLNControlNetModel",
+        "RBLNSD3Transformer2DModel",
+    ],
+    "modeling_diffusers": [
+        "RBLNDiffusionMixin",
     ],
-    "models": ["RBLNAutoencoderKL", "RBLNUNet2DConditionModel", "RBLNControlNetModel"],
 }
 
 if TYPE_CHECKING:
-    from .models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
+    from .modeling_diffusers import RBLNDiffusionMixin
+    from .models import (
+        RBLNAutoencoderKL,
+        RBLNControlNetModel,
+        RBLNSD3Transformer2DModel,
+        RBLNUNet2DConditionModel,
+    )
     from .pipelines import (
         RBLNMultiControlNetModel,
+        RBLNStableDiffusion3Img2ImgPipeline,
+        RBLNStableDiffusion3InpaintPipeline,
+        RBLNStableDiffusion3Pipeline,
         RBLNStableDiffusionControlNetImg2ImgPipeline,
         RBLNStableDiffusionControlNetPipeline,
         RBLNStableDiffusionImg2ImgPipeline,
+        RBLNStableDiffusionInpaintPipeline,
         RBLNStableDiffusionPipeline,
         RBLNStableDiffusionXLControlNetImg2ImgPipeline,
         RBLNStableDiffusionXLControlNetPipeline,
         RBLNStableDiffusionXLImg2ImgPipeline,
+        RBLNStableDiffusionXLInpaintPipeline,
         RBLNStableDiffusionXLPipeline,
     )
 else:

optimum/rbln/{modeling_diffusers.py → diffusers/modeling_diffusers.py}

@@ -20,16 +20,21 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
+
+import copy
 import importlib
 from os import PathLike
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 import torch
 
-from .modeling_base import RBLNModel
-from .modeling_config import ContextRblnConfig, use_rbln_config
-from .utils.decorator_utils import remove_compile_time_kwargs
+from ..modeling import RBLNModel
+from ..modeling_config import RUNTIME_KEYWORDS, ContextRblnConfig, use_rbln_config
+from ..utils.decorator_utils import remove_compile_time_kwargs
+from ..utils.logging import get_logger
+
 
+logger = get_logger(__name__)
 
 if TYPE_CHECKING:
     from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
@@ -74,127 +79,40 @@ class RBLNDiffusionMixin:
 
     @classmethod
     @property
-    def use_encode(cls):
+    def img2img_pipeline(cls):
         return "Img2Img" in cls.__name__
 
     @classmethod
-    def _get_unet_batch_size(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> int:
-        # Calculates the batch size based on guidance scale
-        batch_size = rbln_config.get("batch_size", 1)
-        do_guidance = rbln_config.get("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
-        return batch_size * 2 if do_guidance else batch_size
-
-    @classmethod
-    def _get_vae_sample_size(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> Union[int, Tuple[int, int]]:
-        image_size = (rbln_config.get("img_height"), rbln_config.get("img_width"))
-        if (image_size[0] is None) != (image_size[1] is None):
-            raise ValueError("Both image height and image width must be given or not given")
-        elif image_size[0] is None and image_size[1] is None:
-            if cls.use_encode:
-                sample_size = model.vae.config.sample_size
-            else:
-                # In case of text2img, sample size of vae decoder is determined by unet.
-                unet_sample_size = model.unet.config.sample_size
-                if isinstance(unet_sample_size, int):
-                    sample_size = unet_sample_size * model.vae_scale_factor
-                else:
-                    sample_size = (
-                        unet_sample_size[0] * model.vae_scale_factor,
-                        unet_sample_size[1] * model.vae_scale_factor,
-                    )
-
-        else:
-            sample_size = (image_size[0], image_size[1])
-        return sample_size
-
-    @classmethod
-    def _get_unet_sample_size(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> Union[int, Tuple[int, int]]:
-        image_size = (rbln_config.get("img_height"), rbln_config.get("img_width"))
-        if (image_size[0] is None) != (image_size[1] is None):
-            raise ValueError("Both image height and image width must be given or not given")
-        elif image_size[0] is None and image_size[1] is None:
-            if cls.use_encode:
-                # In case of img2img, sample size of unet is determined by vae encoder.
-                vae_sample_size = model.vae.config.sample_size
-                if isinstance(vae_sample_size, int):
-                    sample_size = vae_sample_size // model.vae_scale_factor
-                else:
-                    sample_size = (
-                        vae_sample_size[0] // model.vae_scale_factor,
-                        vae_sample_size[1] // model.vae_scale_factor,
-                    )
-            else:
-                sample_size = model.unet.config.sample_size
-        else:
-            sample_size = (image_size[0] // model.vae_scale_factor, image_size[1] // model.vae_scale_factor)
-        return sample_size
-
-    @classmethod
-    def _get_default_config(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
-        # default configurations for each submodules
-        return {"img2img_pipeline": cls.use_encode}
+    @property
+    def inpaint_pipeline(cls):
+        return "Inpaint" in cls.__name__
 
     @classmethod
-    def get_default_rbln_config_text_encoder(
-        cls, model: torch.nn.Module, rbln_config: Dict[str, Any]
+    def get_submodule_rbln_config(
+        cls, model: torch.nn.Module, submodule_name: str, rbln_config: Dict[str, Any]
     ) -> Dict[str, Any]:
-        batch_size = rbln_config.get("batch_size", 1)
-        return {"batch_size": batch_size}
+        submodule = getattr(model, submodule_name)
+        submodule_class_name = submodule.__class__.__name__
 
-    @classmethod
-    def get_default_rbln_config_text_encoder_2(
-        cls, model: torch.nn.Module, rbln_config: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        batch_size = rbln_config.get("batch_size", 1)
-        return {"batch_size": batch_size}
+        if submodule_class_name == "MultiControlNetModel":
+            submodule_class_name = "ControlNetModel"
 
-    @classmethod
-    def get_default_rbln_config_unet(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
-        # configuration for unet
-        unet_batch_size = cls._get_unet_batch_size(model, rbln_config)
-        text_model_hidden_size = model.text_encoder_2.config.hidden_size if hasattr(model, "text_encoder_2") else None
-        return {
-            **cls._get_default_config(model, rbln_config),
-            "max_seq_len": model.text_encoder.config.max_position_embeddings,
-            "text_model_hidden_size": text_model_hidden_size,
-            "batch_size": unet_batch_size,
-            "sample_size": cls._get_unet_sample_size(model, rbln_config),
-            "is_controlnet": "controlnet" in model.config.keys(),
-        }
+        submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), f"RBLN{submodule_class_name}")
 
-    @classmethod
-    def get_default_rbln_config_vae(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
-        # configuration for vae
-        batch_size = rbln_config.get("batch_size", 1)
-        return {
-            **cls._get_default_config(model, rbln_config),
-            "sample_size": cls._get_vae_sample_size(model, rbln_config),
-            "batch_size": batch_size,
-        }
+        submodule_config = rbln_config.get(submodule_name, {})
+        submodule_config = copy.deepcopy(submodule_config)
 
-    @classmethod
-    def get_default_rbln_config_controlnet(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
-        # configuration for controlnet
-        unet_batch_size = cls._get_unet_batch_size(model, rbln_config)
-        text_model_hidden_size = model.text_encoder_2.config.hidden_size if hasattr(model, "text_encoder_2") else None
-        return {
-            **cls._get_default_config(model, rbln_config),
-            "max_seq_len": model.text_encoder.config.max_position_embeddings,
-            "vae_sample_size": cls._get_vae_sample_size(model, rbln_config),
-            "unet_sample_size": cls._get_unet_sample_size(model, rbln_config),
-            "batch_size": unet_batch_size,
-            "text_model_hidden_size": text_model_hidden_size,
-        }
+        pipe_global_config = {k: v for k, v in rbln_config.items() if k not in cls._submodules}
 
-    @classmethod
-    def get_default_rbln_config(
-        cls, model: torch.nn.Module, submodule_name: str, rbln_config: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        # Returns the default configuration based on submodule name
-        config_method = f"get_default_rbln_config_{submodule_name}"
-        if hasattr(cls, config_method):
-            return getattr(cls, config_method)(model, rbln_config)
-        raise ValueError(f"Unknown submodule: {submodule_name}")
+        submodule_config.update({k: v for k, v in pipe_global_config.items() if k not in submodule_config})
+        submodule_config.update(
+            {
+                "img2img_pipeline": cls.img2img_pipeline,
+                "inpaint_pipeline": cls.inpaint_pipeline,
+            }
+        )
+        submodule_config = submodule_cls.update_rbln_config_using_pipe(model, submodule_config)
+        return submodule_config
 
     @staticmethod
     def _maybe_apply_and_fuse_lora(
@@ -256,17 +174,46 @@ class RBLNDiffusionMixin:
 
         else:
             # raise error if any of submodules are torch module.
-            for name in cls._submodules:
-                if isinstance(kwargs.get(name), torch.nn.Module):
+            model_index_config = None
+            for submodule_name in cls._submodules:
+                if isinstance(kwargs.get(submodule_name), torch.nn.Module):
                     raise AssertionError(
-                        f"{name} is not compiled torch module. If you want to compile, set `export=True`."
+                        f"{submodule_name} is not compiled torch module. If you want to compile, set `export=True`."
                     )
 
+                submodule_config = rbln_config.get(submodule_name, {})
+
+                for key, value in rbln_config.items():
+                    if key in RUNTIME_KEYWORDS and key not in submodule_config:
+                        submodule_config[key] = value
+
+                if not any(kwd in submodule_config for kwd in RUNTIME_KEYWORDS):
+                    continue
+
+                if model_index_config is None:
+                    model_index_config = cls.load_config(pretrained_model_name_or_path=model_id)
+
+                module_name, class_name = model_index_config[submodule_name]
+                if module_name != "optimum.rbln":
+                    raise ValueError(
+                        f"Invalid module_name '{module_name}' found in model_index.json for "
+                        f"submodule '{submodule_name}'. "
+                        "Expected 'optimum.rbln'. Please check the model_index.json configuration."
+                    )
+
+                submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), class_name)
+
+                submodule = submodule_cls.from_pretrained(
+                    model_id, export=False, subfolder=submodule_name, rbln_config=submodule_config
+                )
+                kwargs[submodule_name] = submodule
+
         with ContextRblnConfig(
             device=rbln_config.get("device"),
             device_map=rbln_config.get("device_map"),
             create_runtimes=rbln_config.get("create_runtimes"),
             optimize_host_mem=rbln_config.get("optimize_host_memory"),
+            activate_profiler=rbln_config.get("activate_profiler"),
         ):
             model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
 
@@ -291,16 +238,11 @@ class RBLNDiffusionMixin:
         model_save_dir: Optional[PathLike],
         rbln_config: Dict[str, Any],
     ) -> Dict[str, RBLNModel]:
-        # Compile submodules based on rbln_config
         compiled_submodules = {}
 
-        # FIXME : Currently, optimum-rbln for transformer does not use base rbln config.
-        base_rbln_config = {k: v for k, v in rbln_config.items() if k not in cls._submodules}
         for submodule_name in cls._submodules:
             submodule = passed_submodules.get(submodule_name) or getattr(model, submodule_name, None)
-            submodule_rbln_config = cls.get_default_rbln_config(model, submodule_name, rbln_config)
-            submodule_rbln_config.update(base_rbln_config)
-            submodule_rbln_config.update(rbln_config.get(submodule_name, {}))
+            submodule_rbln_config = cls.get_submodule_rbln_config(model, submodule_name, rbln_config)
 
             if submodule is None:
                 raise ValueError(f"submodule ({submodule_name}) cannot be accessed since it is not provided.")
@@ -337,8 +279,8 @@ class RBLNDiffusionMixin:
         controlnet_rbln_config: Dict[str, Any],
     ):
         # Compile multiple ControlNet models for a MultiControlNet setup
-        from .diffusers.models.controlnet import RBLNControlNetModel
-        from .diffusers.pipelines.controlnet import RBLNMultiControlNetModel
+        from .models.controlnet import RBLNControlNetModel
+        from .pipelines.controlnet import RBLNMultiControlNetModel
 
         compiled_controlnets = [
             RBLNControlNetModel.from_model(
@@ -349,7 +291,7 @@ class RBLNDiffusionMixin:
             )
             for i, controlnet in enumerate(controlnets.nets)
         ]
-        return RBLNMultiControlNetModel(compiled_controlnets, config=controlnets.nets[0].config)
+        return RBLNMultiControlNetModel(compiled_controlnets)
 
     @classmethod
     def _construct_pipe(cls, model, submodules, model_save_dir, rbln_config):
@@ -395,6 +337,35 @@ class RBLNDiffusionMixin:
 
         return model
 
+    def get_compiled_image_size(self):
+        if hasattr(self, "vae"):
+            compiled_image_size = self.vae.image_size
+        else:
+            compiled_image_size = None
+        return compiled_image_size
+
+    def handle_additional_kwargs(self, **kwargs):
+        """
+        Function to handle additional compile-time parameters during inference.
+
+        If the additional variable is determined by another module, this method should be overrided.
+
+        Example:
+            ```python
+            if hasattr(self, "movq"):
+                compiled_image_size = self.movq.image_size
+                kwargs["height"] = compiled_image_size[0]
+                kwargs["width"] = compiled_image_size[1]
+
+            compiled_num_frames = self.unet.rbln_config.model_cfg.get("num_frames", None)
+            if compiled_num_frames is not None:
+                kwargs["num_frames"] = self.unet.rbln_config.model_cfg.get("num_frames")
+            return kwargs
+            ```
+        """
+        return kwargs
+
     @remove_compile_time_kwargs
     def __call__(self, *args, **kwargs):
+        kwargs = self.handle_additional_kwargs(**kwargs)
         return super().__call__(*args, **kwargs)

optimum/rbln/diffusers/models/__init__.py

@@ -21,6 +21,39 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-from .autoencoder_kl import RBLNAutoencoderKL
-from .controlnet import RBLNControlNetModel
-from .unet_2d_condition import RBLNUNet2DConditionModel
+from typing import TYPE_CHECKING
+
+from transformers.utils import _LazyModule
+
+
+_import_structure = {
+    "autoencoders": [
+        "RBLNAutoencoderKL",
+    ],
+    "unets": [
+        "RBLNUNet2DConditionModel",
+    ],
+    "controlnet": ["RBLNControlNetModel"],
+    "transformers": ["RBLNSD3Transformer2DModel"],
+}
+
+if TYPE_CHECKING:
+    from .autoencoders import (
+        RBLNAutoencoderKL,
+    )
+    from .controlnet import RBLNControlNetModel
+    from .transformers import (
+        RBLNSD3Transformer2DModel,
+    )
+    from .unets import (
+        RBLNUNet2DConditionModel,
+    )
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )

optimum/rbln/{transformers/generation → diffusers/models/autoencoders}/__init__.py

@@ -21,5 +21,4 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-from .streamers import BatchTextIteratorStreamer
-from .utils import RBLNGenerationMixin
+from .autoencoder_kl import RBLNAutoencoderKL