Skip to content

Commit 3dc10a5

Browse files
Update TensorRT txt2img and inpaint community pipelines (huggingface#9037)
* Update TensorRT txt2img and inpaint community pipelines Signed-off-by: Asfiya Baig <[email protected]> * update tensorrt install instructions Signed-off-by: Asfiya Baig <[email protected]> --------- Signed-off-by: Asfiya Baig <[email protected]> Co-authored-by: Sayak Paul <[email protected]>
1 parent c370b90 commit 3dc10a5

File tree

4 files changed

+417
-152
lines changed

4 files changed

+417
-152
lines changed

examples/community/README.md

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1487,17 +1487,16 @@ NOTE: The ONNX conversions and TensorRT engine build may take up to 30 minutes.
14871487
```python
14881488
import torch
14891489
from diffusers import DDIMScheduler
1490-
from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline
1490+
from diffusers.pipelines import DiffusionPipeline
14911491

14921492
# Use the DDIMScheduler scheduler here instead
1493-
scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-2-1",
1494-
subfolder="scheduler")
1493+
scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="scheduler")
14951494

1496-
pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1",
1497-
custom_pipeline="stable_diffusion_tensorrt_txt2img",
1498-
variant='fp16',
1499-
torch_dtype=torch.float16,
1500-
scheduler=scheduler,)
1495+
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1",
1496+
custom_pipeline="stable_diffusion_tensorrt_txt2img",
1497+
variant='fp16',
1498+
torch_dtype=torch.float16,
1499+
scheduler=scheduler,)
15011500

15021501
# re-use cached folder to save ONNX models and TensorRT Engines
15031502
pipe.set_cached_folder("stabilityai/stable-diffusion-2-1", variant='fp16',)
@@ -2231,12 +2230,12 @@ from io import BytesIO
22312230
from PIL import Image
22322231
import torch
22332232
from diffusers import PNDMScheduler
2234-
from diffusers.pipelines.stable_diffusion import StableDiffusionInpaintPipeline
2233+
from diffusers.pipelines import DiffusionPipeline
22352234

22362235
# Use the PNDMScheduler scheduler here instead
22372236
scheduler = PNDMScheduler.from_pretrained("stabilityai/stable-diffusion-2-inpainting", subfolder="scheduler")
22382237

2239-
pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting",
2238+
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting",
22402239
custom_pipeline="stable_diffusion_tensorrt_inpaint",
22412240
variant='fp16',
22422241
torch_dtype=torch.float16,

examples/community/stable_diffusion_tensorrt_img2img.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@
6060
"""
6161
Installation instructions
6262
python3 -m pip install --upgrade transformers diffusers>=0.16.0
63-
python3 -m pip install --upgrade tensorrt-cu12==10.2.0
63+
python3 -m pip install --upgrade tensorrt~=10.2.0
6464
python3 -m pip install --upgrade polygraphy>=0.47.0 onnx-graphsurgeon --extra-index-url https://pypi.ngc.nvidia.com
6565
python3 -m pip install onnxruntime
6666
"""
@@ -659,7 +659,7 @@ class TensorRTStableDiffusionImg2ImgPipeline(DiffusionPipeline):
659659
r"""
660660
Pipeline for image-to-image generation using TensorRT accelerated Stable Diffusion.
661661
662-
This model inherits from [`StableDiffusionImg2ImgPipeline`]. Check the superclass documentation for the generic methods the
662+
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
663663
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
664664
665665
Args:

0 commit comments

Comments
 (0)