reformatted

Author: prishajain1

src/maxdiffusion/models/vae_flax.py

@@ -932,6 +932,7 @@ def __call__(self, sample, sample_posterior=False, deterministic: bool = True, r
 
     return FlaxDecoderOutput(sample=sample)
 
+
 class WanDiagonalGaussianDistribution(FlaxDiagonalGaussianDistribution):
   pass
 

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -24,7 +24,12 @@
 from ...configuration_utils import ConfigMixin
 from ..modeling_flax_utils import FlaxModelMixin, get_activation
 from ... import common_types
-from ..vae_flax import (FlaxAutoencoderKLOutput, FlaxDiagonalGaussianDistribution, FlaxDecoderOutput, WanDiagonalGaussianDistribution)
+from ..vae_flax import (
+    FlaxAutoencoderKLOutput,
+    FlaxDiagonalGaussianDistribution,
+    FlaxDecoderOutput,
+    WanDiagonalGaussianDistribution,
+)
 
 BlockSizes = common_types.BlockSizes
 
@@ -944,8 +949,6 @@ def __call__(self, x: jax.Array, feat_cache=None, feat_idx=0):
     return x, feat_cache, jnp.array(feat_idx, dtype=jnp.int32)
 
 
-
-
 class AutoencoderKLWanCache:
 
   def __init__(self, module):