flaxdiff 0.1.37.6__tar.gz → 0.1.38__tar.gz

{flaxdiff-0.1.37.6 → flaxdiff-0.1.38}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.1.37.6
+Version: 0.1.38
 Summary: A versatile and easy to understand Diffusion library
 Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
 License-Expression: MIT

{flaxdiff-0.1.37.6 → flaxdiff-0.1.38}/flaxdiff/models/autoencoder/diffusers.py

@@ -14,15 +14,15 @@ class StableDiffusionVAE(AutoEncoder):
     def __init__(self, modelname = "CompVis/stable-diffusion-v1-4", revision="bf16", dtype=jnp.bfloat16):
         
         from diffusers.models.vae_flax import FlaxEncoder, FlaxDecoder
-        from diffusers import FlaxStableDiffusionPipeline
+        from diffusers import FlaxStableDiffusionPipeline, FlaxAutoencoderKL
         
-        pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
+        vae, params = FlaxAutoencoderKL.from_pretrained(
             modelname,
-            revision=revision,
+            # revision=revision,
             dtype=dtype,
         )
         
-        vae = pipeline.vae
+        # vae = pipeline.vae
         
         enc = FlaxEncoder(
             in_channels=vae.config.in_channels,

flaxdiff-0.1.38/flaxdiff/models/general.py

@@ -0,0 +1,21 @@
+from flax import linen as nn
+import jax
+import jax.numpy as jnp
+
+class BCHWModelWrapper(nn.Module):
+    model: nn.Module
+
+    @nn.compact
+    def __call__(self, x, temb, textcontext):
+        # Reshape the input to BCHW format from BHWC
+        x = jnp.transpose(x, (0, 3, 1, 2))
+        # Pass the input through the UNet model
+        out = self.model(
+            sample=x,
+            timesteps=temb,
+            encoder_hidden_states=textcontext,
+        )
+        # Reshape the output back to BHWC format
+        out = jnp.transpose(out.sample, (0, 2, 3, 1))
+        return out
+    

{flaxdiff-0.1.37.6 → flaxdiff-0.1.38}/flaxdiff/samplers/common.py

@@ -133,7 +133,7 @@ class DiffusionSampler():
             
         params = params if params is not None else self.params
         
-        @jax.jit
+        # @jax.jit
         def sample_model_fn(x_t, t, *additional_inputs):
             return self.sample_model(params, x_t, t, *additional_inputs)
 

{flaxdiff-0.1.37.6 → flaxdiff-0.1.38}/flaxdiff.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.1.37.6
+Version: 0.1.38
 Summary: A versatile and easy to understand Diffusion library
 Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
 License-Expression: MIT

{flaxdiff-0.1.37.6 → flaxdiff-0.1.38}/flaxdiff.egg-info/SOURCES.txt

@@ -21,6 +21,7 @@ flaxdiff/models/__init__.py
 flaxdiff/models/attention.py
 flaxdiff/models/common.py
 flaxdiff/models/favor_fastattn.py
+flaxdiff/models/general.py
 flaxdiff/models/simple_unet.py
 flaxdiff/models/simple_vit.py
 flaxdiff/models/autoencoder/__init__.py

{flaxdiff-0.1.37.6 → flaxdiff-0.1.38}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "flaxdiff"
-version = "0.1.37.6"
+version = "0.1.38"
 description = "A versatile and easy to understand Diffusion library"
 readme = "README.md"
 authors = [