fix(api): return correct output from tiled VAE

ssube · Apr 29, 2023 · b4dfc44 · b4dfc44
1 parent f420a3a
commit b4dfc44
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/api/onnx_web/diffusers/patches/vae.py b/api/onnx_web/diffusers/patches/vae.py
@@ -123,13 +123,10 @@ def tiled_encode(
             result_rows.append(torch.cat(result_row, dim=3))
 
         moments = torch.cat(result_rows, dim=2)
-        posterior = DiagonalGaussianDistribution(moments)
-        posterior = posterior.numpy()
-
         if not return_dict:
-            return (posterior,)
+            return (moments,)
 
-        return AutoencoderKLOutput(latent_dist=posterior)
+        return AutoencoderKLOutput(latent_dist=moments)
 
     @torch.no_grad()
     def tiled_decode(