make shapes consistent

- output `img_w x img_h x n_channels` from the VAE

Author: younesbelkada

src/diffusers/models/vae_flax.py

@@ -559,7 +559,7 @@ def setup(self):
 
     def init_weights(self, rng: jax.random.PRNGKey) -> FrozenDict:
         # init input tensors
-        sample_shape = (1, self.in_channels, self.sample_size, self.sample_size)
+        sample_shape = (1, self.sample_size, self.sample_size, self.in_channels)
         sample = jnp.zeros(sample_shape, dtype=jnp.float32)
 
         params_rng, dropout_rng, gaussian_rng = jax.random.split(rng, 3)
@@ -568,8 +568,6 @@ def init_weights(self, rng: jax.random.PRNGKey) -> FrozenDict:
         return self.init(rngs, sample)["params"]
 
     def encode(self, sample, deterministic: bool = True, return_dict: bool = True):
-        sample = jnp.transpose(sample, (0, 2, 3, 1))
-
         hidden_states = self.encoder(sample, deterministic=deterministic)
         moments = self.quant_conv(hidden_states)
         posterior = DiagonalGaussianDistribution(moments)
@@ -586,8 +584,6 @@ def decode(self, latents, deterministic: bool = True, return_dict: bool = True):
         hidden_states = self.post_quant_conv(latents)
         hidden_states = self.decoder(hidden_states, deterministic=deterministic)
 
-        hidden_states = jnp.transpose(hidden_states, (0, 3, 1, 2))
-
         if not return_dict:
             return (hidden_states,)