diff --git a/modules/sd_models.py b/modules/sd_models.py index 91fc06168..27791eda6 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1058,32 +1058,6 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No elif "Kandinsky" in sd_model.__class__.__name__: sd_model.scheduler.name = 'DDIM' - base_sent_to_cpu=False - if (shared.opts.cuda_compile and shared.opts.cuda_compile_backend != 'none') or shared.opts.ipex_optimize or shared.opts.nncf_compress_weights: - if op == 'refiner' and not getattr(sd_model, 'has_accelerate', False): - gpu_vram = memory_stats().get('gpu', {}) - free_vram = gpu_vram.get('total', 0) - gpu_vram.get('used', 0) - refiner_enough_vram = free_vram >= 7 if "StableDiffusionXL" in sd_model.__class__.__name__ else 3 - if not shared.opts.diffusers_move_base and refiner_enough_vram: - move_model(sd_model, devices.device) - base_sent_to_cpu=False - else: - if not refiner_enough_vram and not (shared.opts.diffusers_move_base and shared.opts.diffusers_move_refiner): - shared.log.warning(f"Insufficient GPU memory, using system memory as fallback: free={free_vram} GB") - if not shared.opts.shared.opts.diffusers_seq_cpu_offload and not shared.opts.diffusers_model_cpu_offload: - shared.log.debug('Enabled moving base model to CPU') - shared.log.debug('Enabled moving refiner model to CPU') - shared.opts.diffusers_move_base=True - shared.opts.diffusers_move_refiner=True - shared.log.debug('Moving base model to CPU') - move_model(model_data.sd_model, devices.cpu) - devices.torch_gc(force=True) - move_model(sd_model, devices.device) - base_sent_to_cpu=True - else: - move_model(sd_model, devices.device) - sd_models_compile.compile_diffusers(sd_model) - if sd_model is None: shared.log.error('Diffuser model not loaded') return @@ -1097,15 +1071,24 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256 if hasattr(sd_model, "set_progress_bar_config"): sd_model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining}', ncols=80, colour='#327fba') + + set_diffuser_options(sd_model, vae, op) + if op == 'refiner' and shared.opts.diffusers_move_refiner: shared.log.debug('Moving refiner model to CPU') move_model(sd_model, devices.cpu) else: move_model(sd_model, devices.device) - if op == 'refiner' and base_sent_to_cpu: - shared.log.debug('Moving base model back to GPU') - move_model(model_data.sd_model, devices.device) - set_diffuser_options(sd_model, vae, op) # offloading should enabled after all + + if shared.opts.ipex_optimize: + sd_model = sd_models_compile.ipex_optimize(sd_model) + + if shared.opts.nncf_compress_weights and not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"): + sd_model = sd_models_compile.nncf_compress_weights(sd_model) + + if (shared.opts.cuda_compile and shared.opts.cuda_compile_backend != 'none'): + sd_model = sd_models_compile.compile_diffusers(sd_model) + except Exception as e: shared.log.error("Failed to load diffusers model") errors.display(e, "loading Diffusers model") diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index bc114fd18..87ff44978 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -35,7 +35,20 @@ def ipex_optimize(sd_model): import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import model.eval() model.training = False - model = ipex.optimize(model, dtype=devices.dtype, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init + if model.device.type != "meta": + return_device = model.device + model = ipex.optimize(model.to(devices.device), + dtype=devices.dtype, + inplace=True, + weights_prepack=False + ).to(return_device) # pylint: disable=attribute-defined-outside-init + else: + model = ipex.optimize(model, + dtype=devices.dtype, + inplace=True, + weights_prepack=False + ) # pylint: disable=attribute-defined-outside-init + devices.torch_gc() return model if "Model" in shared.opts.ipex_optimize: @@ -79,9 +92,10 @@ def nncf_compress_weights(sd_model): def nncf_compress_model(model): return_device = model.device + model.eval() if hasattr(model, "get_input_embeddings"): backup_embeddings = copy.deepcopy(model.get_input_embeddings()) - model = nncf.compress_weights(model.eval().to(devices.device)).to(return_device) + model = nncf.compress_weights(model.to(devices.device)).to(return_device) if hasattr(model, "set_input_embeddings"): model.set_input_embeddings(backup_embeddings) devices.torch_gc(force=True) @@ -188,7 +202,21 @@ def compile_torch(sd_model): shared.log.debug(f"Model compile available backends: {torch._dynamo.list_backends()}") # pylint: disable=protected-access def torch_compile_model(model): - return torch.compile(model, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph) + if model.device.type != "meta": + return_device = model.device + model = torch.compile(model.to(devices.device), + mode=shared.opts.cuda_compile_mode, + backend=shared.opts.cuda_compile_backend, + fullgraph=shared.opts.cuda_compile_fullgraph + ).to(return_device) + else: + model = torch.compile(model, + mode=shared.opts.cuda_compile_mode, + backend=shared.opts.cuda_compile_backend, + fullgraph=shared.opts.cuda_compile_fullgraph + ) + devices.torch_gc() + return model if shared.opts.cuda_compile_backend == "openvino_fx": sd_model = optimize_openvino(sd_model) @@ -273,10 +301,6 @@ def compile_deepcache(sd_model): def compile_diffusers(sd_model): - if shared.opts.ipex_optimize: - sd_model = ipex_optimize(sd_model) - if shared.opts.nncf_compress_weights and not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"): - sd_model = nncf_compress_weights(sd_model) if not shared.opts.cuda_compile: return sd_model if shared.opts.cuda_compile_backend == 'none':