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

optimum/rbln/diffusers/pipelines/stable_diffusion/__init__.py

@@ -21,6 +21,5 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-from .pipeline_controlnet_img2img import RBLNStableDiffusionControlNetImg2ImgPipeline
 from .pipeline_stable_diffusion import RBLNStableDiffusionPipeline
 from .pipeline_stable_diffusion_img2img import RBLNStableDiffusionImg2ImgPipeline

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

@@ -22,7 +22,6 @@
 # from Rebellions Inc.
 """RBLNStableDiffusionPipeline class for inference of diffusion models on rbln devices."""
 
-
 from diffusers import StableDiffusionPipeline
 
 from ....modeling_base import RBLNBaseModel
@@ -50,17 +49,22 @@ class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
                 - A path to a *directory* containing a model saved using [`~OptimizedModel.save_pretrained`],
         """
         export = kwargs.pop("export", None)
+        model_save_dir = kwargs.pop("model_save_dir", None)
+        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
         model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
+
         if export is None or export is False:
             return model
 
-        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
+        do_classifier_free_guidance = (
+            rbln_config_kwargs.pop("rbln_guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
+        )
 
-        # compile model, create runtime
         vae = RBLNAutoencoderKL.from_pretrained(
             model_id=model_id,
             subfolder="vae",
             export=True,
+            model_save_dir=model_save_dir,
             rbln_unet_sample_size=model.unet.config.sample_size,
             rbln_use_encode=False,
             **rbln_config_kwargs,
@@ -70,16 +74,19 @@ class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
             model_id=model_id,
             subfolder="text_encoder",
             export=True,
+            model_save_dir=model_save_dir,
             **rbln_config_kwargs,
             **rbln_constructor_kwargs,
         )
 
         batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
-        unet_batch_size = batch_size * 2
+        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
+
         unet = RBLNUNet2DConditionModel.from_pretrained(
             model_id=model_id,
             subfolder="unet",
             export=True,
+            model_save_dir=model_save_dir,
             rbln_max_seq_len=text_encoder.config.max_position_embeddings,
             rbln_batch_size=unet_batch_size,
             rbln_use_encode=False,
@@ -88,6 +95,14 @@ class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
             **rbln_constructor_kwargs,
         )
 
+        if model_save_dir is not None:
+            # To skip saving original pytorch modules
+            del (model.vae, model.text_encoder, model.unet)
+
+            # Direct calling of `save_pretrained` causes config.unet = (None, None).
+            # So config must be saved again, later.
+            model.save_pretrained(model_save_dir)
+
         # replace modules
         model.vae = vae
         model.text_encoder = text_encoder
@@ -101,6 +116,10 @@ class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
         }
         model.register_to_config(**update_dict)
 
+        if model_save_dir is not None:
+            # overwrite to replace incorrect config
+            model.save_config(model_save_dir)
+
         model.models = [vae.model[0], text_encoder.model[0], unet.model[0]]
 
         return model

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

@@ -22,9 +22,6 @@
 # from Rebellions Inc.
 """RBLNStableDiffusionPipeline class for inference of diffusion models on rbln devices."""
 
-from pathlib import Path
-from tempfile import TemporaryDirectory
-
 from diffusers import StableDiffusionImg2ImgPipeline
 
 from ....modeling_base import RBLNBaseModel
@@ -52,21 +49,23 @@ class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
                 - A path to a *directory* containing a model saved using [`~OptimizedModel.save_pretrained`],
         """
         export = kwargs.pop("export", None)
+        model_save_dir = kwargs.pop("model_save_dir", None)
+        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
         model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
+
         if export is None or export is False:
             return model
 
