Unverified Commit c14c141b authored by asfiyab-nvidia's avatar asfiyab-nvidia Committed by GitHub
Browse files

TensorRT Inpaint pipeline: minor fixes (#4457)


Signed-off-by: default avatarAsfiya Baig <asfiyab@nvidia.com>
parent 79ef9e52
...@@ -1646,13 +1646,13 @@ from io import BytesIO ...@@ -1646,13 +1646,13 @@ from io import BytesIO
from PIL import Image from PIL import Image
import torch import torch
from diffusers import PNDMScheduler from diffusers import PNDMScheduler
from diffusers.pipelines.stable_diffusion import StableDiffusionImg2ImgPipeline from diffusers.pipelines.stable_diffusion import StableDiffusionInpaintPipeline
# Use the PNDMScheduler scheduler here instead # Use the PNDMScheduler scheduler here instead
scheduler = PNDMScheduler.from_pretrained("stabilityai/stable-diffusion-2-inpainting", subfolder="scheduler") 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", custom_pipeline="stable_diffusion_tensorrt_inpaint",
revision='fp16', revision='fp16',
torch_dtype=torch.float16, torch_dtype=torch.float16,
......
...@@ -1054,7 +1054,7 @@ class TensorRTStableDiffusionInpaintPipeline(StableDiffusionInpaintPipeline): ...@@ -1054,7 +1054,7 @@ class TensorRTStableDiffusionInpaintPipeline(StableDiffusionInpaintPipeline):
return_image=True, return_image=True,
), ),
) )
# print(mask)
mask = torch.nn.functional.interpolate(mask, size=(latent_height, latent_width)) mask = torch.nn.functional.interpolate(mask, size=(latent_height, latent_width))
mask = torch.cat([mask] * 2) mask = torch.cat([mask] * 2)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment