OpenVINO add experimental support for Text Encoder compiling

This commit is contained in:
Disty0
2024-01-06 01:37:04 +03:00
parent 3b14e1e70e
commit f3cebcb9dd
6 changed files with 55 additions and 23 deletions
+12 -4
View File
@@ -21,7 +21,8 @@ class CompiledModelState:
self.lora_compile = False
self.compiled_cache = {}
self.partitioned_modules = {}
self.compiling_vae = False
self.subgraph_type = []
self.compile_dont_use_4bit = False
def ipex_optimize(sd_model):
@@ -149,10 +150,10 @@ def compile_torch(sd_model):
t0 = time.time()
if shared.opts.cuda_compile:
if shared.opts.cuda_compile and (not hasattr(sd_model, 'unet') or not hasattr(sd_model.unet, 'config')):
shared.log.warning('Model compile enabled but model has no Unet')
else:
if hasattr(sd_model, 'unet') and hasattr(sd_model.unet, 'config'):
sd_model.unet = torch.compile(sd_model.unet, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
else:
shared.log.warning('Model compile enabled but model has no Unet')
if shared.opts.cuda_compile_vae:
if hasattr(sd_model, 'vae') and hasattr(sd_model.vae, 'decode'):
sd_model.vae.decode = torch.compile(sd_model.vae.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
@@ -160,6 +161,13 @@ def compile_torch(sd_model):
sd_model.movq.decode = torch.compile(sd_model.movq.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
else:
shared.log.warning('Model compile enabled but model has no VAE')
if shared.opts.cuda_compile_text_encoder:
if hasattr(sd_model, 'text_encoder'):
sd_model.text_encoder = torch.compile(sd_model.text_encoder, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
if hasattr(sd_model, 'text_encoder_2'):
sd_model.text_encoder_2 = torch.compile(sd_model.text_encoder_2, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
else:
shared.log.warning('Model compile enabled but model has no Unet')
setup_logging() # compile messes with logging so reset is needed
if shared.opts.cuda_compile_precompile:
sd_model("dummy prompt")