-        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
-
-        save_dir = TemporaryDirectory()
-        save_dir_path = Path(save_dir.name)
-        model.save_pretrained(save_directory=save_dir_path, **kwargs)
+        do_classifier_free_guidance = (
+            rbln_config_kwargs.pop("rbln_guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
+        )
 
         # compile model, create runtime
         vae = RBLNAutoencoderKL.from_pretrained(
             model_id=model_id,
             subfolder="vae",
             export=True,
+            model_save_dir=model_save_dir,
             rbln_unet_sample_size=model.unet.config.sample_size,
             rbln_use_encode=True,
             rbln_vae_scale_factor=model.vae_scale_factor,
@@ -77,17 +76,19 @@ class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
             model_id=model_id,
             subfolder="text_encoder",
             export=True,
+            model_save_dir=model_save_dir,
             **rbln_config_kwargs,
             **rbln_constructor_kwargs,
         )
 
         batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
-        unet_batch_size = batch_size * 2
+        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
 
         unet = RBLNUNet2DConditionModel.from_pretrained(
             model_id=model_id,
             subfolder="unet",
             export=True,
+            model_save_dir=model_save_dir,
             rbln_max_seq_len=text_encoder.config.max_position_embeddings,
             rbln_batch_size=unet_batch_size,
             rbln_use_encode=True,
@@ -97,6 +98,14 @@ class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
             **rbln_constructor_kwargs,
         )
 
+        if model_save_dir is not None:
+            # To skip saving original pytorch modules
+            del (model.vae, model.text_encoder, model.unet)
+
+            # Direct calling of `save_pretrained` causes config.unet = (None, None).
+            # So config must be saved again, later.
+            model.save_pretrained(model_save_dir)
+
         # replace modules
         model.vae = vae
         model.text_encoder = text_encoder
@@ -110,6 +119,10 @@ class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
         }
         model.register_to_config(**update_dict)
 
+        if model_save_dir is not None:
+            # overwrite to replace incorrect config
+            model.save_config(model_save_dir)
+
         # vae encoder, vae decoder, text_encoder, unet
         model.models = [vae.model[0], vae.model[1], text_encoder.model[0], unet.model[0]]
 

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 """RBLNStableDiffusionXLPipeline class for inference of diffusion models on rbln devices."""
 
-
 from diffusers import StableDiffusionXLPipeline
 
 from ....modeling_base import RBLNBaseModel
@@ -42,12 +41,13 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
                 - A path to a *directory* containing a model saved using [`~OptimizedModel.save_pretrained`],
         """
         export = kwargs.pop("export", None)
+        model_save_dir = kwargs.pop("model_save_dir", None)
+        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
         model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
+
         if export is None or export is False:
             return model
 
-        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
-
         do_classifier_free_guidance = (
             rbln_config_kwargs.pop("rbln_guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
@@ -56,6 +56,7 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
             model_id=model_id,
             subfolder="vae",
             export=True,
+            model_save_dir=model_save_dir,
             rbln_unet_sample_size=model.unet.config.sample_size,
             rbln_use_encode=False,
             **rbln_config_kwargs,
@@ -65,6 +66,7 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
             model_id=model_id,
             subfolder="text_encoder",
             export=True,
+            model_save_dir=model_save_dir,
             **rbln_config_kwargs,
             **rbln_constructor_kwargs,
         )
@@ -72,6 +74,7 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
             model_id=model_id,
             subfolder="text_encoder_2",
             export=True,
+            model_save_dir=model_save_dir,
             **rbln_config_kwargs,
             **rbln_constructor_kwargs,
         )
@@ -83,6 +86,7 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
             model_id=model_id,
             subfolder="unet",
             export=True,
+            model_save_dir=model_save_dir,
             rbln_max_seq_len=model.text_encoder.config.max_position_embeddings,
             rbln_text_model_hidden_size=model.text_encoder_2.config.hidden_size,
             rbln_batch_size=unet_batch_size,
@@ -92,6 +96,14 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
             **rbln_constructor_kwargs,
         )
 
+        if model_save_dir is not None:
+            # To skip saving original pytorch modules
+            del (model.vae, model.text_encoder, model.unet, model.text_encoder_2)
+
+            # Direct calling of `save_pretrained` causes config.unet = (None, None).
+            # So config must be saved again, later.
+            model.save_pretrained(model_save_dir)
+
         model.vae = vae
         model.text_encoder = text_encoder
         model.unet = unet
@@ -104,6 +116,10 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
         }
         model.register_to_config(**update_dict)
 
+        if model_save_dir is not None:
+            # overwrite to replace incorrect config
+            model.save_config(model_save_dir)
+
         model.models = [vae.model[0], unet.model[0], text_encoder.model[0], text_encoder_2.model[0]]
 
         return model

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 """RBLNStableDiffusionXLPipeline class for inference of diffusion models on rbln devices."""
 
-
 from diffusers import StableDiffusionXLImg2ImgPipeline
 
 from ....modeling_base import RBLNBaseModel
@@ -42,12 +41,13 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
                 - A path to a *directory* containing a model saved using [`~OptimizedModel.save_pretrained`],
         """
         export = kwargs.pop("export", None)
