OpenVINO fix cache loading

modules/intel/openvino/__init__.py

@@ -358,41 +358,41 @@ def openvino_fx(subgraph, example_inputs):
         maybe_fs_cached_name = cached_model_name(model_hash_str + "_fs", get_device(), example_inputs, shared.opts.openvino_cache_path)
 
         if os.path.isfile(maybe_fs_cached_name + ".xml") and os.path.isfile(maybe_fs_cached_name + ".bin"):
-            if (shared.compiled_model_state.cn_model != [] and str(shared.compiled_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
+            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
 
-                # Deleting unused subgraphs doesn't do anything, so we cast it down to fp8
-                subgraph = subgraph.to(dtype=torch.float8_e4m3fn)
+            # Deleting unused subgraphs doesn't do anything, so we cast it down to fp8
+            subgraph = subgraph.to(dtype=torch.float8_e4m3fn)
+            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)
+            # 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 (shared.compiled_model_state.cn_model != [] and str(shared.compiled_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
+            def _call(*args):
+                if (shared.compiled_model_state.cn_model != [] and str(shared.compiled_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)
-                    shared.compiled_model_state.partition_id = shared.compiled_model_state.partition_id + 1
-                    return res
-                return _call
+                res = execute_cached(compiled_model, *args)
+                shared.compiled_model_state.partition_id = shared.compiled_model_state.partition_id + 1
+                return res
+            return _call
     else:
         os.environ.setdefault('OPENVINO_TORCH_MODEL_CACHING', "0")
         maybe_fs_cached_name = None

modules/processing_diffusers.py

@@ -345,10 +345,20 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
                 shared.compiled_model_state.height = compile_height
                 shared.compiled_model_state.width = compile_width
                 shared.compiled_model_state.batch_size = p.batch_size
-            else:
-                pass #Can be implemented for TensorRT or Olive
-        else:
-            pass #Do nothing if compile is disabled
+
+    # Downcast UNET after OpenVINO compile
+    def downcast_openvino(op="base"):
+        if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx":
+            if shared.compiled_model_state.first_pass and op == "base":
+                shared.compiled_model_state.first_pass = False
+                if hasattr(shared.sd_model, "unet"):
+                    shared.sd_model.unet.to(dtype=torch.float8_e4m3fn)
+                    devices.torch_gc(force=True)
+            if shared.compiled_model_state.first_pass_refiner and op == "refiner":
+                shared.compiled_model_state.first_pass_refiner = False
+                if hasattr(shared.sd_refiner, "unet"):
+                    shared.sd_refiner.unet.to(dtype=torch.float8_e4m3fn)
+                    devices.torch_gc(force=True)
 
     def update_sampler(sd_model, second_pass=False):
         sampler_selection = p.latent_sampler if second_pass else p.sampler_name
@@ -448,12 +458,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
     try:
         t0 = time.time()
         output = shared.sd_model(**base_args) # pylint: disable=not-callable
-        # Downcast UNET after OpenVINO compile
-        if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx" and shared.compiled_model_state.first_pass:
-            shared.compiled_model_state.first_pass = False
-            if hasattr(shared.sd_model, "unet"):
-                shared.sd_model.unet.to(dtype=torch.float8_e4m3fn)
-                devices.torch_gc(force=True)
+        downcast_openvino(op="base")
         if shared.cmd_opts.profile:
             t1 = time.time()
             shared.log.debug(f'Profile: pipeline call: {t1-t0:.2f}')
@@ -517,12 +522,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
                 shared.state.sampling_steps = hires_args['num_inference_steps']
                 try:
                     output = shared.sd_model(**hires_args) # pylint: disable=not-callable
-                    # Downcast UNET after OpenVINO compile
-                    if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx" and shared.compiled_model_state.first_pass:
-                        shared.compiled_model_state.first_pass = False
-                        if hasattr(shared.sd_model, "unet"):
-                            shared.sd_model.unet.to(dtype=torch.float8_e4m3fn)
-                            devices.torch_gc(force=True)
+                    downcast_openvino(op="base")
                 except AssertionError as e:
                     shared.log.info(e)
                 p.init_images = []
@@ -582,12 +582,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
             try:
                 shared.sd_refiner.register_to_config(requires_aesthetics_score=shared.opts.diffusers_aesthetics_score)
                 refiner_output = shared.sd_refiner(**refiner_args) # pylint: disable=not-callable
-                # Downcast UNET after OpenVINO compile
-                if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx" and shared.compiled_model_state.first_pass_refiner:
-                    shared.compiled_model_state.first_pass_refiner = False
-                    if hasattr(shared.sd_refiner, "unet"):
-                        shared.sd_refiner.unet.to(dtype=torch.float8_e4m3fn)
-                        devices.torch_gc(force=True)
+                downcast_openvino(op="refiner")
             except AssertionError as e:
                 shared.log.info(e)
 

modules/processing_vae.py

@@ -54,7 +54,7 @@ def full_vae_decode(latents, model):
     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"):
-            shared.sd_model.vae.to(dtype=torch.float8_e4m3fn)
+            model.vae.to(dtype=torch.float8_e4m3fn)
             devices.torch_gc(force=True)
 
     if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'):