From 26a5a93e310d0d92b49d6b1c1afba56b0366d323 Mon Sep 17 00:00:00 2001 From: Asfiya Baig Date: Thu, 3 Aug 2023 17:41:25 -0700 Subject: [PATCH] TensorRT Inpaint pipeline: minor fixes Signed-off-by: Asfiya Baig --- examples/community/README.md | 4 ++-- examples/community/stable_diffusion_tensorrt_inpaint.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/community/README.md b/examples/community/README.md index 8c8461f05b96..498023307512 100755 --- a/examples/community/README.md +++ b/examples/community/README.md @@ -1646,13 +1646,13 @@ from io import BytesIO from PIL import Image import torch from diffusers import PNDMScheduler -from diffusers.pipelines.stable_diffusion import StableDiffusionImg2ImgPipeline +from diffusers.pipelines.stable_diffusion import StableDiffusionInpaintPipeline # Use the PNDMScheduler scheduler here instead scheduler = PNDMScheduler.from_pretrained("stabilityai/stable-diffusion-2-inpainting", subfolder="scheduler") -pipe = StableDiffusionImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting", +pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting", custom_pipeline="stable_diffusion_tensorrt_inpaint", revision='fp16', torch_dtype=torch.float16, diff --git a/examples/community/stable_diffusion_tensorrt_inpaint.py b/examples/community/stable_diffusion_tensorrt_inpaint.py index d17a69111048..0a5aa827d99c 100755 --- a/examples/community/stable_diffusion_tensorrt_inpaint.py +++ b/examples/community/stable_diffusion_tensorrt_inpaint.py @@ -1054,7 +1054,7 @@ def __call__( return_image=True, ), ) - # print(mask) + mask = torch.nn.functional.interpolate(mask, size=(latent_height, latent_width)) mask = torch.cat([mask] * 2)