+        model_save_dir = kwargs.pop("model_save_dir", None)
+        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
         model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
+
         if export is None or export is False:
             return model
 
-        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
-
         do_classifier_free_guidance = (
             rbln_config_kwargs.pop("rbln_guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
@@ -56,6 +56,7 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
             model_id=model_id,
             subfolder="vae",
             export=True,
+            model_save_dir=model_save_dir,
             rbln_unet_sample_size=model.unet.config.sample_size,
             rbln_use_encode=True,
             rbln_vae_scale_factor=model.vae_scale_factor,
@@ -66,6 +67,7 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
             model_id=model_id,
             subfolder="text_encoder",
             export=True,
+            model_save_dir=model_save_dir,
             **rbln_config_kwargs,
             **rbln_constructor_kwargs,
         )
@@ -73,6 +75,7 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
             model_id=model_id,
             subfolder="text_encoder_2",
             export=True,
+            model_save_dir=model_save_dir,
             **rbln_config_kwargs,
             **rbln_constructor_kwargs,
         )
@@ -84,6 +87,7 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
             model_id=model_id,
             subfolder="unet",
             export=True,
+            model_save_dir=model_save_dir,
             rbln_max_seq_len=model.text_encoder.config.max_position_embeddings,
             rbln_text_model_hidden_size=model.text_encoder_2.config.hidden_size,
             rbln_batch_size=unet_batch_size,
@@ -94,6 +98,14 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
             **rbln_constructor_kwargs,
         )
 
+        if model_save_dir is not None:
+            # To skip saving original pytorch modules
+            del (model.vae, model.text_encoder, model.unet, model.text_encoder_2)
+
+            # Direct calling of `save_pretrained` causes config.unet = (None, None).
+            # So config must be saved again, later.
+            model.save_pretrained(model_save_dir)
+
         model.vae = vae
         model.text_encoder = text_encoder
         model.unet = unet
@@ -106,6 +118,10 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
         }
         model.register_to_config(**update_dict)
 
+        if model_save_dir is not None:
+            # overwrite to replace incorrect config
+            model.save_config(model_save_dir)
+
         model.models = [vae.model[0], vae.model[1], unet.model[0], text_encoder.model[0], text_encoder_2.model[0]]
 
         return model

optimum/rbln/modeling_base.py

@@ -99,7 +99,7 @@ class RBLNBaseModel(OptimizedModel, ABC):
 
     model_type = "rbln_model"
     auto_model_class = AutoModel  # feature extraction
-    config_name = "model_index.json"
+    config_name = "config.json"
 
     def __init__(
         self,
@@ -109,7 +109,8 @@ class RBLNBaseModel(OptimizedModel, ABC):
         rbln_config: Optional[RBLNConfig],
         rbln_device: Optional[List[int]] = None,
         rbln_device_map: Optional[Dict[str, int]] = None,
-        rbln_create_runtimes: Optional[bool] = True,
+        rbln_create_runtimes: Optional[bool] = None,
+        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         **kwargs,
     ):
         super().__init__(models, config)
@@ -145,9 +146,24 @@ class RBLNBaseModel(OptimizedModel, ABC):
 
         self.device = torch.device("cpu")
 
+        if rbln_create_runtimes is None:
+            rbln_create_runtimes = rebel.npu_is_available()
+
         # create runtimes only if `rbln_create_runtimes` is enabled
         self.runtimes = self._create_runtimes(self.rbln_device_map) if rbln_create_runtimes else UnavailableRuntime()
 
+        # FIXME :: model_save_dir is not used after initialized. (This can be used when save/load)
+        # This attribute is needed to keep one reference on the temporary directory, since garbage collecting it
+        # would end-up removing the directory containing the underlying ONNX model.
+        self._model_save_dir_tempdirectory_instance = None
+        if isinstance(model_save_dir, TemporaryDirectory):
+            self._model_save_dir_tempdirectory_instance = model_save_dir
+            self.model_save_dir = Path(model_save_dir.name)
+        elif isinstance(model_save_dir, str):
+            self.model_save_dir = Path(model_save_dir)
+        else:
+            self.model_save_dir = model_save_dir
+
         self.__post_init__(**kwargs)
 
     def __post_init__(self, **kwargs):
