optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py [725:789]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - if use_warmup_inference_steps and i == throughput_warmup_steps: t1_inf = time.time() t1 += t1_inf - t0_inf if self.interrupt: continue timestep = timesteps[0] timesteps = torch.roll(timesteps, shifts=-1, dims=0) # expand the latents if we are doing classifier free guidance latent_model_input = ( torch.cat([latents_batch] * 2) if self.do_classifier_free_guidance else latents_batch ) latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep) # predict the noise residual added_cond_kwargs = {"text_embeds": add_text_embeddings_batch, "time_ids": add_time_ids_batch} if ip_adapter_image is not None: added_cond_kwargs["image_embeds"] = image_embeds noise_pred = self.unet_hpu( latent_model_input, timestep, text_embeddings_batch, timestep_cond, self.cross_attention_kwargs, added_cond_kwargs, ) # perform guidance if self.do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) if self.do_classifier_free_guidance and self.guidance_rescale > 0.0: # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf noise_pred = rescale_noise_cfg( noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale ) # compute the previous noisy sample x_t -> x_t-1 latents_batch = self.scheduler.step( noise_pred, timestep, latents_batch, **extra_step_kwargs, return_dict=False )[0] if not self.use_hpu_graphs: self.htcore.mark_step() if callback_on_step_end is not None: callback_kwargs = {} for k in callback_on_step_end_tensor_inputs: callback_kwargs[k] = locals()[k] callback_outputs = callback_on_step_end(self, i, timestep, callback_kwargs) latents_batch = callback_outputs.pop("latents", latents_batch) _prompt_embeds = callback_outputs.pop("prompt_embeds", None) _negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", None) if _prompt_embeds is not None and _negative_prompt_embeds is not None: text_embeddings_batch = torch.cat([_negative_prompt_embeds, _prompt_embeds]) _add_text_embeds = callback_outputs.pop("add_text_embeds", None) _negative_pooled_prompt_embeds = callback_outputs.pop("negative_pooled_prompt_embeds", None) if _add_text_embeds is not None and _negative_pooled_prompt_embeds is not None: add_text_embeddings_batch = torch.cat([_negative_pooled_prompt_embeds, _add_text_embeds]) _add_time_ids = callback_outputs.pop("add_time_ids", None) _negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", None) if _add_time_ids is not None and _negative_add_time_ids is not None: add_time_ids_batch = torch.cat([_add_time_ids, _negative_add_time_ids]) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py [571:635]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - if use_warmup_inference_steps and i == throughput_warmup_steps: t1_inf = time.time() t1 += t1_inf - t0_inf if self.interrupt: continue timestep = timesteps[0] timesteps = torch.roll(timesteps, shifts=-1, dims=0) # expand the latents if we are doing classifier free guidance latent_model_input = ( torch.cat([latents_batch] * 2) if self.do_classifier_free_guidance else latents_batch ) latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep) # predict the noise residual added_cond_kwargs = {"text_embeds": add_text_embeddings_batch, "time_ids": add_time_ids_batch} if ip_adapter_image is not None: added_cond_kwargs["image_embeds"] = image_embeds noise_pred = self.unet_hpu( latent_model_input, timestep, text_embeddings_batch, timestep_cond, self.cross_attention_kwargs, added_cond_kwargs, ) # perform guidance if self.do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) if self.do_classifier_free_guidance and self.guidance_rescale > 0.0: # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf noise_pred = rescale_noise_cfg( noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale ) # compute the previous noisy sample x_t -> x_t-1 latents_batch = self.scheduler.step( noise_pred, timestep, latents_batch, **extra_step_kwargs, return_dict=False )[0] if not self.use_hpu_graphs: self.htcore.mark_step() if callback_on_step_end is not None: callback_kwargs = {} for k in callback_on_step_end_tensor_inputs: callback_kwargs[k] = locals()[k] callback_outputs = callback_on_step_end(self, i, timestep, callback_kwargs) latents_batch = callback_outputs.pop("latents", latents_batch) _prompt_embeds = callback_outputs.pop("prompt_embeds", None) _negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", None) if _prompt_embeds is not None and _negative_prompt_embeds is not None: text_embeddings_batch = torch.cat([_negative_prompt_embeds, _prompt_embeds]) _add_text_embeds = callback_outputs.pop("add_text_embeds", None) _negative_pooled_prompt_embeds = callback_outputs.pop("negative_pooled_prompt_embeds", None) if _add_text_embeds is not None and _negative_pooled_prompt_embeds is not None: add_text_embeddings_batch = torch.cat([_negative_pooled_prompt_embeds, _add_text_embeds]) _add_time_ids = callback_outputs.pop("add_time_ids", None) _negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", None) if _add_time_ids is not None and _negative_add_time_ids is not None: add_time_ids_batch = torch.cat([_add_time_ids, _negative_add_time_ids]) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -