🐛 [Bug] The example of torch tensorrt cannot generate images correctly

[pangyoki]

Bug Description

The example of pytorch tensorrt cannot generate images correctly. There were no issues with version v2.3.0, but version v2.4.0 cannot generate images correctly.
https://github.com/pytorch/TensorRT/blob/main/examples/dynamo/torch_compile_stable_diffusion.py

To Reproduce

Steps to reproduce the behavior:

1. python3 examples/dynamo/torch_compile_stable_diffusion.py

Expected behavior

the example can generate normal image.

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

• Torch-TensorRT Version (e.g. 1.0.0): 2.4.0
• PyTorch Version (e.g. 1.0): 2.4.1
• OS (e.g., Linux): linux
• How you installed PyTorch (conda, pip, libtorch, source): pip
• Python version: 3.9
• CUDA version: 12.4

Additional context

[lanluo-nvidia]

I have verified on my linux machine with the same torch/torch_tensorrt/cuda/python version,
it is not working, it first pop up the error:
Potential NSFW content was detected in one or more images. A black image will be returned instead. Try again with a different prompt and/or seed.
if I added the safety_checker = None, requires_safety_checker = False, flag then it will generate a nonsense image like this in the attachment.

however I verified it does work on torch_tensorrt latest main branch. it can successfully generate the image attached here

[pangyoki]

@lanluo-nvidia Thank you for your reply. I tried latest main branch, and "CompVis/stable-diffusion-v1-4" model in example works well. But both the "stabilityai/stable-diffusion-xl-base-1.0" and "stabilityai/stable-diffusion-3-medium" models are unable to generate normal images. Below is the test code.

import torch
import torch_tensorrt
from diffusers import StableDiffusion3Pipeline

model_id = "stabilityai/stable-diffusion-3-medium-diffusers"

pipe = StableDiffusion3Pipeline.from_pretrained(
    model_id, torch_dtype=torch.float16
)
pipe.to("cuda")

backend = "torch_tensorrt"
pipe.transformer = torch.compile(
    pipe.transformer,
    backend=backend,
    options={
        "truncate_long_and_double": True,
        "enabled_precisions": {torch.float32, torch.float16},
    },
    dynamic=False,
)

generator = torch.Generator(pipe.device).manual_seed(7)

image = pipe(prompt=["Dog."],
    height=1280,
    width=720,
    num_inference_steps=30,
    guidance_scale=7,
    num_images_per_prompt=1,
    generator=generator).images[0]

image.save("./dog.png")

[lanluo-nvidia]

@pangyoki
Currently I am facing a CUDA out of memory error when try to run this model in my RTX4080 which has 16G of GPU memory.
I will verify on another machine with more GPU memory later.

[lanluo-nvidia]

I have verified from the latest main branch that I am not able to generate correct image from "stabilityai/stable-diffusion-3-medium-diffusers" model, it gives me a nonsense image.
however I can successfully generate image from this model without torch_tensorrt compile.
@peri044 @narendasan how can I investigate further whether it is a torch_tensorrt issue or tensorrt issue?