@@ -179,6 +195,7 @@ class RBLNBaseModel(OptimizedModel, ABC):
         cache_dir: Optional[str] = None,
         subfolder: str = "",
         local_files_only: bool = False,
+        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         **kwargs,
     ) -> "RBLNBaseModel":
         model_path = Path(model_id)
@@ -216,6 +233,7 @@ class RBLNBaseModel(OptimizedModel, ABC):
                 rebel.RBLNCompiledModel(model_path / f"{compiled_model_name}.rbln")
                 for compiled_model_name in rbln_config
             ]
+            new_model_save_dir = model_path
 
         else:
             rbln_config_filename = rbln_config_filenames[0]
@@ -243,14 +261,19 @@ class RBLNBaseModel(OptimizedModel, ABC):
                     local_files_only=local_files_only,
                 )
                 models.append(rebel.RBLNCompiledModel(model_cache_path))
+            new_model_save_dir = Path(rbln_config_cache_path).parent
 
         preprocessors = maybe_load_preprocessors(model_id, subfolder=subfolder)
 
+        if model_save_dir is None:
+            model_save_dir = new_model_save_dir
+
         return cls(
             models,
             config,
             preprocessors,
             rbln_config=rbln_config,
+            model_save_dir=model_save_dir,
             **kwargs,
         )
 
@@ -370,6 +393,7 @@ class RBLNBaseModel(OptimizedModel, ABC):
         subfolder: str = "",
         local_files_only: bool = False,
         trust_remote_code: bool = False,
+        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         **kwargs,
     ):
         """
@@ -408,6 +432,7 @@ class RBLNModel(RBLNBaseModel):
         subfolder: str = "",
         local_files_only: bool = False,
         trust_remote_code: bool = False,
+        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         **kwargs,
     ) -> "RBLNModel":
         """
@@ -417,8 +442,16 @@ class RBLNModel(RBLNBaseModel):
         if task is None:
             task = TasksManager.infer_task_from_model(cls.auto_model_class)
 
