forked from huggingface/optimum-intel
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmodeling_diffusion.py
1640 lines (1392 loc) · 64.5 KB
/
modeling_diffusion.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import importlib
import inspect
import logging
import os
import shutil
from abc import abstractmethod
from collections import OrderedDict
from copy import deepcopy
from pathlib import Path
from tempfile import gettempdir
from typing import Any, Dict, List, Optional, Union
import numpy as np
import openvino
import torch
from diffusers import (
AutoPipelineForImage2Image,
AutoPipelineForInpainting,
AutoPipelineForText2Image,
DiffusionPipeline,
LatentConsistencyModelImg2ImgPipeline,
LatentConsistencyModelPipeline,
StableDiffusionImg2ImgPipeline,
StableDiffusionInpaintPipeline,
StableDiffusionPipeline,
StableDiffusionXLImg2ImgPipeline,
StableDiffusionXLInpaintPipeline,
StableDiffusionXLPipeline,
pipelines,
)
from diffusers.configuration_utils import ConfigMixin
from diffusers.schedulers import SchedulerMixin
from diffusers.schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME
from diffusers.utils.constants import CONFIG_NAME
from huggingface_hub import snapshot_download
from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
from huggingface_hub.utils import validate_hf_hub_args
from openvino._offline_transformations import compress_model_transformation
from openvino.runtime import Core
from transformers import CLIPFeatureExtractor, CLIPTokenizer
from transformers.modeling_outputs import ModelOutput
from optimum.utils import (
DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER,
DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER,
DIFFUSION_MODEL_UNET_SUBFOLDER,
DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER,
DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER,
)
from ...exporters.openvino import main_export
from ..utils.import_utils import is_diffusers_version, is_openvino_version
from .configuration import OVConfig, OVQuantizationMethod, OVWeightQuantizationConfig
from .loaders import OVTextualInversionLoaderMixin
from .modeling_base import OVBaseModel
from .utils import (
ONNX_WEIGHTS_NAME,
OV_TO_PT_TYPE,
OV_XML_FILE_NAME,
TemporaryDirectory,
_print_compiled_model_properties,
check_scale_available,
model_has_dynamic_inputs,
np_to_pt_generators,
)
if is_diffusers_version(">=", "0.25.0"):
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
else:
from diffusers.models.vae import DiagonalGaussianDistribution
if is_diffusers_version(">=", "0.29.0"):
from diffusers import StableDiffusion3Img2ImgPipeline, StableDiffusion3Pipeline
else:
StableDiffusion3Pipeline, StableDiffusion3Img2ImgPipeline = object, object
if is_diffusers_version(">=", "0.30.0"):
from diffusers import FluxPipeline, StableDiffusion3InpaintPipeline
else:
StableDiffusion3InpaintPipeline = object
FluxPipeline = object
if is_diffusers_version(">=", "0.31.0"):
from diffusers import FluxImg2ImgPipeline, FluxInpaintPipeline
else:
FluxImg2ImgPipeline = object
FluxInpaintPipeline = object
if is_diffusers_version(">=", "0.32.0"):
from diffusers import FluxFillPipeline
else:
FluxFillPipeline = object
DIFFUSION_MODEL_TRANSFORMER_SUBFOLDER = "transformer"
DIFFUSION_MODEL_TEXT_ENCODER_3_SUBFOLDER = "text_encoder_3"
core = Core()
logger = logging.getLogger(__name__)
# TODO: support DiffusionPipeline.from_pipe()
# TODO: makes more sense to have a compositional OVMixin class
# TODO: instead of one bloated __init__, we should consider an __init__ per pipeline
class OVDiffusionPipeline(OVBaseModel, DiffusionPipeline):
auto_model_class = DiffusionPipeline
config_name = "model_index.json"
_library_name = "diffusers"
def __init__(
self,
scheduler: SchedulerMixin,
unet: Optional[openvino.runtime.Model] = None,
vae_decoder: Optional[openvino.runtime.Model] = None,
# optional pipeline models
vae_encoder: Optional[openvino.runtime.Model] = None,
text_encoder: Optional[openvino.runtime.Model] = None,
text_encoder_2: Optional[openvino.runtime.Model] = None,
text_encoder_3: Optional[openvino.runtime.Model] = None,
transformer: Optional[openvino.runtime.Model] = None,
# optional pipeline submodels
tokenizer: Optional[CLIPTokenizer] = None,
tokenizer_2: Optional[CLIPTokenizer] = None,
tokenizer_3: Optional[CLIPTokenizer] = None,
feature_extractor: Optional[CLIPFeatureExtractor] = None,
# stable diffusion xl specific arguments
force_zeros_for_empty_prompt: bool = True,
requires_aesthetics_score: bool = False,
add_watermarker: Optional[bool] = None,
# openvino specific arguments
device: str = "CPU",
compile: bool = True,
compile_only: bool = False,
dynamic_shapes: bool = True,
ov_config: Optional[Dict[str, str]] = None,
model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
quantization_config: Optional[Union[OVWeightQuantizationConfig, Dict]] = None,
**kwargs,
):
self._device = device.upper()
self.is_dynamic = dynamic_shapes
self._compile_only = compile_only
self.model_save_dir = model_save_dir
self.ov_config = {} if ov_config is None else {**ov_config}
self.preprocessors = kwargs.get("preprocessors", [])
if self._compile_only:
if not compile:
raise ValueError(
"`compile_only` mode does not support disabling compilation."
"Please provide `compile=True` if you want to use `compile_only=True` or set `compile_only=False`"
)
main_model = unet if unet is not None else transformer
if not isinstance(main_model, openvino.runtime.CompiledModel):
raise ValueError("`compile_only` expect that already compiled model will be provided")
model_is_dynamic = model_has_dynamic_inputs(main_model)
if dynamic_shapes ^ model_is_dynamic:
requested_shapes = "dynamic" if dynamic_shapes else "static"
compiled_shapes = "dynamic" if model_is_dynamic else "static"
raise ValueError(
f"Provided compiled model with {compiled_shapes} shapes but requested to use {requested_shapes}. "
f"Please set `compile_only=False` or `dynamic_shapes={model_is_dynamic}`"
)
self.unet = OVModelUnet(unet, self, DIFFUSION_MODEL_UNET_SUBFOLDER) if unet is not None else None
self.transformer = (
OVModelTransformer(transformer, self, DIFFUSION_MODEL_TRANSFORMER_SUBFOLDER)
if transformer is not None
else None
)
if unet is None and transformer is None:
raise ValueError("`unet` or `transformer` model should be provided for pipeline work")
self.vae_decoder = OVModelVaeDecoder(vae_decoder, self, DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER)
self.vae_encoder = (
OVModelVaeEncoder(vae_encoder, self, DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER)
if vae_encoder is not None
else None
)
self.text_encoder = (
OVModelTextEncoder(text_encoder, self, DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER)
if text_encoder is not None
else None
)
self.text_encoder_2 = (
OVModelTextEncoder(text_encoder_2, self, DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER)
if text_encoder_2 is not None
else None
)
self.text_encoder_3 = (
OVModelTextEncoder(text_encoder_3, self, DIFFUSION_MODEL_TEXT_ENCODER_3_SUBFOLDER)
if text_encoder_3 is not None
else None
)
# We wrap the VAE Decoder & Encoder in a single object to simulate diffusers API
self.vae = OVModelVae(decoder=self.vae_decoder, encoder=self.vae_encoder)
self.scheduler = scheduler
self.tokenizer = tokenizer
self.tokenizer_2 = tokenizer_2
self.tokenizer_3 = tokenizer_3
self.feature_extractor = feature_extractor
# we allow passing these as torch models for now
self.image_encoder = kwargs.pop("image_encoder", None) # TODO: maybe mplement OVModelImageEncoder
self.safety_checker = kwargs.pop("safety_checker", None) # TODO: maybe mplement OVModelSafetyChecker
all_pipeline_init_args = {
"vae": self.vae,
"unet": self.unet,
"transformer": self.transformer,
"text_encoder": self.text_encoder,
"text_encoder_2": self.text_encoder_2,
"text_encoder_3": self.text_encoder_3,
"safety_checker": self.safety_checker,
"image_encoder": self.image_encoder,
"scheduler": self.scheduler,
"tokenizer": self.tokenizer,
"tokenizer_2": self.tokenizer_2,
"tokenizer_3": self.tokenizer_3,
"feature_extractor": self.feature_extractor,
"requires_aesthetics_score": requires_aesthetics_score,
"force_zeros_for_empty_prompt": force_zeros_for_empty_prompt,
"add_watermarker": add_watermarker,
}
diffusers_pipeline_args = {}
for key in inspect.signature(self.auto_model_class).parameters.keys():
if key in all_pipeline_init_args:
diffusers_pipeline_args[key] = all_pipeline_init_args[key]
# inits diffusers pipeline specific attributes (registers modules and config)
self.auto_model_class.__init__(self, **diffusers_pipeline_args)
# we use auto_model_class.__init__ here because we can't call super().__init__
# as OptimizedModel already defines an __init__ which is the first in the MRO
self._openvino_config = None
if quantization_config:
self._openvino_config = OVConfig(quantization_config=quantization_config)
self._set_ov_config_parameters()
if self.is_dynamic and not self._compile_only:
self.reshape(batch_size=-1, height=-1, width=-1, num_images_per_prompt=-1)
if compile and not self._compile_only:
self.compile()
def _save_pretrained(self, save_directory: Union[str, Path]):
"""
Saves the model to the OpenVINO IR format so that it can be re-loaded using the
[`~optimum.intel.openvino.modeling.OVModel.from_pretrained`] class method.
Arguments:
save_directory (`str` or `Path`):
The directory where to save the model files
"""
if self._compile_only:
raise ValueError(
"`save_pretrained()` is not supported with `compile_only` mode, please intialize model without this option"
)
save_directory = Path(save_directory)
models_to_save_paths = {
(self.unet, save_directory / DIFFUSION_MODEL_UNET_SUBFOLDER),
(self.vae_decoder, save_directory / DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER),
(self.vae_encoder, save_directory / DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER),
(self.text_encoder, save_directory / DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER),
(self.text_encoder_2, save_directory / DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER),
(self.text_encoder_3, save_directory / DIFFUSION_MODEL_TEXT_ENCODER_3_SUBFOLDER),
(self.transformer, save_directory / DIFFUSION_MODEL_TRANSFORMER_SUBFOLDER),
}
for model, save_path in models_to_save_paths:
if model is not None:
dst_path = save_path / OV_XML_FILE_NAME
dst_path.parent.mkdir(parents=True, exist_ok=True)
openvino.save_model(model.model, dst_path, compress_to_fp16=False)
model_dir = model.config.get("_name_or_path", None) or model.model_save_dir
config_path = Path(model_dir) / CONFIG_NAME
if config_path.is_file():
config_save_path = save_path / CONFIG_NAME
shutil.copyfile(config_path, config_save_path)
else:
if hasattr(model, "save_config"):
model.save_config(save_path)
elif hasattr(model, "config") and hasattr(model.config, "save_pretrained"):
model.config.save_pretrained(save_path)
self.scheduler.save_pretrained(save_directory / "scheduler")
if self.tokenizer is not None:
self.tokenizer.save_pretrained(save_directory / "tokenizer")
if self.tokenizer_2 is not None:
self.tokenizer_2.save_pretrained(save_directory / "tokenizer_2")
if self.tokenizer_3 is not None:
self.tokenizer_3.save_pretrained(save_directory / "tokenizer_3")
if self.feature_extractor is not None:
self.feature_extractor.save_pretrained(save_directory / "feature_extractor")
if getattr(self, "safety_checker", None) is not None:
self.safety_checker.save_pretrained(save_directory / "safety_checker")
self._save_openvino_config(save_directory)
def _save_config(self, save_directory):
"""
Saves a model configuration into a directory, so that it can be re-loaded using the
[`from_pretrained`] class method.
"""
model_dir = (
self.model_save_dir
if not isinstance(self.model_save_dir, TemporaryDirectory)
else self.model_save_dir.name
)
save_dir = Path(save_directory)
original_config = Path(model_dir) / self.config_name
if original_config.exists():
if not save_dir.exists():
save_dir.mkdir(parents=True)
shutil.copy(original_config, save_dir)
else:
self.config.save_pretrained(save_dir)
@classmethod
def _from_pretrained(
cls,
model_id: Union[str, Path],
config: Dict[str, Any],
token: Optional[Union[bool, str]] = None,
revision: Optional[str] = None,
force_download: bool = False,
local_files_only: bool = False,
cache_dir: str = HUGGINGFACE_HUB_CACHE,
unet_file_name: Optional[str] = None,
vae_decoder_file_name: Optional[str] = None,
vae_encoder_file_name: Optional[str] = None,
text_encoder_file_name: Optional[str] = None,
text_encoder_2_file_name: Optional[str] = None,
text_encoder_3_file_name: Optional[str] = None,
transformer_file_name: Optional[str] = None,
from_onnx: bool = False,
load_in_8bit: bool = False,
quantization_config: Union[OVWeightQuantizationConfig, Dict] = None,
model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
**kwargs,
):
# same as DiffusionPipeline.from_pretraoned, if called directly, it loads the class in the config
if cls.__name__ == "OVDiffusionPipeline":
class_name = config["_class_name"]
ov_pipeline_class = _get_ov_class(class_name)
else:
ov_pipeline_class = cls
default_file_name = ONNX_WEIGHTS_NAME if from_onnx else OV_XML_FILE_NAME
unet_file_name = unet_file_name or default_file_name
vae_encoder_file_name = vae_encoder_file_name or default_file_name
vae_decoder_file_name = vae_decoder_file_name or default_file_name
text_encoder_file_name = text_encoder_file_name or default_file_name
text_encoder_2_file_name = text_encoder_2_file_name or default_file_name
text_encoder_3_file_name = text_encoder_3_file_name or default_file_name
transformer_file_name = transformer_file_name or default_file_name
if not os.path.isdir(str(model_id)):
all_components = {key for key in config.keys() if not key.startswith("_")} | {"vae_encoder", "vae_decoder"}
allow_patterns = {os.path.join(component, "*") for component in all_components}
allow_patterns.update(
{
unet_file_name,
transformer_file_name,
vae_encoder_file_name,
vae_decoder_file_name,
text_encoder_file_name,
text_encoder_2_file_name,
text_encoder_3_file_name,
unet_file_name.replace(".xml", ".bin"),
transformer_file_name.replace(".xml", ".bin"),
vae_encoder_file_name.replace(".xml", ".bin"),
vae_decoder_file_name.replace(".xml", ".bin"),
text_encoder_file_name.replace(".xml", ".bin"),
text_encoder_2_file_name.replace(".xml", ".bin"),
text_encoder_3_file_name.replace(".xml", ".bin"),
SCHEDULER_CONFIG_NAME,
cls.config_name,
CONFIG_NAME,
}
)
ignore_patterns = ["*.msgpack", "*.safetensors", "*pytorch_model.bin"]
if not from_onnx:
ignore_patterns.extend(["*.onnx", "*.onnx_data"])
model_save_folder = snapshot_download(
model_id,
cache_dir=cache_dir,
force_download=force_download,
local_files_only=local_files_only,
revision=revision,
token=token,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
)
else:
model_save_folder = str(model_id)
model_save_path = Path(model_save_folder)
if model_save_dir is None:
model_save_dir = model_save_path
submodels = {
"scheduler": None,
"tokenizer": None,
"tokenizer_2": None,
"tokenizer_3": None,
"feature_extractor": None,
"safety_checker": None,
"image_encoder": None,
}
for name in submodels.keys():
if name in kwargs:
submodels[name] = kwargs.pop(name)
elif config.get(name, (None, None))[0] is not None:
module_name, module_class = config.get(name)
if hasattr(pipelines, module_name):
module = getattr(pipelines, module_name)
else:
module = importlib.import_module(module_name)
class_obj = getattr(module, module_class)
load_method = getattr(class_obj, "from_pretrained")
# Check if the module is in a subdirectory
if (model_save_path / name).is_dir():
submodels[name] = load_method(model_save_path / name)
# For backward compatibility with models exported using previous optimum version, where safety_checker saving was disabled
elif name == "safety_checker":
logger.warning(
"Pipeline config contains `safety_checker` subcomponent, while `safety_checker` is not available in model directory. "
"`safety_checker` will be disabled. If you want to enable it please set it explicitly to `from_pretrained` method "
"or reexport model with new optimum-intel version"
)
submodels[name] = None
else:
submodels[name] = load_method(model_save_path)
models = {
"unet": model_save_path / DIFFUSION_MODEL_UNET_SUBFOLDER / unet_file_name,
"transformer": model_save_path / DIFFUSION_MODEL_TRANSFORMER_SUBFOLDER / transformer_file_name,
"vae_decoder": model_save_path / DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER / vae_decoder_file_name,
"vae_encoder": model_save_path / DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER / vae_encoder_file_name,
"text_encoder": model_save_path / DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER / text_encoder_file_name,
"text_encoder_2": model_save_path / DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER / text_encoder_2_file_name,
"text_encoder_3": model_save_path / DIFFUSION_MODEL_TEXT_ENCODER_3_SUBFOLDER / text_encoder_3_file_name,
}
for config_key, value in config.items():
if config_key not in models and config_key not in kwargs and config_key not in submodels:
kwargs[config_key] = value
compile_only = kwargs.get("compile_only", False)
quantization_config = cls._prepare_weight_quantization_config(quantization_config, load_in_8bit)
if (quantization_config is None or quantization_config.dataset is None) and not compile_only:
for name, path in models.items():
if name in kwargs:
models[name] = kwargs.pop(name)
else:
models[name] = cls.load_model(path, quantization_config) if path.is_file() else None
elif compile_only:
ov_config = kwargs.get("ov_config", {})
device = kwargs.get("device", "CPU")
vae_ov_conifg = {**ov_config}
if (
"GPU" in device.upper()
and "INFERENCE_PRECISION_HINT" not in vae_ov_conifg
and is_openvino_version("<=", "2025.0")
):
vae_model_path = models["vae_decoder"]
required_upcast = check_scale_available(vae_model_path)
if required_upcast:
vae_ov_conifg["INFERENCE_PRECISION_HINT"] = "f32"
for name, path in models.items():
if name in kwargs:
models[name] = kwargs.pop(name)
else:
models[name] = (
cls._compile_model(
path,
device,
ov_config if "vae" not in name else vae_ov_conifg,
Path(model_save_dir) / name,
)
if path.is_file()
else None
)
else:
# why is this quantization not performed in __init__?
if ov_pipeline_class.export_feature != "text-to-image":
raise NotImplementedError(f"Quantization in hybrid mode is not supported for {cls.__name__}")
from optimum.intel import OVQuantizer
for name, path in models.items():
if name in kwargs:
models[name] = kwargs.pop(name)
else:
models[name] = cls.load_model(path) if path.is_file() else None
ov_pipeline = ov_pipeline_class(**models, **submodels, model_save_dir=model_save_dir, **kwargs)
# same as in DiffusionPipeline.from_pretrained, we save where the model was instantiated from
ov_pipeline.register_to_config(_name_or_path=config.get("_name_or_path", str(model_id)))
hybrid_quantization_config = deepcopy(quantization_config)
hybrid_quantization_config.quant_method = OVQuantizationMethod.HYBRID
quantizer = OVQuantizer(ov_pipeline)
quantizer.quantize(ov_config=OVConfig(quantization_config=hybrid_quantization_config))
return ov_pipeline
ov_pipeline = ov_pipeline_class(
**models,
**submodels,
model_save_dir=model_save_dir,
quantization_config=quantization_config,
**kwargs,
)
# same as in DiffusionPipeline.from_pretrained, we save where the model was instantiated from
ov_pipeline.register_to_config(_name_or_path=config.get("_name_or_path", str(model_id)))
return ov_pipeline
@classmethod
def _from_transformers(
cls,
model_id: str,
config: Dict[str, Any],
token: Optional[Union[bool, str]] = None,
revision: Optional[str] = None,
force_download: bool = False,
cache_dir: str = HUGGINGFACE_HUB_CACHE,
local_files_only: bool = False,
load_in_8bit: Optional[bool] = None,
quantization_config: Union[OVWeightQuantizationConfig, Dict] = None,
compile_only: bool = False,
**kwargs,
):
if compile_only:
logger.warning(
"`compile_only` mode will be disabled because it does not support model export."
"Please provide openvino model obtained using optimum-cli or saved on disk using `save_pretrained`"
)
compile_only = False
# If load_in_8bit and quantization_config not specified then ov_config is set
# to None and will be set by default in convert depending on the model size
if load_in_8bit is None and not quantization_config:
ov_config = None
else:
ov_config = OVConfig(dtype="fp32")
model_save_dir = TemporaryDirectory()
model_save_path = Path(model_save_dir.name)
main_export(
model_name_or_path=model_id,
output=model_save_path,
do_validation=False,
no_post_process=True,
revision=revision,
cache_dir=cache_dir,
task=cls.export_feature,
token=token,
local_files_only=local_files_only,
force_download=force_download,
ov_config=ov_config,
library_name=cls._library_name,
)
return cls._from_pretrained(
model_id=model_save_path,
config=config,
from_onnx=False,
token=token,
revision=revision,
cache_dir=cache_dir,
force_download=force_download,
local_files_only=local_files_only,
model_save_dir=model_save_dir,
quantization_config=quantization_config,
load_in_8bit=load_in_8bit,
compile_only=compile_only,
**kwargs,
)
def to(self, *args, device: Optional[str] = None, dtype: Optional[torch.dtype] = None):
for arg in args:
if isinstance(arg, str):
device = arg
elif isinstance(arg, torch.dtype):
dtype = arg
if isinstance(device, str):
self._device = device.upper()
self.clear_requests()
elif device is not None:
raise ValueError(
"The `device` argument should be a string representing the device on which the model should be loaded."
)
if dtype is not None and dtype != self.dtype:
raise NotImplementedError(
f"Cannot change the dtype of the model from {self.dtype} to {dtype}. "
f"Please export the model with the desired dtype."
)
return self
@property
def height(self) -> int:
model = self.vae.decoder.model
height = model.inputs[0].get_partial_shape()[2]
if height.is_dynamic:
return -1
return height.get_length() * self.vae_scale_factor
@property
def width(self) -> int:
model = self.vae.decoder.model
width = model.inputs[0].get_partial_shape()[3]
if width.is_dynamic:
return -1
return width.get_length() * self.vae_scale_factor
@property
def batch_size(self) -> int:
model = self.unet.model if self.unet is not None else self.transformer.model
batch_size = model.inputs[0].get_partial_shape()[0]
if batch_size.is_dynamic:
return -1
return batch_size.get_length()
def _reshape_unet(
self,
model: openvino.runtime.Model,
batch_size: int = -1,
height: int = -1,
width: int = -1,
num_images_per_prompt: int = -1,
tokenizer_max_length: int = -1,
):
if batch_size == -1 or num_images_per_prompt == -1:
batch_size = -1
else:
batch_size *= num_images_per_prompt
# The factor of 2 comes from the guidance scale > 1
if "timestep_cond" not in {inputs.get_any_name() for inputs in model.inputs}:
batch_size *= 2
height = height // self.vae_scale_factor if height > 0 else height
width = width // self.vae_scale_factor if width > 0 else width
shapes = {}
for inputs in model.inputs:
shapes[inputs] = inputs.get_partial_shape()
if inputs.get_any_name() == "timestep":
if shapes[inputs].rank == 1:
shapes[inputs][0] = 1
elif inputs.get_any_name() == "sample":
in_channels = self.unet.config.get("in_channels", None)
if in_channels is None:
in_channels = shapes[inputs][1]
if in_channels.is_dynamic:
logger.warning(
"Could not identify `in_channels` from the unet configuration, to statically reshape the unet please provide a configuration."
)
self.is_dynamic = True
shapes[inputs] = [batch_size, in_channels, height, width]
elif inputs.get_any_name() == "text_embeds":
shapes[inputs] = [batch_size, self.text_encoder_2.config["projection_dim"]]
elif inputs.get_any_name() == "time_ids":
shapes[inputs] = [batch_size, inputs.get_partial_shape()[1]]
elif inputs.get_any_name() == "timestep_cond":
shapes[inputs] = [batch_size, self.unet.config["time_cond_proj_dim"]]
else:
shapes[inputs][0] = batch_size
shapes[inputs][1] = tokenizer_max_length
model.reshape(shapes)
return model
def _reshape_transformer(
self,
model: openvino.runtime.Model,
batch_size: int = -1,
height: int = -1,
width: int = -1,
num_images_per_prompt: int = -1,
tokenizer_max_length: int = -1,
):
if batch_size == -1 or num_images_per_prompt == -1:
batch_size = -1
else:
# The factor of 2 comes from the guidance scale > 1
batch_size *= num_images_per_prompt
if "img_ids" not in {inputs.get_any_name() for inputs in model.inputs}:
batch_size *= 2
height = height // self.vae_scale_factor if height > 0 else height
width = width // self.vae_scale_factor if width > 0 else width
packed_height = height // 2 if height > 0 else height
packed_width = width // 2 if width > 0 else width
packed_height_width = packed_width * packed_height if height > 0 and width > 0 else -1
shapes = {}
for inputs in model.inputs:
shapes[inputs] = inputs.get_partial_shape()
if inputs.get_any_name() in ["timestep", "guidance"]:
shapes[inputs][0] = batch_size
elif inputs.get_any_name() == "hidden_states":
in_channels = self.transformer.config.get("in_channels", None)
if in_channels is None:
in_channels = (
shapes[inputs][1] if inputs.get_partial_shape().rank.get_length() == 4 else shapes[inputs][2]
)
if in_channels.is_dynamic:
logger.warning(
"Could not identify `in_channels` from the unet configuration, to statically reshape the unet please provide a configuration."
)
self.is_dynamic = True
if inputs.get_partial_shape().rank.get_length() == 4:
shapes[inputs] = [batch_size, in_channels, height, width]
else:
shapes[inputs] = [batch_size, packed_height_width, in_channels]
elif inputs.get_any_name() == "pooled_projections":
shapes[inputs] = [batch_size, self.transformer.config["pooled_projection_dim"]]
elif inputs.get_any_name() == "img_ids":
shapes[inputs] = (
[batch_size, packed_height_width, 3]
if is_diffusers_version("<", "0.31.0")
else [packed_height_width, 3]
)
elif inputs.get_any_name() == "txt_ids":
shapes[inputs] = [batch_size, -1, 3] if is_diffusers_version("<", "0.31.0") else [-1, 3]
else:
shapes[inputs][0] = batch_size
shapes[inputs][1] = -1 # text_encoder_3 may have vary input length
model.reshape(shapes)
return model
def _reshape_text_encoder(
self, model: openvino.runtime.Model, batch_size: int = -1, tokenizer_max_length: int = -1
):
if batch_size != -1:
shapes = {model.inputs[0]: [batch_size, tokenizer_max_length]}
model.reshape(shapes)
return model
def _reshape_vae_encoder(
self, model: openvino.runtime.Model, batch_size: int = -1, height: int = -1, width: int = -1
):
in_channels = self.vae_encoder.config.get("in_channels", None)
if in_channels is None:
in_channels = model.inputs[0].get_partial_shape()[1]
if in_channels.is_dynamic:
logger.warning(
"Could not identify `in_channels` from the VAE encoder configuration, to statically reshape the VAE encoder please provide a configuration."
)
self.is_dynamic = True
shapes = {model.inputs[0]: [batch_size, in_channels, height, width]}
model.reshape(shapes)
return model
def _reshape_vae_decoder(
self, model: openvino.runtime.Model, height: int = -1, width: int = -1, num_images_per_prompt: int = -1
):
height = height // self.vae_scale_factor if height > -1 else height
width = width // self.vae_scale_factor if width > -1 else width
latent_channels = self.vae_decoder.config.get("latent_channels", None)
if latent_channels is None:
latent_channels = model.inputs[0].get_partial_shape()[1]
if latent_channels.is_dynamic:
logger.warning(
"Could not identify `latent_channels` from the VAE decoder configuration, to statically reshape the VAE decoder please provide a configuration."
)
self.is_dynamic = True
shapes = {model.inputs[0]: [num_images_per_prompt, latent_channels, height, width]}
model.reshape(shapes)
return model
def reshape(
self,
batch_size: int,
height: int,
width: int,
num_images_per_prompt: int = -1,
):
if self._compile_only:
raise ValueError(
"`reshape()` is not supported with `compile_only` mode, please intialize model without this option"
)
self.is_dynamic = -1 in {batch_size, height, width, num_images_per_prompt}
if self.tokenizer is None and self.tokenizer_2 is None:
tokenizer_max_len = -1
else:
tokenizer_max_len = (
self.tokenizer.model_max_length if self.tokenizer is not None else self.tokenizer_2.model_max_length
)
if self.unet is not None:
self.unet.model = self._reshape_unet(
self.unet.model, batch_size, height, width, num_images_per_prompt, tokenizer_max_len
)
if self.transformer is not None:
self.transformer.model = self._reshape_transformer(
self.transformer.model, batch_size, height, width, num_images_per_prompt, tokenizer_max_len
)
self.vae_decoder.model = self._reshape_vae_decoder(
self.vae_decoder.model, height, width, num_images_per_prompt
)
if self.vae_encoder is not None:
self.vae_encoder.model = self._reshape_vae_encoder(self.vae_encoder.model, batch_size, height, width)
if self.text_encoder is not None:
self.text_encoder.model = self._reshape_text_encoder(
self.text_encoder.model, batch_size, self.tokenizer.model_max_length
)
if self.text_encoder_2 is not None:
self.text_encoder_2.model = self._reshape_text_encoder(
self.text_encoder_2.model, batch_size, self.tokenizer_2.model_max_length
)
if self.text_encoder_3 is not None:
self.text_encoder_3.model = self._reshape_text_encoder(
self.text_encoder_3.model, batch_size, self.tokenizer_3.model_max_length
)
self.clear_requests()
return self
def half(self):
"""
Converts all the model weights to FP16 for more efficient inference on GPU.
"""
if self._compile_only:
raise ValueError(
"`half()` is not supported with `compile_only` mode, please intialize model without this option"
)
for component in {
self.unet,
self.transformer,
self.vae_encoder,
self.vae_decoder,
self.text_encoder,
self.text_encoder_2,
self.text_encoder_3,
}:
if component is not None:
compress_model_transformation(component.model)
self.clear_requests()
return self
def clear_requests(self):
if self._compile_only:
raise ValueError(
"`clear_requests()` is not supported with `compile_only` mode, please intialize model without this option"
)
for component in [
self.unet,
self.transformer,
self.vae_encoder,
self.vae_decoder,
self.text_encoder,
self.text_encoder_2,
self.text_encoder_3,
]:
if component is not None:
component.request = None
def compile(self):
for component in [
self.unet,
self.transformer,
self.vae_encoder,
self.vae_decoder,
self.text_encoder,
self.text_encoder_2,
self.text_encoder_3,
]:
if component is not None:
component._compile()
@classmethod
def _load_config(cls, config_name_or_path: Union[str, os.PathLike], **kwargs):
return cls.load_config(config_name_or_path, **kwargs)
@property
def components(self) -> Dict[str, Any]:
components = {
"vae": self.vae,
"unet": self.unet,
"transformer": self.transformer,
"text_encoder": self.text_encoder,
"text_encoder_2": self.text_encoder_2,
"text_encoder_3": self.text_encoder_2,
"safety_checker": self.safety_checker,
"image_encoder": self.image_encoder,
}
components = {k: v for k, v in components.items() if v is not None}
return components
def __call__(self, *args, **kwargs):
# we do this to keep numpy random states support for now
# TODO: deprecate and add warnings when a random state is passed
args = list(args)
for i in range(len(args)):
args[i] = np_to_pt_generators(args[i], self.device)
for k, v in kwargs.items():
kwargs[k] = np_to_pt_generators(v, self.device)
# we use auto_model_class.__call__ here because we can't call super().__call__
# as OptimizedModel already defines a __call__ which is the first in the MRO
return self.auto_model_class.__call__(self, *args, **kwargs)
class OVPipelinePart(ConfigMixin):
config_name: str = CONFIG_NAME
def __init__(
self,
model: openvino.runtime.Model,
parent_pipeline: OVDiffusionPipeline,
model_name: str = "",
):
self.model = model
self.model_name = model_name
self.parent_pipeline = parent_pipeline
self.request = None if not parent_pipeline._compile_only else self.model
self.ov_config = parent_pipeline.ov_config
if isinstance(parent_pipeline.model_save_dir, TemporaryDirectory):
self.model_save_dir = Path(parent_pipeline.model_save_dir.name) / self.model_name
else:
self.model_save_dir = Path(parent_pipeline.model_save_dir) / self.model_name
config_file_path = self.model_save_dir / self.config_name
if not config_file_path.is_file():
# config is mandatory for the model part to be used for inference
raise ValueError(f"Configuration file for {self.__class__.__name__} not found at {config_file_path}")
config_dict = self._dict_from_json_file(config_file_path)
self.register_to_config(**config_dict)
@property
def _device(self) -> str:
return self.parent_pipeline._device
@property
def device(self) -> torch.device:
return self.parent_pipeline.device
@property
def dtype(self) -> torch.dtype:
return OV_TO_PT_TYPE[self.ov_config.get("dtype", "f32")]
def _compile(self):
if self.request is None:
if (
"CACHE_DIR" not in self.ov_config.keys()
and not str(self.model_save_dir).startswith(gettempdir())
and "GPU" in self._device
):
self.ov_config["CACHE_DIR"] = os.path.join(self.model_save_dir, "model_cache")
logger.info(f"Compiling the {self.model_name} to {self._device} ...")
self.request = core.compile_model(self.model, self._device, self.ov_config)
# OPENVINO_LOG_LEVEL can be found in https://docs.openvino.ai/2023.2/openvino_docs_OV_UG_supported_plugins_AUTO_debugging.html