diff --git a/installer.py b/installer.py index 1f853cc24..ae227723d 100644 --- a/installer.py +++ b/installer.py @@ -429,8 +429,8 @@ def check_torch(): torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.0.0a0 intel_extension_for_pytorch==2.0.110+gitba7f6c1 -f https://developer.intel.com/ipex-whl-stable-xpu') elif allow_openvino and args.use_openvino: #Remove this after 2.1.0 releases - log.info('Using OpenVINO with Torch Nightly CPU') - torch_command = os.environ.get('TORCH_COMMAND', '--pre torch==2.1.0.dev20230713+cpu torchvision==0.16.0.dev20230713+cpu -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html') + log.info('Using OpenVINO') + torch_command = os.environ.get('TORCH_COMMAND', '--pre torch==2.1.0.dev20230726+cpu torchvision==0.16.0.dev20230726+cpu -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html') else: machine = platform.machine() if sys.platform == 'darwin': diff --git a/modules/intel/openvino/__init__.py b/modules/intel/openvino/__init__.py index b6bcec4ab..6162b674b 100644 --- a/modules/intel/openvino/__init__.py +++ b/modules/intel/openvino/__init__.py @@ -7,24 +7,20 @@ from torch._dynamo.backends.common import fake_tensor_unsupported from torch._dynamo.backends.registry import register_backend from torch.fx.experimental.proxy_tensor import make_fx from torch._inductor.compile_fx import compile_fx +from torch.utils._pytree import tree_flatten from hashlib import sha256 +import functools from modules import shared, devices -@register_backend -@fake_tensor_unsupported -def openvino_fx(subgraph, example_inputs): - executor_parameters = None +def openvino_clear_caches(): + global partitioned_modules + global compiled_cache + + compiled_cache.clear() + partitioned_modules.clear() + +def get_device(): core = Core() - if os.getenv("OPENVINO_TORCH_MODEL_CACHING") != "0": - os.environ.setdefault('OPENVINO_TORCH_MODEL_CACHING', "1") - model_hash_str = sha256(subgraph.code.encode('utf-8')).hexdigest() - executor_parameters = {"model_hash_str": model_hash_str} - - example_inputs.reverse() - cache_root = "./cache/" - if os.getenv("OPENVINO_TORCH_CACHE_DIR") is not None: - cache_root = os.getenv("OPENVINO_TORCH_CACHE_DIR") - if os.getenv("OPENVINO_TORCH_BACKEND_DEVICE") is not None: device = os.getenv("OPENVINO_TORCH_BACKEND_DEVICE") elif any(openvino_cpu in cpu_module.lower() for cpu_module in shared.cmd_opts.use_cpu for openvino_cpu in ["openvino", "all"]): @@ -43,83 +39,169 @@ def openvino_fx(subgraph, example_inputs): os.environ.setdefault('OPENVINO_TORCH_BACKEND_DEVICE', device) shared.log.debug(f"OpenVINO Device: {device}") if shared.opts.cuda_compile_errors and device not in core.available_devices: - shared.log.warning(f"OpenVINO: Specified device {device} is not in the list of OpenVINO Available Devices") + shared.log.error(f"OpenVINO: Specified device {device} is not in the list of OpenVINO Available Devices") assert device in core.available_devices, f"OpenVINO: Specified device {device} is not in the list of OpenVINO Available Devices" - #Cache saving keeps increasing the partition id - #This loop check if non 0 partition id caches exist - #Takes 0.002 seconds when nothing is found - use_cached_file = False - for i in range(100): - file_name = get_cached_file_name(*example_inputs, model_hash_str=str(model_hash_str + str(i)), device=device, cache_root=cache_root) - if file_name is not None and os.path.isfile(file_name + ".xml") and os.path.isfile(file_name + ".bin"): - use_cached_file = True - break + return device - if use_cached_file: - om = core.read_model(file_name + ".xml") +def cache_root_path(): + cache_root = "./cache/" + if os.getenv("OPENVINO_TORCH_CACHE_DIR") is not None: + cache_root = os.getenv("OPENVINO_TORCH_CACHE_DIR") + return cache_root - dtype_mapping = { - torch.float32: Type.f32, - torch.float64: Type.f64, - torch.float16: Type.f16, - torch.int64: Type.i64, - torch.int32: Type.i32, - torch.uint8: Type.u8, - torch.int8: Type.i8, - torch.bool: Type.boolean - } +def cached_model_name(model_hash_str, device, args, cache_root, reversed = False): + if model_hash_str is None: + return None - for idx, input_data in enumerate(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() + model_cache_dir = cache_root + "/model/" - if model_hash_str is not None: - core.set_property({'CACHE_DIR': cache_root + '/blob'}) + try: + os.makedirs(model_cache_dir, exist_ok=True) + file_name = model_cache_dir + model_hash_str + "_" + device + except OSError as error: + shared.log.error(f"Cache directory {cache_root} cannot be created. Model caching is disabled. Error: {error}") + return None - compiled_model = core.compile_model(om, device) - def _call(*args): - ov_inputs = [a.detach().cpu().numpy() for a in args] - ov_inputs.reverse() - res = compiled_model(ov_inputs) - result = [torch.from_numpy(res[out]) for out in compiled_model.outputs] - return result - return _call - else: - example_inputs.reverse() - model = make_fx(subgraph)(*example_inputs) - with devices.inference_context(): - model.eval() - partitioner = Partitioner() - compiled_model = partitioner.make_partitions(model) + inputs_str = "" + for input_data in args: + if reversed: + inputs_str = "_" + str(input_data.type()) + str(input_data.size())[11:-1].replace(" ", "") + inputs_str + else: + inputs_str += "_" + str(input_data.type()) + str(input_data.size())[11:-1].replace(" ", "") + inputs_str = sha256(inputs_str.encode('utf-8')).hexdigest() + file_name += inputs_str - def _call(*args): - res = execute(compiled_model, *args, executor="openvino", - executor_parameters=executor_parameters) - return res - return _call - - -def get_cached_file_name(*args, model_hash_str, device, cache_root): - file_name = None - if model_hash_str is not None: - model_cache_dir = cache_root + "/model/" - try: - os.makedirs(model_cache_dir, exist_ok=True) - file_name = model_cache_dir + model_hash_str + "_" + device - for input_data in args: - if file_name is not None: - file_name += "_" + str(input_data.type()) + str(input_data.size())[11:-1].replace(" ", "") - except OSError as error: - print("Cache directory ", cache_root, " cannot be created. Model caching is disabled. Error: ", error) - file_name = None - model_hash_str = None return file_name -def openvino_clear_caches(): - global partitioned_modules - global compiled_cache +def openvino_compile_cached_model(cached_model_path, *example_inputs): + core = Core() + om = core.read_model(cached_model_path + ".xml") - compiled_cache.clear() - partitioned_modules.clear() + dtype_mapping = { + torch.float32: Type.f32, + torch.float64: Type.f64, + torch.float16: Type.f16, + torch.int64: Type.i64, + torch.int32: Type.i32, + torch.uint8: Type.u8, + torch.int8: Type.i8, + torch.bool: Type.boolean + } + + for idx, input_data in enumerate(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() + + core.set_property({'CACHE_DIR': cache_root_path() + '/blob'}) + + compiled_model = core.compile_model(om, get_device()) + + return compiled_model + +def execute_cached(compiled_model, *args): + model_state = shared.compiled_model_state + flat_args, _ = tree_flatten(args) + ov_inputs = [a.detach().cpu().numpy() for a in flat_args] + + if (model_state.cn_model == "None"): + ov_inputs.reverse() + + res = compiled_model(ov_inputs) + result = [torch.from_numpy(res[out]) for out in compiled_model.outputs] + return result + +def check_fully_supported(self, graph_module): + num_fused = 0 + for node in graph_module.graph.nodes: + if node.op == "call_module" and "fused_" in node.name: + num_fused += 1 + elif node.op != "placeholder" and node.op != "output": + return False + if num_fused == 1: + return True + return False + +Partitioner.check_fully_supported = functools.partial(check_fully_supported, Partitioner) + +@register_backend +@fake_tensor_unsupported +def openvino_fx(subgraph, example_inputs): + model_state = shared.compiled_model_state + executor_parameters = None + inputs_reversed = False + if os.getenv("OPENVINO_TORCH_MODEL_CACHING") != "0": + os.environ.setdefault('OPENVINO_TORCH_MODEL_CACHING', "1") + # Create a hash to be used for caching + model_hash_str = sha256(subgraph.code.encode('utf-8')).hexdigest() + if (model_state.cn_model != "None" and model_state.partition_id == 0): + model_hash_str = model_hash_str + model_state.cn_model + + if (model_state.lora_model != "None"): + model_hash_str = model_hash_str + model_state.lora_model + + executor_parameters = {"model_hash_str": model_hash_str} + # Check if the model was fully supported and already cached + example_inputs.reverse() + inputs_reversed = True + maybe_fs_cached_name = cached_model_name(model_hash_str + "_fs", get_device(), example_inputs, cache_root_path()) + + if os.path.isfile(maybe_fs_cached_name + ".xml") and os.path.isfile(maybe_fs_cached_name + ".bin"): + if (model_state.cn_model != "None" and model_state.cn_model in maybe_fs_cached_name): + example_inputs_reordered = [] + if (os.path.isfile(maybe_fs_cached_name + ".txt")): + f = open(maybe_fs_cached_name + ".txt", "r") + for input_data in example_inputs: + shape = f.readline() + if (str(input_data.size()) != shape): + for idx1, input_data1 in enumerate(example_inputs): + if (str(input_data1.size()).strip() == str(shape).strip()): + example_inputs_reordered.append(example_inputs[idx1]) + example_inputs = example_inputs_reordered + + # 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) + + def _call(*args): + if (model_state.cn_model != "None" and model_state.cn_model in maybe_fs_cached_name): + args_reordered = [] + if (os.path.isfile(maybe_fs_cached_name + ".txt")): + f = open(maybe_fs_cached_name + ".txt", "r") + for input_data in args: + shape = f.readline() + if (str(input_data.size()) != shape): + for idx1, input_data1 in enumerate(args): + if (str(input_data1.size()).strip() == str(shape).strip()): + args_reordered.append(args[idx1]) + args = args_reordered + + res = execute_cached(compiled_model, *args) + model_state.partition_id = model_state.partition_id + 1 + return res + return _call + else: + maybe_fs_cached_name = None + + if inputs_reversed: + example_inputs.reverse() + model = make_fx(subgraph)(*example_inputs) + for node in model.graph.nodes: + if node.target == torch.ops.aten.mul_.Tensor: + node.target = torch.ops.aten.mul.Tensor + with torch.no_grad(): + model.eval() + partitioner = Partitioner() + compiled_model = partitioner.make_partitions(model) + + if executor_parameters is not None and 'model_hash_str' in executor_parameters: + # Check if the model is fully supported. + fully_supported = partitioner.check_fully_supported(compiled_model) + if fully_supported: + executor_parameters["model_hash_str"] += "_fs" + + def _call(*args): + res = execute(compiled_model, *args, executor="openvino", + executor_parameters=executor_parameters) #, file_name=maybe_fs_cached_name) + return res + return _call diff --git a/modules/lora_diffusers.py b/modules/lora_diffusers.py index 82d7a10b0..72af05722 100644 --- a/modules/lora_diffusers.py +++ b/modules/lora_diffusers.py @@ -35,6 +35,8 @@ def unload_diffusers_lora(): lora_state['loaded'].clear() lora_state['all_loras'] = [] lora_state['multiplier'] = [] + if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx": + shared.compiled_model_state.lora_model = "None" except Exception as e: shared.log.error(f"LoRA unload failed: {e}") @@ -74,6 +76,8 @@ def load_diffusers_lora(name, lora, strength = 1.0, num_loras = 1): lora_state['loaded'].append(f'{lora.filename}:{strength}') t1 = time.time() fuse = f'fuse={fuse:.2f}s' if fuse > 0 else '' + if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx": + shared.compiled_model_state.lora_model = str(lora_state['loaded']) shared.log.info(f'LoRA loaded: {name} strength={strength} loader="{shared.opts.diffusers_lora_loader}" lora={t1-t0:.2f}s {fuse}') except Exception as e: lines = str(e).splitlines() diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index d913c4291..6b03640f3 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -260,9 +260,9 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro if shared.opts.cuda_compile_backend == "openvino_fx": compile_height = p.height if not hires else p.hr_upscale_to_y compile_width = p.width if not hires else p.hr_upscale_to_x - if (not hasattr(shared.sd_model, "compiled_model_state") or (not shared.sd_model.compiled_model_state.first_pass - and (shared.sd_model.compiled_model_state.height != compile_height or shared.sd_model.compiled_model_state.width != compile_width - or shared.sd_model.compiled_model_state.batch_size != p.batch_size))): + if (shared.compiled_model_state is None or (not shared.compiled_model_state.first_pass + and (shared.compiled_model_state.height != compile_height or shared.compiled_model_state.width != compile_width + or shared.compiled_model_state.batch_size != p.batch_size))): shared.log.info("OpenVINO: Resolution change detected") shared.log.info("OpenVINO: Recompiling base model") sd_models.unload_model_weights(op='model') @@ -271,15 +271,17 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro shared.log.info("OpenVINO: Recompiling refiner") sd_models.unload_model_weights(op='refiner') sd_models.reload_model_weights(op='refiner') - shared.sd_model.compiled_model_state.height = compile_height - shared.sd_model.compiled_model_state.width = compile_width - shared.sd_model.compiled_model_state.batch_size = p.batch_size - shared.sd_model.compiled_model_state.first_pass = False + shared.compiled_model_state.height = compile_height + shared.compiled_model_state.width = compile_width + shared.compiled_model_state.batch_size = p.batch_size + shared.compiled_model_state.first_pass = False else: pass #Can be implemented for TensorRT or Olive else: pass #Do nothing if compile is disabled + recompile_model() + is_karras_compatible = shared.sd_model.__class__.__init__.__annotations__.get("scheduler", None) == diffusers.schedulers.scheduling_utils.KarrasDiffusionSchedulers if (not hasattr(shared.sd_model.scheduler, 'name')) or (shared.sd_model.scheduler.name != p.sampler_name) and (p.sampler_name != 'Default') and is_karras_compatible: sampler = sd_samplers.all_samplers_map.get(p.sampler_name, None) @@ -316,8 +318,6 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro unload_diffusers_lora() return results - recompile_model() - if shared.opts.diffusers_move_base and not shared.sd_model.has_accelerate: shared.sd_model.to(devices.device) @@ -379,12 +379,12 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro output.images = hires_resize(latents=output.images) if latent_scale_mode is not None or p.hr_force: p.ops.append('hires') + recompile_model(hires=True) if (not hasattr(shared.sd_model.scheduler, 'name')) or (shared.sd_model.scheduler.name != p.latent_sampler) and (p.latent_sampler != 'Default') and is_karras_compatible: sampler = sd_samplers.all_samplers_map.get(p.latent_sampler, None) if sampler is None: sampler = sd_samplers.all_samplers_map.get("UniPC") sd_samplers.create_sampler(sampler.name, shared.sd_model) # TODO(Patrick): For wrapped pipelines this is currently a no-op - recompile_model(hires=True) sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE) hires_args = set_pipeline_args( model=shared.sd_model, diff --git a/modules/sd_models.py b/modules/sd_models.py index b9acd87a4..f27068914 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -118,10 +118,13 @@ class CheckpointInfo: #Used by OpenVINO, can be used with TensorRT or Olive class CompiledModelState: def __init__(self): + self.first_pass = True self.height = 512 self.width = 512 self.batch_size = 1 - self.first_pass = True + self.partition_id = 0 + self.cn_model = "None" + self.lora_model = "None" class NoWatermark: @@ -888,8 +891,8 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No from modules.intel.openvino import openvino_fx, openvino_clear_caches # pylint: disable=unused-import openvino_clear_caches() torch._dynamo.eval_frame.check_if_dynamo_supported = lambda: True # pylint: disable=protected-access - sd_model.compiled_model_state = CompiledModelState() - sd_model.compiled_model_state.first_pass = True if not shared.opts.cuda_compile_precompile else False + shared.compiled_model_state = CompiledModelState() + shared.compiled_model_state.first_pass = True if not shared.opts.cuda_compile_precompile else False log_level = logging.WARNING if shared.opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access if hasattr(torch, '_logging'): torch._logging.set_logs(dynamo=log_level, aot=log_level, inductor=log_level) # pylint: disable=protected-access diff --git a/modules/shared.py b/modules/shared.py index bc8f896b5..616ea76d9 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -1048,4 +1048,5 @@ sd_model = None sd_refiner = None sd_model_type = '' sd_refiner_type = '' +compiled_model_state = None sys.modules[__name__].__class__ = Shared