-        save_dir = TemporaryDirectory()
-        save_dir_path = Path(save_dir.name)
+        if model_save_dir is None:
+            save_dir = TemporaryDirectory()
+            save_dir_path = Path(save_dir.name)
+        else:
+            save_dir = model_save_dir
+            if isinstance(save_dir, TemporaryDirectory):
+                save_dir_path = Path(model_save_dir.name)
+            else:
+                save_dir_path = Path(model_save_dir)
+                save_dir_path.mkdir(exist_ok=True)
 
         kwargs.update(
             {
@@ -457,7 +490,7 @@ class RBLNModel(RBLNBaseModel):
         preprocessors = maybe_save_preprocessors(model_id, save_dir_path, src_subfolder=subfolder)
 
         # Get compilation arguments
-        if rbln_config_kwargs.get("rbln_config", None) is None:
+        if (rbln_config := rbln_config_kwargs.pop("rbln_config", None)) is None:
             rbln_config = cls.get_rbln_config(preprocessors=preprocessors, model_config=config, **rbln_config_kwargs)
 
         rbln_runtime_configs = list(rbln_config.values())
@@ -562,7 +595,7 @@ class RBLNModelForImageClassification(RBLNModel):
                     rbln_image_size = processor.size["shortest_edge"]
                     break
             if rbln_image_size is None:
-                raise ValueError("`rbln_rbln_image_size` should be specified!")
+                raise ValueError("`rbln_image_size` should be specified!")
 
         if rbln_batch_size is None:
             rbln_batch_size = 1

optimum/rbln/modeling_seq2seq.py

@@ -160,6 +160,7 @@ class RBLNModelForSeq2SeqLM(RBLNBaseModel):
         subfolder: str = "",
         local_files_only: bool = False,
         trust_remote_code: bool = False,
+        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         **kwargs,
     ) -> "AutoModelForSeq2SeqLM":
         """
@@ -169,8 +170,16 @@ class RBLNModelForSeq2SeqLM(RBLNBaseModel):
         if task is None:
             task = TasksManager.infer_task_from_model(cls.auto_model_class)
 
-        save_dir = TemporaryDirectory()
-        save_dir_path = Path(save_dir.name)
+        if model_save_dir is None:
+            save_dir = TemporaryDirectory()
+            save_dir_path = Path(save_dir.name)
+        else:
+            save_dir = model_save_dir
+            if isinstance(save_dir, TemporaryDirectory):
+                save_dir_path = Path(model_save_dir.name)
+            else:
+                save_dir_path = Path(model_save_dir)
+                save_dir_path.mkdir(exist_ok=True)
 
         kwargs.update(
             {
@@ -339,6 +348,8 @@ class RBLNModelForSeq2SeqLM(RBLNBaseModel):
         if max_position_embeddings is not None and rbln_dec_max_seq_len > max_position_embeddings:
             raise ValueError("`rbln_dec_max_seq_len` should be less or equal than max_position_embeddings!")
 
+        rbln_batch_size = 1 if rbln_batch_size is None else rbln_batch_size
+
         meta["rbln_enc_max_seq_len"] = rbln_enc_max_seq_len
         meta["rbln_dec_max_seq_len"] = rbln_dec_max_seq_len
         meta["rbln_batch_size"] = rbln_batch_size
@@ -429,9 +440,13 @@ class RBLNModelForSeq2SeqLM(RBLNBaseModel):
         return repr(self.runtimes[0]) + "\n" + repr(self.runtimes[1])
 
     def _prepare_encoder_decoder_kwargs_for_generation(
-        self, inputs_tensor: torch.Tensor, model_kwargs, model_input_name: Optional[str] = None
+        self,
+        inputs_tensor: torch.Tensor,
+        model_kwargs,
+        model_input_name: Optional[str] = None,
+        *args,
+        **kwargs,
     ) -> Dict[str, Any]:
-
         ########## thkim change start ###################
         # padding input_ids & attention_mask regardless of user's tokenizer usage
         batch_size, input_len = inputs_tensor.shape

optimum/rbln/transformers/__init__.py

@@ -35,6 +35,7 @@ _import_structure = {
         "RBLNWav2Vec2ForCTC",
         "RBLNWhisperForConditionalGeneration",
         "RBLNLlamaForCausalLM",
+        "RBLNMidmLMHeadModel",
     ],
 }
 
@@ -45,6 +46,7 @@ if TYPE_CHECKING:
         RBLNCLIPTextModelWithProjection,
         RBLNGPT2LMHeadModel,
         RBLNLlamaForCausalLM,
+        RBLNMidmLMHeadModel,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,
     )

optimum/rbln/transformers/generation/__init__.py

@@ -22,3 +22,4 @@
 # from Rebellions Inc.
 
 from .streamers import BatchTextIteratorStreamer
+from .utils import RBLNGenerationMixin

optimum/rbln/transformers/generation/streamers.py

@@ -64,11 +64,13 @@ class BatchTextIteratorStreamer(TextIteratorStreamer):
         self.batch_size: int = batch_size
         self.token_cache: List[List[int]] = [[] for _ in range(batch_size)]
         self.print_len = [0] * batch_size
+        self.blocked = False
 
     def put(self, value):
         """
         Receives tokens, decodes them, and prints them to buffer as soon as they form entire words.
         """
+
         if len(value.shape) < 2:
             value = torch.reshape(value, (self.batch_size, value.shape[0] // self.batch_size))
 
@@ -115,8 +117,23 @@ class BatchTextIteratorStreamer(TextIteratorStreamer):
 
         self.next_tokens_are_prompt = True
         self.on_finalized_text(batch_printable_text, stream_end=True)
+        self.blocked = False
 
     def on_finalized_text(self, texts: List[str], stream_end: bool = False):
         self.text_queue.put(texts, timeout=self.timeout)
         if stream_end:
             self.text_queue.put(self.stop_signal, timeout=self.timeout)
+
+    # thkim change for demo
+    def __next__(self):
+        value = self.text_queue.get(timeout=self.timeout)
+        if value == self.stop_signal:
+            raise StopIteration()
+        else:
+            return value
+
+    def block(self):
+        self.blocked = True
+
+    def is_blocked(self):
+        return self.blocked