From 8af29e884c1bcfc48808c76af04c1e18a4783f0a Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Thu, 19 Oct 2023 16:46:55 +0200 Subject: [PATCH] format --- optimum/onnxruntime/modeling_diffusion.py | 8 +++---- .../diffusers/pipeline_latent_consistency.py | 23 ++++++++++--------- optimum/utils/input_generators.py | 2 +- 3 files changed, 17 insertions(+), 16 deletions(-) diff --git a/optimum/onnxruntime/modeling_diffusion.py b/optimum/onnxruntime/modeling_diffusion.py index 3e9adcccde..88e2339b64 100644 --- a/optimum/onnxruntime/modeling_diffusion.py +++ b/optimum/onnxruntime/modeling_diffusion.py @@ -47,7 +47,9 @@ from ..pipelines.diffusers.pipeline_stable_diffusion_inpaint import StableDiffusionInpaintPipelineMixin from ..pipelines.diffusers.pipeline_stable_diffusion_xl import StableDiffusionXLPipelineMixin from ..pipelines.diffusers.pipeline_stable_diffusion_xl_img2img import StableDiffusionXLImg2ImgPipelineMixin -from ..pipelines.diffusers.pipeline_stable_diffusion_latent_consistency import LatentConsistencyModelPipelinePipelineMixin +from ..pipelines.diffusers.pipeline_stable_diffusion_latent_consistency import ( + LatentConsistencyModelPipelinePipelineMixin, +) from ..pipelines.diffusers.pipeline_utils import VaeImageProcessor from ..utils import ( DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER, @@ -565,18 +567,16 @@ class ORTStableDiffusionInpaintPipeline(ORTStableDiffusionPipelineBase, StableDi __call__ = StableDiffusionInpaintPipelineMixin.__call__ - @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTLatentConsistencyModelPipeline(ORTStableDiffusionPipelineBase, LatentConsistencyModelPipelinePipelineMixin): """ ONNX Runtime-powered stable diffusion pipeline corresponding to [diffusers.LatentConsistencyModelPipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/latent_consistency#diffusers.LatentConsistencyModelPipeline). """ + auto_model_class = LatentConsistencyModelPipeline __call__ = LatentConsistencyModelPipelinePipelineMixin.__call__ - - class ORTStableDiffusionXLPipelineBase(ORTStableDiffusionPipelineBase): auto_model_class = StableDiffusionXLImg2ImgPipeline diff --git a/optimum/pipelines/diffusers/pipeline_latent_consistency.py b/optimum/pipelines/diffusers/pipeline_latent_consistency.py index b9287a178d..13bd165e37 100644 --- a/optimum/pipelines/diffusers/pipeline_latent_consistency.py +++ b/optimum/pipelines/diffusers/pipeline_latent_consistency.py @@ -28,12 +28,11 @@ class LatentConsistencyModelPipelinePipelineMixin(StableDiffusionPipelineMixin): - # Adapted from https://github.com/huggingface/diffusers/blob/v0.22.0/src/diffusers/pipelines/latent_consistency/pipeline_latent_consistency.py#L264 def __call__( self, prompt: Optional[Union[str, List[str]]] = None, - height: Optional[int] = 768, # TODO : default to None + height: Optional[int] = 768, # TODO : default to None width: Optional[int] = 768, num_inference_steps: int = 4, guidance_scale: float = 7.5, @@ -110,10 +109,9 @@ def __call__( height = height or self.unet.config.get("sample_size", 64) * self.vae_scale_factor width = width or self.unet.config.get("sample_size", 64) * self.vae_scale_factor - - # Don't need to get negative prompts due to LCM guided distillation - negative_prompt=None - negative_prompt_embeds=None + # Don't need to get negative prompts due to LCM guided distillation + negative_prompt = None + negative_prompt_embeds = None # check inputs. Raise error if not correct self.check_inputs( @@ -164,10 +162,14 @@ def __call__( num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order for i, t in enumerate(self.progress_bar(timesteps)): - # predict the noise residual timestep = np.full(bs, t, dtype=timestep_dtype) - noise_pred = self.unet(sample=latent_model_input, timestep=timestep, encoder_hidden_states=prompt_embeds, timestep_cond=w_embedding) + noise_pred = self.unet( + sample=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + timestep_cond=w_embedding, + ) noise_pred = noise_pred[0] # compute the previous noisy sample x_t -> x_t-1 @@ -202,12 +204,11 @@ def __call__( return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) - # Adapted from https://github.com/huggingface/diffusers/blob/v0.22.0/src/diffusers/pipelines/latent_consistency/pipeline_latent_consistency.py#L264 def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=None): """ See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - + Args: timesteps (`torch.Tensor`): generate embedding vectors at these timesteps @@ -215,7 +216,7 @@ def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=None): dimension of the embeddings to generate dtype: data type of the generated embeddings - + Returns: `torch.FloatTensor`: Embedding vectors with shape `(len(timesteps), embedding_dim)` """ diff --git a/optimum/utils/input_generators.py b/optimum/utils/input_generators.py index 700aa4fccb..1028307f8c 100644 --- a/optimum/utils/input_generators.py +++ b/optimum/utils/input_generators.py @@ -710,7 +710,7 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int if input_name == "timestep": shape = [self.batch_size] return self.random_int_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=int_dtype) - + if input_name == "text_embeds": dim = self.text_encoder_projection_dim elif input_name == "timestep_cond":