diff --git a/scripts/openvino_accelerate.py b/scripts/openvino_accelerate.py index 2bcd5002681..759d38eabb3 100644 --- a/scripts/openvino_accelerate.py +++ b/scripts/openvino_accelerate.py @@ -53,6 +53,7 @@ StableDiffusionXLImg2ImgPipeline, StableDiffusionXLInpaintPipeline, ControlNetModel, + StableDiffusionLatentUpscalePipeline, DDIMScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler, @@ -535,6 +536,23 @@ class NoWatermark: def apply_watermark(self, img): return img + +def get_diffusers_upscaler(upscaler: str): + torch._dynamo.reset() + openvino_clear_caches() + model_name = "stabilityai/sd-x2-latent-upscaler" + print("OpenVINO Script: loading upscaling model: " + model_name) + sd_model = StableDiffusionLatentUpscalePipeline.from_pretrained(model_name, torch_dtype=torch.float32) + sd_model.safety_checker = None + sd_model.cond_stage_key = functools.partial(cond_stage_key, shared.sd_model) + sd_model.unet = torch.compile(sd_model.unet, backend="openvino") + sd_model.vae.decode = torch.compile(sd_model.vae.decode, backend="openvino") + shared.sd_diffusers_model = sd_model + del sd_model + + return shared.sd_diffusers_model + + def get_diffusers_sd_model(model_config, vae_ckpt, sampler_name, enable_caching, openvino_device, mode, is_xl_ckpt, refiner_ckpt, refiner_frac): if (model_state.recompile == 1): model_state.partition_id = 0 @@ -770,7 +788,8 @@ def init_new(self, all_prompts, all_seeds, all_subseeds): else: raise RuntimeError(f"bad number of images passed: {len(imgs)}; expecting {self.batch_size} or less") -def process_images_openvino(p: StableDiffusionProcessing, model_config, vae_ckpt, sampler_name, enable_caching, openvino_device, mode, is_xl_ckpt, refiner_ckpt, refiner_frac) -> Processed: +def process_images_openvino(p: StableDiffusionProcessing, model_config, vae_ckpt, sampler_name, enable_caching, override_hires, upscaler, hires_steps, d_strength, openvino_device, mode, is_xl_ckpt, refiner_ckpt, refiner_frac) -> Processed: + """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch""" """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch""" if (mode == 0 and p.enable_hr): @@ -1092,6 +1111,23 @@ def callback(iter, t, latents): devices.torch_gc() + # Hight resolutuon mode + if override_hires: + if upscaler == "Latent": + model_state.mode = -1 + shared.sd_diffusers_model = get_diffusers_upscaler(upscaler) + img_idx = slice(len(output_images)) if p.batch_size == 1 else slice(1, len(output_images)) + output_images[img_idx] = shared.sd_diffusers_model( + image=output_images[img_idx], + prompt=p.prompts, + negative_prompt=p.negative_prompts, + num_inference_steps=hires_steps, + guidance_scale=p.cfg_scale, + generator=generator, + callback = callback, + callback_steps = 1, + ).images + res = Processed( p, images_list=output_images, @@ -1102,7 +1138,8 @@ def callback(iter, t, latents): index_of_first_image=index_of_first_image, infotexts=infotexts, ) - + if override_hires: + res.info = res.info + f", Hires upscaler: {upscaler}, Denoising strength: {d_strength}" res.info = res.info + ", Warm up time: " + str(round(warmup_duration, 2)) + " secs " if (generation_rate >= 1.0): @@ -1116,6 +1153,9 @@ def callback(iter, t, latents): return res +def on_change(mode): + return gr.update(visible=mode) + class Script(scripts.Script): def title(self): return "Accelerate with OpenVINO" @@ -1170,6 +1210,12 @@ def get_refiner_list(): override_sampler = gr.Checkbox(label="Override the sampling selection from the main UI (Recommended as only below sampling methods have been validated for OpenVINO)", value=True) sampler_name = gr.Radio(label="Select a sampling method", choices=["Euler a", "Euler", "LMS", "Heun", "DPM++ 2M", "LMS Karras", "DPM++ 2M Karras", "DDIM", "PLMS"], value="Euler a") enable_caching = gr.Checkbox(label="Cache the compiled models on disk for faster model load in subsequent launches (Recommended)", value=True, elem_id=self.elem_id("enable_caching")) + override_hires = gr.Checkbox(label="Override the Hires.fix selection from the main UI (Recommended as only below upscalers have been validated for OpenVINO)", value=False, visible=self.is_txt2img) + with gr.Group(visible=False) as hires: + with gr.Row(): + upscaler = gr.Dropdown(label="Upscaler", choices=["Latent"], value="Latent") + hires_steps = gr.Slider(1, 150, value=10, step=1, label="Steps") + d_strength = gr.Slider(0, 1, value=0.5, step=0.01, label="Strength") warmup_status = gr.Textbox(label="Device", interactive=False, visible=False) vae_status = gr.Textbox(label="VAE", interactive=False, visible=False) gr.Markdown( @@ -1184,6 +1230,8 @@ def get_refiner_list(): So it's normal for the first inference after a settings change to be slower, while subsequent inferences use the optimized compiled model and run faster. """) + override_hires.change(on_change, override_hires, hires) + def device_change(choice): if (model_state.device == choice): return gr.update(value="Device selected is " + choice, visible=True) @@ -1206,9 +1254,9 @@ def refiner_ckpt_change(choice): else: model_state.refiner_ckpt = choice refiner_ckpt.change(refiner_ckpt_change, refiner_ckpt) - return [model_config, vae_ckpt, openvino_device, override_sampler, sampler_name, enable_caching, is_xl_ckpt, refiner_ckpt, refiner_frac] + return [model_config, vae_ckpt, openvino_device, override_sampler, sampler_name, enable_caching, override_hires, upscaler, hires_steps, d_strength, is_xl_ckpt, refiner_ckpt, refiner_frac] - def run(self, p, model_config, vae_ckpt, openvino_device, override_sampler, sampler_name, enable_caching, is_xl_ckpt, refiner_ckpt, refiner_frac): + def run(self, p, model_config, vae_ckpt, openvino_device, override_sampler, sampler_name, enable_caching, override_hires, upscaler, hires_steps, d_strength, is_xl_ckpt, refiner_ckpt, refiner_frac): os.environ["OPENVINO_TORCH_BACKEND_DEVICE"] = str(openvino_device) if enable_caching: @@ -1225,14 +1273,12 @@ def run(self, p, model_config, vae_ckpt, openvino_device, override_sampler, samp mode = 0 if self.is_txt2img: mode = 0 - processed = process_images_openvino(p, model_config, vae_ckpt, p.sampler_name, enable_caching, openvino_device, mode, is_xl_ckpt, refiner_ckpt, refiner_frac) + processed = process_images_openvino(p, model_config, vae_ckpt, p.sampler_name, enable_caching, override_hires, upscaler, hires_steps, d_strength, openvino_device, mode, is_xl_ckpt, refiner_ckpt, refiner_frac) else: if p.image_mask is None: mode = 1 else: mode = 2 p.init = functools.partial(init_new, p) - processed = process_images_openvino(p, model_config, vae_ckpt, p.sampler_name, enable_caching, openvino_device, mode, is_xl_ckpt, refiner_ckpt, refiner_frac) + processed = process_images_openvino(p, model_config, vae_ckpt, p.sampler_name, enable_caching, override_hires, upscaler, hires_steps, d_strength, openvino_device, mode, is_xl_ckpt, refiner_ckpt, refiner_frac) return processed - -