From 978126316ba5c0112e0b20f81ececefc824178cd Mon Sep 17 00:00:00 2001 From: Jingya HUANG <44135271+JingyaHuang@users.noreply.github.com> Date: Mon, 9 Sep 2024 21:35:33 +0000 Subject: [PATCH] make style --- .../pipelines/diffusers/pipeline_controlnet.py | 6 +++++- tests/inference/test_stable_diffusion_pipeline.py | 12 ++++++------ 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/optimum/neuron/pipelines/diffusers/pipeline_controlnet.py b/optimum/neuron/pipelines/diffusers/pipeline_controlnet.py index e304c4a83..daaa4f9a4 100644 --- a/optimum/neuron/pipelines/diffusers/pipeline_controlnet.py +++ b/optimum/neuron/pipelines/diffusers/pipeline_controlnet.py @@ -510,7 +510,11 @@ def __call__( ) else: for i, scale in enumerate(cond_scale): - new_scale = torch.tensor([scale]).repeat(2) if self.data_parallel_mode == "unet" else torch.tensor(scale) + new_scale = ( + torch.tensor([scale]).repeat(2) + if self.data_parallel_mode == "unet" + else torch.tensor(scale) + ) cond_scale[i] = new_scale down_block_res_samples, mid_block_res_sample = self.controlnet( diff --git a/tests/inference/test_stable_diffusion_pipeline.py b/tests/inference/test_stable_diffusion_pipeline.py index 028deab89..baca33668 100644 --- a/tests/inference/test_stable_diffusion_pipeline.py +++ b/tests/inference/test_stable_diffusion_pipeline.py @@ -37,11 +37,11 @@ ) from optimum.neuron.modeling_diffusion import ( NeuronControlNetModel, - NeuronMultiControlNetModel, NeuronModelTextEncoder, NeuronModelUnet, NeuronModelVaeDecoder, NeuronModelVaeEncoder, + NeuronMultiControlNetModel, ) from optimum.neuron.utils.testing_utils import is_inferentia_test, requires_neuronx from optimum.utils import logging @@ -225,9 +225,9 @@ def prepare_canny_image(image_url=None): image = image[:, :, None] image = np.concatenate([image, image, image], axis=2) canny_image = PIL.Image.fromarray(image) - - return canny_image - + + return canny_image + @parameterized.expand(SUPPORTED_ARCHITECTURES, skip_on_empty=True) def test_export_and_inference_with_single_controlnet(self, model_arch): input_shapes = copy.deepcopy(self.STATIC_INPUTS_SHAPES) @@ -251,13 +251,13 @@ def test_export_and_inference_with_single_controlnet(self, model_arch): image = neuron_pipeline(prompt, image=canny_image).images[0] neuron_pipeline.scheduler = UniPCMultistepScheduler.from_config(neuron_pipeline.scheduler.config) self.assertIsInstance(image, PIL.Image.Image) - + @parameterized.expand(SUPPORTED_ARCHITECTURES, skip_on_empty=True) def test_export_and_inference_with_multiple_controlnet(self, model_arch): input_shapes = copy.deepcopy(self.STATIC_INPUTS_SHAPES) input_shapes.update({"num_images_per_prompt": 1}) controlnet_id = "hf-internal-testing/tiny-controlnet" - + neuron_pipeline = NeuronStableDiffusionControlNetPipeline.from_pretrained( MODEL_NAMES[model_arch], controlnet_ids=[controlnet_id, controlnet_id],