diff --git a/CHANGELOG.md b/CHANGELOG.md index 800b6fece..2ef93e289 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -67,6 +67,7 @@ And it also includes fixes for all reported issues so far requires nightly versions of `torch` and `torchao` > pip install -U --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121 > pip install -U git+https://github.com/pytorch-labs/ao + - new option: **compile text encoder** (experimental) - **IPEX**, thanks @disty0 - rewrite ipex hijacks without CondFunc improves compatibilty and performance @@ -77,6 +78,8 @@ And it also includes fixes for all reported issues so far - **4-bit support with NNCF** enable *Compress Model weights with NNCF* from *Compute Settings* and set a 4-bit NNCF mode 4-bit and 8-bit with OpenVINO is CPU only for now + - experimental support for *Text Encoder* compiling + OpenVINO is faster than IPEX now - reduce system memory usage after compile - fix cache loading with multiple models - **Fixes** diff --git a/modules/intel/openvino/__init__.py b/modules/intel/openvino/__init__.py index e5440cc69..6984db16f 100644 --- a/modules/intel/openvino/__init__.py +++ b/modules/intel/openvino/__init__.py @@ -229,8 +229,8 @@ def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None, file_na om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype]) om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape))) om.validate_nodes_and_infer_types() - if shared.opts.nncf_compress_weights and not (shared.compiled_model_state.compiling_vae and not shared.opts.nncf_compress_vae_weights): - if shared.compiled_model_state.compiling_vae or shared.opts.nncf_compress_weights_mode == "INT8": + if shared.opts.nncf_compress_weights and not (shared.compiled_model_state.compile_dont_use_4bit and not shared.opts.nncf_compress_vae_weights): + if shared.compiled_model_state.compile_dont_use_4bit or shared.opts.nncf_compress_weights_mode == "INT8": om = nncf.compress_weights(om) else: om = nncf.compress_weights(om, mode=getattr(nncf.CompressWeightsMode, shared.opts.nncf_compress_weights_mode), group_size=8, ratio=shared.opts.nncf_compress_weights_raito) @@ -238,7 +238,7 @@ def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None, file_na if model_hash_str is not None: core.set_property({'CACHE_DIR': cache_root + '/blob'}) - shared.compiled_model_state.compiling_vae = False + shared.compiled_model_state.compile_dont_use_4bit = False compiled_model = core.compile_model(om, device) return compiled_model @@ -261,15 +261,15 @@ def openvino_compile_cached_model(cached_model_path, *example_inputs): om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype]) om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape))) om.validate_nodes_and_infer_types() - if shared.opts.nncf_compress_weights and not (shared.compiled_model_state.compiling_vae and not shared.opts.nncf_compress_vae_weights): - if shared.compiled_model_state.compiling_vae or shared.opts.nncf_compress_weights_mode == "INT8": + if shared.opts.nncf_compress_weights and not (shared.compiled_model_state.compile_dont_use_4bit and not shared.opts.nncf_compress_vae_weights): + if shared.compiled_model_state.compile_dont_use_4bit or shared.opts.nncf_compress_weights_mode == "INT8": om = nncf.compress_weights(om) else: om = nncf.compress_weights(om, mode=getattr(nncf.CompressWeightsMode, shared.opts.nncf_compress_weights_mode), group_size=8, ratio=shared.opts.nncf_compress_weights_raito) core.set_property({'CACHE_DIR': shared.opts.openvino_cache_path + '/blob'}) - shared.compiled_model_state.compiling_vae = False + shared.compiled_model_state.compile_dont_use_4bit = False compiled_model = core.compile_model(om, get_device()) return compiled_model @@ -345,6 +345,7 @@ def partition_graph(gm: GraphModule, use_python_fusion_cache: bool, model_hash_s def generate_subgraph_str(tensor): if hasattr(tensor, "weight"): shared.compiled_model_state.model_str = shared.compiled_model_state.model_str + str(tensor.weight) + shared.compiled_model_state.subgraph_type.append(type(tensor)) return tensor @register_backend @@ -353,11 +354,22 @@ def openvino_fx(subgraph, example_inputs): executor_parameters = None inputs_reversed = False maybe_fs_cached_name = None + + shared.compiled_model_state.model_str = "" + shared.compiled_model_state.subgraph_type = [] + subgraph.apply(generate_subgraph_str) + + # SD 1.5 / SDXL VAE + if (shared.compiled_model_state.subgraph_type[0] is torch.nn.modules.conv.Conv2d and + shared.compiled_model_state.subgraph_type[1] is torch.nn.modules.conv.Conv2d and + shared.compiled_model_state.subgraph_type[2] is torch.nn.modules.normalization.GroupNorm and + shared.compiled_model_state.subgraph_type[3] is torch.nn.modules.activation.SiLU): + + shared.compiled_model_state.compile_dont_use_4bit = True + if not shared.opts.openvino_disable_model_caching: os.environ.setdefault('OPENVINO_TORCH_MODEL_CACHING', "1") # Create a hash to be used for caching - shared.compiled_model_state.model_str = "" - subgraph.apply(generate_subgraph_str) model_hash_str = sha256(shared.compiled_model_state.model_str.encode('utf-8')).hexdigest() if (shared.compiled_model_state.cn_model != [] and shared.compiled_model_state.partition_id == 0): model_hash_str = model_hash_str + str(shared.compiled_model_state.cn_model) @@ -383,9 +395,18 @@ def openvino_fx(subgraph, example_inputs): example_inputs_reordered.append(example_inputs[idx1]) example_inputs = example_inputs_reordered - # Delete unused subgraphs - subgraph = subgraph.apply(sd_models.convert_to_faketensors) - devices.torch_gc(force=True) + # SD 1.5 / SDXL Text Encoder + if (shared.compiled_model_state.subgraph_type[0] is torch.nn.modules.sparse.Embedding and + shared.compiled_model_state.subgraph_type[1] is torch.nn.modules.sparse.Embedding and + shared.compiled_model_state.subgraph_type[2] is torch.nn.modules.normalization.LayerNorm and + shared.compiled_model_state.subgraph_type[3] is torch.nn.modules.linear.Linear): + + pass # Fails with FakeTensors or Downcast + else: + # Delete unused subgraphs + print(shared.compiled_model_state.subgraph_type) + subgraph = subgraph.apply(sd_models.convert_to_faketensors) + devices.torch_gc(force=True) # Model is fully supported and already cached. Run the cached OV model directly. compiled_model = openvino_compile_cached_model(maybe_fs_cached_name, *example_inputs) diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 7b5fba688..8462e4a8d 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -48,13 +48,9 @@ def full_vae_decode(latents, model): model.upcast_vae() latents = latents.to(next(iter(model.vae.post_quant_conv.parameters())).dtype) - # OpenVINO with INT4 doesn't work with VAE decode so we pass that we are using VAE right now to OpenVINO - if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx" and shared.compiled_model_state.first_pass_vae: - shared.compiled_model_state.compiling_vae = True - decoded = model.vae.decode(latents / model.vae.config.scaling_factor, return_dict=False)[0] - # Downcast VAE after OpenVINO compile + # Delete PyTorch VAE after OpenVINO compile if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx" and shared.compiled_model_state.first_pass_vae: shared.compiled_model_state.first_pass_vae = False if hasattr(shared.sd_model, "vae"): diff --git a/modules/sd_models.py b/modules/sd_models.py index e6cbf8235..4267f4cbc 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1263,9 +1263,9 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model') def convert_to_faketensors(tensor): - fake = torch._subclasses.fake_tensor.FakeTensorMode() + fake_module = torch._subclasses.fake_tensor.FakeTensorMode(allow_non_fake_inputs=True) if hasattr(tensor, "weight"): - tensor.weight = torch.nn.Parameter(fake.from_tensor(tensor.weight)) + tensor.weight = torch.nn.Parameter(fake_module.from_tensor(tensor.weight)) return tensor @@ -1280,6 +1280,9 @@ def disable_offload(sd_model): def unload_model_weights(op='model', change_from='none'): + if shared.compiled_model_state is not None: + shared.compiled_model_state.compiled_cache.clear() + shared.compiled_model_state.partitioned_modules.clear() if op == 'model' or op == 'dict': if model_data.sd_model: if (shared.backend == shared.Backend.ORIGINAL and change_from != shared.Backend.DIFFUSERS) or change_from == shared.Backend.ORIGINAL: diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 3a5742bea..1199be162 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -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") diff --git a/modules/shared.py b/modules/shared.py index a40338671..93b317c30 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -327,7 +327,8 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "cuda_compile_sep": OptionInfo("