Repository: huggingface/diffusers

Issue:
TensorRT pipeline: Returns an error when compiling VAE when using AutoencoderTiny
### Describe the bug

TensorRT returns an error when compiling VAE when using AutoencoderTiny.

### Reproduction

```
import random
from PIL import Image
import torch
from diffusers import (
    LCMScheduler,
    AutoencoderTiny,
    DPMSolverMultistepScheduler,
)
from diffusers.pipelines.stable_diffusion import StableDiffusionImg2ImgPipeline
from utils import TensorRTStableDiffusionImg2ImgPipeline
import time


class Config:
    strength = 0.3
    steps = 20
    guidance_scale = 7.5
    prompt = "1girl, masterpiece, high quality, best quality, white background,"
    negative_prompt = "2girl,((EasyNegative)),monochrome,lowres,bad anatomy,bad hands,text,error,missing fingers,extra digit,fewer digits,cropped,worst quality,low quality,normal quality,jpeg artifacts,signature,watermark,username,blurry,lowres,bad anatomy,bad hands,text,error,extra digit,fewer digits,cropped,worst quality,low quality,normal quality,jpeg artifacts,signature,watermark,username,blurry,ugly,pregnant,vore,duplicate,morbid,mut ilated,tran nsexual,hermaphrodite,long neck,mutated hands,poorly drawn hands,poorly drawn face,mutation,deformed,blurry,bad anatomy,bad proportions,malformed limbs,extra limbs,cloned face,disfigured,gross proportions,(((missing arms))),(((missing legs))),(((extra arms))),(((extra legs))),pubic hair,plump,bad legs,error legs,username,blurry,bad feet"
    input_image_src = "assets/input-3.jpg"
    use_tiny_vae = True
    use_lcm = False
    use_easynegative = True
    use_tensorrt = True


model = "KBlueLeaf/kohaku-v2.1"
if Config.use_tensorrt:
    pipe = TensorRTStableDiffusionImg2ImgPipeline.from_pretrained(
        model,
        torch_dtype=torch.float16,
    )
    pipe.scheduler = DPMSolverMultistepScheduler.from_pretrained(
        model, subfolder="scheduler", use_karras_sigmas=True
    )
else:
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
        model,
        torch_dtype=torch.float16,
    )
    pipe.scheduler = DPMSolverMultistepScheduler.from_pretrained(
        model, subfolder="scheduler", use_karras_sigmas=True
    )

# EasyNegative
if Config.use_easynegative:
    pipe.load_textual_inversion(
        "sayakpaul/EasyNegative-test",
        weight_name="EasyNegative.safetensors",
        token="EasyNegative",
    )

# LCM
if Config.use_lcm:
    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
    pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")

# Tiny VAE
if Config.use_tiny_vae:
    pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesd").to(
        device=pipe.device, dtype=pipe.dtype
    )

# Safety Checker
pipe.safety_checker = None

# re-use cached folder to save ONNX models and TensorRT Engines
if Config.use_tensorrt:
    pipe.set_cached_folder(model)

pipe = pipe.to("cuda")


input_image = (
    Image.open(Config.input_image_src).convert("RGB").resize((512, 512), Image.LANCZOS)
)
input_image = input_image

while True:
    start = time.time()
    image = pipe(
        prompt=Config.prompt,
        negative_prompt=Config.negative_prompt,
        image=input_image,
        strength=Config.strength,
        # clip_skip=2,
        num_inference_steps=Config.steps,
        guidance_scale=Config.guidance_scale,
    ).images[0]

    print(f"elapsed time: {time.time() - start:.3f}s")
    image.save("output.png")
    try:
        guidance_scale = int(input("Press new guidance_scale to continue..."))
    except:
        pass

```

### Logs

