optimum-rbln 0.1.0__py3-none-any.whl → 0.1.1__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 (29) hide show
  1. optimum/rbln/__init__.py +6 -0
  2. optimum/rbln/__version__.py +1 -1
  3. optimum/rbln/diffusers/__init__.py +7 -0
  4. optimum/rbln/diffusers/models/autoencoder_kl.py +30 -9
  5. optimum/rbln/diffusers/models/controlnet.py +93 -23
  6. optimum/rbln/diffusers/models/unet_2d_condition.py +78 -61
  7. optimum/rbln/diffusers/pipelines/__init__.py +7 -2
  8. optimum/rbln/diffusers/pipelines/controlnet/__init__.py +4 -0
  9. optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet.py +768 -0
  10. optimum/rbln/diffusers/pipelines/{stable_diffusion → controlnet}/pipeline_controlnet_img2img.py +25 -16
  11. optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +942 -0
  12. optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +955 -0
  13. optimum/rbln/diffusers/pipelines/stable_diffusion/__init__.py +0 -1
  14. optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +23 -4
  15. optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +22 -9
  16. optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +19 -3
  17. optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +19 -3
  18. optimum/rbln/modeling_base.py +36 -3
  19. optimum/rbln/modeling_seq2seq.py +19 -4
  20. optimum/rbln/transformers/generation/__init__.py +1 -0
  21. optimum/rbln/transformers/generation/streamers.py +17 -0
  22. optimum/rbln/transformers/generation/utils.py +399 -0
  23. optimum/rbln/transformers/models/gpt2/modeling_gpt2.py +24 -333
  24. optimum/rbln/transformers/models/llama/modeling_llama.py +63 -45
  25. optimum/rbln/transformers/models/whisper/modeling_whisper.py +13 -3
  26. {optimum_rbln-0.1.0.dist-info → optimum_rbln-0.1.1.dist-info}/METADATA +1 -1
  27. {optimum_rbln-0.1.0.dist-info → optimum_rbln-0.1.1.dist-info}/RECORD +29 -25
  28. {optimum_rbln-0.1.0.dist-info → optimum_rbln-0.1.1.dist-info}/WHEEL +0 -0
  29. {optimum_rbln-0.1.0.dist-info → optimum_rbln-0.1.1.dist-info}/licenses/LICENSE +0 -0
@@ -46,28 +46,43 @@ logger = logging.get_logger(__name__)
46
46
  class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2ImgPipeline):
47
47
  @classmethod
48
48
  def from_pretrained(cls, model_id, **kwargs):
49
+ """
50
+ Pipeline for image-to-image generation using Stable Diffusion with ControlNet.
51
+
52
+ This model inherits from [`StableDiffusionControlNetImg2ImgPipeline`]. Check the superclass documentation for the generic methods
53
+ implemented for all pipelines (downloading, saving, running on a particular device, etc.).
54
+
55
+ It implements the methods to convert a pre-trained Stable Diffusion Controlnet pipeline into a RBLNStableDiffusionControlNetImg2Img pipeline by:
56
+ - transferring the checkpoint weights of the original into an optimized RBLN graph,
57
+ - compiling the resulting graph using the RBLN compiler.
58
+
59
+ Args:
60
+ model_id (`Union[str, Path]`):
61
+ Can be either:
62
+ - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
63
+ - A path to a *directory* containing a model saved using [`~OptimizedModel.save_pretrained`],
64
+ """
49
65
  export = kwargs.pop("export", None)
50
66
  text_encoder = kwargs.pop("text_encoder", None)
51
67
  controlnets = kwargs.pop("controlnet", None)
52
68
 
69
+ rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
70
+
53
71
  kwargs_dict = {
54
72
  "pretrained_model_name_or_path": model_id,
55
73
  "text_encoder": text_encoder,
56
74
  "controlnet": controlnets,
57
75
  **kwargs,
58
76
  }
77
+
59
78
  model = super().from_pretrained(**{k: v for k, v in kwargs_dict.items() if v is not None})
60
79
 
61
80
  if export is None or export is False:
62
81
  return model
63
82
 
64
- rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
65
-
66
- rbln_callback = rbln_config_kwargs.pop("rbln_callback", None)
67
- img_width = kwargs.pop("img_width", None)
68
- img_height = kwargs.pop("img_height", None)
69
-
70
- model = rbln_callback(model)
83
+ do_classifier_free_guidance = (
84
+ rbln_config_kwargs.pop("rbln_guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
85
+ )
71
86
 
72
87
  save_dir = TemporaryDirectory()
73
88
  save_dir_path = Path(save_dir.name)
@@ -80,8 +95,6 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
80
95
  export=True,
81
96
  rbln_unet_sample_size=model.unet.config.sample_size,
82
97
  rbln_use_encode=True,
83
- rbln_img_width=img_width,
84
- rbln_img_height=img_height,
85
98
  rbln_vae_scale_factor=model.vae_scale_factor,
86
99
  **rbln_config_kwargs,
87
100
  **rbln_constructor_kwargs,
@@ -95,7 +108,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
95
108
  )
96
109
 
97
110
  batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
98
- unet_batch_size = batch_size * 2
111
+ unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
99
112
 
100
113
  unet = RBLNUNet2DConditionModel.from_pretrained(
101
114
  model_id=save_dir_path / "unet",
@@ -103,8 +116,6 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
103
116
  rbln_max_seq_len=text_encoder.config.max_position_embeddings,
104
117
  rbln_batch_size=unet_batch_size,
105
118
  rbln_use_encode=True,
106
- rbln_img_width=img_width,
107
- rbln_img_height=img_height,
108
119
  rbln_vae_scale_factor=model.vae_scale_factor,
109
120
  rbln_is_controlnet=True if "controlnet" in model.config.keys() else False,
110
121
  **rbln_config_kwargs,
@@ -116,8 +127,6 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
116
127
  model_id=str(save_dir_path / "controlnet"),
117
128
  export=True,
118
129
  rbln_batch_size=unet_batch_size,
119
- rbln_img_width=img_width,
120
- rbln_img_height=img_height,
121
130
  rbln_vae_scale_factor=model.vae_scale_factor,
122
131
  **rbln_config_kwargs,
123
132
  **rbln_constructor_kwargs,
@@ -128,8 +137,6 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
128
137
  model_id=save_dir_path / "controlnet",
129
138
  export=True,
130
139
  rbln_batch_size=unet_batch_size,
131
- rbln_img_width=img_width,
132
- rbln_img_height=img_height,
133
140
  rbln_vae_scale_factor=model.vae_scale_factor,
134
141
  **rbln_config_kwargs,
135
142
  **rbln_constructor_kwargs,
@@ -151,6 +158,8 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
151
158
  }
152
159
  model.register_to_config(**update_dict)
153
160
 
161
+ model.models = [vae.model[0], vae.model[1], text_encoder.model[0], unet.model[0], controlnet.model[0]]
162
+
154
163
  return model
155
164
 
156
165
  def check_inputs(