fix: use retrieve_latents (#6337)

huggingface · Dec 27, 2023 · c751449 · c751449
1 parent c1e8bdf
commit c751449
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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):