```shell
Engine build can take a while to complete
Exporting model: C:\Users\Jannchie\.cache\huggingface\hub\models--KBlueLeaf--kohaku-v2.1\snapshots\0d9ad923d447081af3db92005e14718509af5fda\onnx\vae.onnx
Generating optimizing model: C:\Users\Jannchie\.cache\huggingface\hub\models--KBlueLeaf--kohaku-v2.1\snapshots\0d9ad923d447081af3db92005e14718509af5fda\onnx\vae.opt.onnx
[I] Folding Constants | Pass 1
[I]     Total Nodes | Original:    92, After Folding:    85 |     7 Nodes Folded
[I] Folding Constants | Pass 2
[I]     Total Nodes | Original:    85, After Folding:    85 |     0 Nodes Folded
Building Engines...
Engine build can take a while to complete
Exporting model: C:\Users\Jannchie\.cache\huggingface\hub\models--KBlueLeaf--kohaku-v2.1\snapshots\0d9ad923d447081af3db92005e14718509af5fda\onnx\vae_encoder.onnx
Traceback (most recent call last):
  File "C:\Code\diffusers-playground\img2img\optimize.py", line 72, in <module>
    pipe = pipe.to("cuda")
  File "C:\Code\diffusers-playground\img2img\utils\__init__.py", line 972, in to
    self.engine = build_engines(
  File "C:\Code\diffusers-playground\img2img\utils\__init__.py", line 405, in build_engines
    torch.onnx.export(
  File "C:\Users\Jannchie\anaconda3\envs\streamdiffusion\lib\site-packages\torch\onnx\utils.py", line 516, in export
    _export(
  File "C:\Users\Jannchie\anaconda3\envs\streamdiffusion\lib\site-packages\torch\onnx\utils.py", line 1596, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "C:\Users\Jannchie\anaconda3\envs\streamdiffusion\lib\site-packages\torch\onnx\utils.py", line 1135, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File "C:\Users\Jannchie\anaconda3\envs\streamdiffusion\lib\site-packages\torch\onnx\utils.py", line 1011, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "C:\Users\Jannchie\anaconda3\envs\streamdiffusion\lib\site-packages\torch\onnx\utils.py", line 915, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
  File "C:\Users\Jannchie\anaconda3\envs\streamdiffusion\lib\site-packages\torch\jit\_trace.py", line 1285, in _get_trace_graph
    outs = ONNXTracedModule(
  File "C:\Users\Jannchie\anaconda3\envs\streamdiffusion\lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Users\Jannchie\anaconda3\envs\streamdiffusion\lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Users\Jannchie\anaconda3\envs\streamdiffusion\lib\site-packages\torch\jit\_trace.py", line 133, in forward
    graph, out = torch._C._create_graph_by_tracing(
  File "C:\Users\Jannchie\anaconda3\envs\streamdiffusion\lib\site-packages\torch\jit\_trace.py", line 124, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "C:\Users\Jannchie\anaconda3\envs\streamdiffusion\lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Users\Jannchie\anaconda3\envs\streamdiffusion\lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Users\Jannchie\anaconda3\envs\streamdiffusion\lib\site-packages\torch\nn\modules\module.py", line 1508, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "C:\Code\diffusers-playground\img2img\utils\__init__.py", line 734, in forward
    return self.vae_encoder.encode(x).latent_dist.sample()
AttributeError: 'AutoencoderTinyOutput' object has no attribute 'latent_dist'
```


### System Info

- `diffusers` version: 0.24.0
- Platform: Windows-10-10.0.22631-SP0
- Python version: 3.10.13
- PyTorch version (GPU?): 2.1.0+cu121 (True)
- Huggingface_hub version: 0.20.1
- Transformers version: 4.36.2
- Accelerate version: 0.25.0
- xFormers version: 0.0.22.post7
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

cc: @asfiyab-nvidia 


Discussion:

Resolving patch:
diff --git a/examples/community/stable_diffusion_tensorrt_img2img.py b/examples/community/stable_diffusion_tensorrt_img2img.py
--- a/examples/community/stable_diffusion_tensorrt_img2img.py
+++ b/examples/community/stable_diffusion_tensorrt_img2img.py
@@ -50,6 +50,7 @@
     StableDiffusionPipelineOutput,
     StableDiffusionSafetyChecker,
 )
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
 from diffusers.schedulers import DDIMScheduler
 from diffusers.utils import logging
 
@@ -608,7 +609,7 @@ def __init__(self, model):
         self.vae_encoder = model
 
     def forward(self, x):
-        return self.vae_encoder.encode(x).latent_dist.sample()
+        return retrieve_latents(self.vae_encoder.encode(x))
 
 
 class VAEEncoder(BaseModel):


Test patch: