flaxdiff 0.1.9__tar.gz → 0.1.10__tar.gz

{flaxdiff-0.1.9 → flaxdiff-0.1.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flaxdiff
-Version: 0.1.9
+Version: 0.1.10
 Summary: A versatile and easy to understand Diffusion library
 Author: Ashish Kumar Singh
 Author-email: ashishkmr472@gmail.com

{flaxdiff-0.1.9 → flaxdiff-0.1.10}/flaxdiff/models/common.py

@@ -267,15 +267,17 @@ class ResidualBlock(nn.Module):
     kernel_init:Callable=kernel_init(1.0)
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
+    named_norms:bool=False
     
     def setup(self):
         if self.norm_groups > 0:
             norm = partial(nn.GroupNorm, self.norm_groups)
+            self.norm1 = norm(name="GroupNorm_0") if self.named_norms else norm()
+            self.norm2 = norm(name="GroupNorm_1") if self.named_norms else norm()
         else:
             norm = partial(nn.RMSNorm, 1e-5)
-        
-        self.norm1 = norm()
-        self.norm2 = norm()
+            self.norm1 = norm()
+            self.norm2 = norm()
 
     @nn.compact
     def __call__(self, x:jax.Array, temb:jax.Array, textemb:jax.Array=None, extra_features:jax.Array=None):

{flaxdiff-0.1.9 → flaxdiff-0.1.10}/flaxdiff/models/simple_unet.py

@@ -19,15 +19,15 @@ class Unet(nn.Module):
     norm_groups:int=8
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
+    named_norms: bool = False # This is for backward compatibility reasons; older checkpoints have named norms
 
     def setup(self):
         if self.norm_groups > 0:
             norm = partial(nn.GroupNorm, self.norm_groups)
+            self.conv_out_norm = norm(name="GroupNorm_0") if self.named_norms else norm()
         else:
             norm = partial(nn.RMSNorm, 1e-5)
-        
-        # self.last_up_norm = norm()
-        self.conv_out_norm = norm()
+            self.conv_out_norm = norm()
         
     @nn.compact
     def __call__(self, x, temb, textcontext):
@@ -70,7 +70,8 @@ class Unet(nn.Module):
                     activation=self.activation,
                     norm_groups=self.norm_groups,
                     dtype=self.dtype,
-                    precision=self.precision
+                    precision=self.precision,
+                    named_norms=self.named_norms
                 )(x, temb)
                 if attention_config is not None and j == self.num_res_blocks - 1:   # Apply attention only on the last block
                     x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32),
@@ -108,7 +109,8 @@ class Unet(nn.Module):
                 activation=self.activation,
                 norm_groups=self.norm_groups,
                 dtype=self.dtype,
-                precision=self.precision
+                precision=self.precision,
+                named_norms=self.named_norms
             )(x, temb)
             if middle_attention is not None and j == self.num_middle_res_blocks - 1:   # Apply attention only on the last block
                 x = TransformerBlock(heads=middle_attention['heads'], dtype=middle_attention.get('dtype', jnp.float32), 
@@ -130,7 +132,8 @@ class Unet(nn.Module):
                 activation=self.activation,
                 norm_groups=self.norm_groups,
                 dtype=self.dtype,
-                precision=self.precision
+                precision=self.precision,
+                named_norms=self.named_norms
             )(x, temb)
 
         # Upscaling Blocks
@@ -151,7 +154,8 @@ class Unet(nn.Module):
                     activation=self.activation,
                     norm_groups=self.norm_groups,
                     dtype=self.dtype,
-                    precision=self.precision
+                    precision=self.precision,
+                    named_norms=self.named_norms
                 )(x, temb)
                 if attention_config is not None and j == self.num_res_blocks - 1:   # Apply attention only on the last block
                     x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32), 
@@ -196,7 +200,8 @@ class Unet(nn.Module):
             activation=self.activation,
             norm_groups=self.norm_groups,
             dtype=self.dtype,
-            precision=self.precision
+            precision=self.precision,
+            named_norms=self.named_norms
         )(x, temb)
 
         x = self.conv_out_norm(x)

{flaxdiff-0.1.9 → flaxdiff-0.1.10}/flaxdiff/trainer/diffusion_trainer.py

@@ -16,6 +16,7 @@ from flaxdiff.utils import RandomMarkovState
 from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics
 
 from flaxdiff.models.autoencoder.autoencoder import AutoEncoder
+from flax.training.dynamic_scale import DynamicScale
 
 class TrainState(SimpleTrainState):
     rngs: jax.random.PRNGKey
@@ -83,7 +84,8 @@ class DiffusionTrainer(SimpleTrainer):
             new_state = existing_state
 
         if param_transforms is not None:
-            params = param_transforms(params)
+            new_state['params'] = param_transforms(new_state['params'])
+            new_state['ema_params'] = param_transforms(new_state['ema_params'])
 
         state = TrainState.create(
             apply_fn=model.apply,
@@ -92,7 +94,7 @@ class DiffusionTrainer(SimpleTrainer):
             tx=optimizer,
             rngs=rngs,
             metrics=Metrics.empty(),
-            dynamic_scale = flax.training.dynamic_scale.DynamicScale() if use_dynamic_scale else None
+            dynamic_scale = DynamicScale() if use_dynamic_scale else None
         )
             
         if existing_best_state is not None:

{flaxdiff-0.1.9 → flaxdiff-0.1.10}/flaxdiff/trainer/simple_trainer.py

@@ -22,7 +22,7 @@ from jax.experimental.shard_map import shard_map
 from orbax.checkpoint.utils import fully_replicated_host_local_array_to_global_array
 from termcolor import colored
 from typing import Dict, Callable, Sequence, Any, Union, Tuple
-
+from flax.training.dynamic_scale import DynamicScale
 from flaxdiff.utils import RandomMarkovState
 
 PROCESS_COLOR_MAP = {
@@ -68,7 +68,7 @@ class Metrics(metrics.Collection):
 # Define the TrainState
 class SimpleTrainState(train_state.TrainState):
     metrics: Metrics
-    dynamic_scale: flax.training.dynamic_scale.DynamicScale
+    dynamic_scale: DynamicScale
 
 class SimpleTrainer:
     state: SimpleTrainState
@@ -177,13 +177,16 @@ class SimpleTrainer:
             params = model.init(subkey, **input_vars)
         else:
             params = existing_state['params']
+            
+        if param_transforms is not None:
+            params = param_transforms(params)
 
         state = SimpleTrainState.create(
             apply_fn=model.apply,
             params=params,
             tx=optimizer,
             metrics=Metrics.empty(),
-            dynamic_scale = flax.training.dynamic_scale.DynamicScale() if use_dynamic_scale else None
+            dynamic_scale = DynamicScale() if use_dynamic_scale else None
         )
         if existing_best_state is not None:
             best_state = state.replace(

{flaxdiff-0.1.9 → flaxdiff-0.1.10}/flaxdiff.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flaxdiff
-Version: 0.1.9
+Version: 0.1.10
 Summary: A versatile and easy to understand Diffusion library
 Author: Ashish Kumar Singh
 Author-email: ashishkmr472@gmail.com

{flaxdiff-0.1.9 → flaxdiff-0.1.10}/setup.py

@@ -11,7 +11,7 @@ required_packages=[
 setup(
     name='flaxdiff',
     packages=find_packages(),
-    version='0.1.9',
+    version='0.1.10',
     description='A versatile and easy to understand Diffusion library',
     long_description=open('README.md').read(),
     long_description_content_type='text/markdown',