flaxdiff 0.1.35.5__py3-none-any.whl → 0.1.36__py3-none-any.whl

flaxdiff/trainer/simple_trainer.py

@@ -24,6 +24,7 @@ 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
+from flax.training import dynamic_scale as dynamic_scale_lib
 
 PROCESS_COLOR_MAP = {
     0: "green",
@@ -63,12 +64,12 @@ def convert_to_global_tree(global_mesh, pytree):
 @struct.dataclass
 class Metrics(metrics.Collection):
     accuracy: metrics.Accuracy
-    loss: metrics.Average.from_output('loss')
+    loss: metrics.Average#.from_output('loss')
 
 # Define the TrainState
 class SimpleTrainState(train_state.TrainState):
     metrics: Metrics
-    dynamic_scale: DynamicScale
+    dynamic_scale: dynamic_scale_lib.DynamicScale
 
 class SimpleTrainer:
     state: SimpleTrainState
@@ -110,6 +111,7 @@ class SimpleTrainer:
         
         
         if wandb_config is not None and jax.process_index() == 0:
+            import wandb
             run = wandb.init(**wandb_config)
             self.wandb = run
             
@@ -177,16 +179,13 @@ 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 = DynamicScale() if use_dynamic_scale else None
+            dynamic_scale = dynamic_scale_lib.DynamicScale() if use_dynamic_scale else None
         )
         if existing_best_state is not None:
             best_state = state.replace(
@@ -207,16 +206,16 @@ class SimpleTrainer:
         self.best_state = best_state
 
     def get_state(self):
-        # return fully_replicated_host_local_array_to_global_array()
-        return jax.tree_util.tree_map(lambda x : np.array(x), self.state)
+        return self.get_np_tree(self.state)
 
     def get_best_state(self):
-        # return convert_to_global_tree(self.mesh, flax.jax_utils.replicate(self.best_state, jax.local_devices()))
-        return jax.tree_util.tree_map(lambda x : np.array(x), self.best_state)
+        return self.get_np_tree(self.best_state)
         
     def get_rngstate(self):
-        # return convert_to_global_tree(self.mesh, flax.jax_utils.replicate(self.rngstate, jax.local_devices()))
-        return jax.tree_util.tree_map(lambda x : np.array(x), self.rngstate)
+        return self.get_np_tree(self.rngstate)
+    
+    def get_np_tree(self, pytree):
+        return jax.tree_util.tree_map(lambda x : np.array(x), pytree)
 
     def checkpoint_path(self):
         path = os.path.join(self.checkpoint_base_path, self.name.replace(' ', '_').lower())
@@ -253,29 +252,35 @@ class SimpleTrainer:
         rngstate = ckpt['rngs']
         # Convert the state to a TrainState
         self.best_loss = ckpt['best_loss']
+        if self.best_loss == 0:
+            # It cant be zero as that must have been some problem
+            self.best_loss = 1e9
         current_epoch = ckpt.get('epoch', step) # Must be a checkpoint from an older version which used epochs instead of steps
         print(
             f"Loaded model from checkpoint at epoch {current_epoch} step {step}", ckpt['best_loss'])
         return current_epoch, step, state, best_state, rngstate
 
-    def save(self, epoch=0, step=0):
+    def save(self, epoch=0, step=0, state=None, rngstate=None):
         print(f"Saving model at epoch {epoch} step {step}")
-        ckpt = {
-            # 'model': self.model,
-            'rngs': self.get_rngstate(),
-            'state': self.get_state(),
-            'best_state': self.get_best_state(),
-            'best_loss': np.array(self.best_loss),
-            'epoch': epoch,
-        }
         try:
-            save_args = orbax_utils.save_args_from_target(ckpt)
-            self.checkpointer.save(step, ckpt, save_kwargs={
-                                   'save_args': save_args}, force=True)
-            self.checkpointer.wait_until_finished()
-            pass
+            ckpt = {
+                # 'model': self.model,
+                'rngs': self.get_rngstate() if rngstate is None else self.get_np_tree(rngstate),
+                'state': self.get_state() if state is None else self.get_np_tree(state),
+                'best_state': self.get_best_state(),
+                'best_loss': np.array(self.best_loss),
+                'epoch': epoch,
+            }
+            try:
+                save_args = orbax_utils.save_args_from_target(ckpt)
+                self.checkpointer.save(step, ckpt, save_kwargs={
+                                    'save_args': save_args}, force=True)
+                self.checkpointer.wait_until_finished()
+                pass
+            except Exception as e:
+                print("Error saving checkpoint", e)
         except Exception as e:
-            print("Error saving checkpoint", e)
+            print("Error saving checkpoint outer", e)
 
     def _define_train_step(self, **kwargs):
         model = self.model
@@ -304,21 +309,26 @@ class SimpleTrainer:
             train_step = jax.pmap(train_step)
         return train_step
 
-    def _define_compute_metrics(self):
+    def _define_vaidation_step(self):
         model = self.model
         loss_fn = self.loss_fn
+        distributed_training = self.distributed_training
 
-        @jax.jit
-        def compute_metrics(state: SimpleTrainState, batch):
+        def validation_step(state: SimpleTrainState, batch):
             preds = model.apply(state.params, batch['image'])
             expected_output = batch['label']
             loss = jnp.mean(loss_fn(preds, expected_output))
+            if distributed_training:
+                loss = jax.lax.pmean(loss, "data")
             metric_updates = state.metrics.single_from_model_output(
                 loss=loss, logits=preds, labels=expected_output)
             metrics = state.metrics.merge(metric_updates)
             state = state.replace(metrics=metrics)
             return state
-        return compute_metrics
+        if distributed_training:
+            validation_step = shard_map(validation_step, mesh=self.mesh, in_specs=(P(), P('data')), out_specs=(P()))
+            validation_step = jax.pmap(validation_step)
+        return validation_step
 
     def summary(self):
         input_vars = self.get_input_ones()
@@ -343,17 +353,53 @@ class SimpleTrainer:
             "batch_size": batch_size
         })
         return summary_writer
+    
+    def validation_loop(
+        self,
+        val_state: SimpleTrainState,
+        val_step_fn: Callable,
+        val_ds,
+        val_steps_per_epoch,
+        current_step,
+    ):
+        global_device_count = jax.device_count()
+        local_device_count = jax.local_device_count()
+        process_index = jax.process_index()
+        
+        val_ds = iter(val_ds()) if val_ds else None
+        # Evaluation step
+        try:
+            for i in range(val_steps_per_epoch):
+                if val_ds is None:
+                    batch = None
+                else:
+                    batch = next(val_ds)
+                    if self.distributed_training and global_device_count > 1:
+                        batch = convert_to_global_tree(self.mesh, batch)
+                if i == 0:
+                    print(f"Evaluation started for process index {process_index}")
+                metrics = val_step_fn(val_state, batch)
+                if self.wandb is not None:
+                    # metrics is a dict of metrics
+                    if metrics and type(metrics) == dict:
+                        for key, value in metrics.items():
+                            if isinstance(value, jnp.ndarray):
+                                value = np.array(value)
+                            self.wandb.log({
+                                f"val/{key}": value,
+                            }, step=current_step)
+        except Exception as e:
+            print("Error logging images to wandb", e)
 
-    def fit(self, data, steps_per_epoch, epochs, train_step_args={}):
-        train_ds = iter(data['train']())
-        if 'test' in data:
-            test_ds = data['test']
-        else:
-            test_ds = None
-        train_step = self._define_train_step(**train_step_args)
-        compute_metrics = self._define_compute_metrics()
-        train_state = self.state
-        rng_state = self.rngstate
+    def train_loop(
+        self,
+        train_state: SimpleTrainState,
+        train_step_fn: Callable,
+        train_ds,
+        train_steps_per_epoch,
+        current_step,
+        rng_state
+    ):
         global_device_count = jax.device_count()
         local_device_count = jax.local_device_count()
         process_index = jax.process_index()
@@ -361,67 +407,105 @@ class SimpleTrainer:
             global_device_indexes = jnp.arange(global_device_count)
         else:
             global_device_indexes = 0
+            
+        epoch_loss = 0
+        current_epoch = current_step // train_steps_per_epoch
+        last_save_time = time.time()
+        
+        if process_index == 0:
+            pbar = tqdm.tqdm(total=train_steps_per_epoch, desc=f'\t\tEpoch {current_epoch}', ncols=100, unit='step')
+            
+        for i in range(train_steps_per_epoch):
+            batch = next(train_ds)
+            if i == 0:
+                print(f"First batch loaded at step {current_step}")
+                
+            if self.distributed_training and global_device_count > 1:
+            #     # Convert the local device batches to a unified global jax.Array 
+                batch = convert_to_global_tree(self.mesh, batch)
+            train_state, loss, rng_state = train_step_fn(train_state, rng_state, batch, global_device_indexes)
 
-        def train_loop(current_step, pbar: tqdm.tqdm, train_state, rng_state):
-            epoch_loss = 0
-            current_epoch = current_step // steps_per_epoch
-            last_save_time = time.time()
-            for i in range(steps_per_epoch):
-                batch = next(train_ds)
-                if self.distributed_training and global_device_count > 1:
-                    # Convert the local device batches to a unified global jax.Array 
-                    batch = convert_to_global_tree(self.mesh, batch)
-                train_state, loss, rng_state = train_step(train_state, rng_state, batch, global_device_indexes)
-
-                if self.distributed_training:
-                    loss = jax.experimental.multihost_utils.process_allgather(loss)
-                    loss = jnp.mean(loss) # Just to make sure its a scaler value
+            if i == 0:
+                print(f"Training started for process index {process_index} at step {current_step}")
+                
+            if self.distributed_training:
+                # loss = jax.experimental.multihost_utils.process_allgather(loss)
+                loss = jnp.mean(loss) # Just to make sure its a scaler value
                     
-                if loss <= 1e-6:
-                    # If the loss is too low, we can assume the model has diverged
-                    print(colored(f"Loss too low at step {current_step} => {loss}", 'red'))
-                    # Exit the training loop
-                    exit(1)
+            if loss <= 1e-6:
+                # If the loss is too low, we can assume the model has diverged
+                print(colored(f"Loss too low at step {current_step} => {loss}", 'red'))
+                # Reset the model to the old state
+                exit(1)
                             
-                epoch_loss += loss
-                current_step += 1
-                if i % 100 == 0:
-                    if pbar is not None:
-                        pbar.set_postfix(loss=f'{loss:.4f}')
-                        pbar.update(100)
-                        if self.wandb is not None:
-                            self.wandb.log({
-                                "train/step" : current_step,
-                                "train/loss": loss,
-                            }, step=current_step)
-                    # Save the model every 40 minutes
-                    if time.time() - last_save_time > 40 * 60:
-                        print(f"Saving model after 40 minutes at step {current_step}")
-                        self.save(current_epoch, current_step)
-                        last_save_time = time.time()
-            print(colored(f"Epoch done on index {process_index} => {current_epoch} Loss: {epoch_loss/steps_per_epoch}", 'green'))
-            return epoch_loss, current_step, train_state, rng_state
-
-        while self.latest_step < epochs * steps_per_epoch:
-            current_epoch = self.latest_step // steps_per_epoch
+            epoch_loss += loss
+            current_step += 1
+            if i % 100 == 0:
+                if pbar is not None:
+                    pbar.set_postfix(loss=f'{loss:.4f}')
+                    pbar.update(100)
+                    if self.wandb is not None:
+                        self.wandb.log({
+                            "train/step" : current_step,
+                            "train/loss": loss,
+                        }, step=current_step)
+                # Save the model every few steps
+                if i % 10000 == 0 and i > 0:
+                    print(f"Saving model after 10000 step {current_step}")
+                    print(f"Devices: {len(jax.devices())}") # To sync the devices
+                    self.save(current_epoch, current_step, train_state, rng_state)
+                    print(f"Saving done by process index {process_index}")
+                    last_save_time = time.time()
+        print(colored(f"Epoch done on index {process_index} => {current_epoch} Loss: {epoch_loss/train_steps_per_epoch}", 'green'))
+        if pbar is not None:
+            pbar.close()
+        return epoch_loss, current_step, train_state, rng_state
+
+
+    def fit(self, data, train_steps_per_epoch, epochs, train_step_args={}, val_steps_per_epoch=5, validation_step_args={}):
+        train_ds = iter(data['train']())
+        train_step = self._define_train_step(**train_step_args)
+        val_step = self._define_vaidation_step(**validation_step_args)
+        train_state = self.state
+        rng_state = self.rngstate
+        process_index = jax.process_index()
+        
+        if val_steps_per_epoch > 0:
+            # We should first run a validation step to make sure the model is working
+            print(f"Validation run for sanity check for process index {process_index}")
+            # Validation step
+            self.validation_loop(
+                train_state,
+                val_step,
+                data.get('test', data.get('val', None)),
+                val_steps_per_epoch,
+                self.latest_step,
+            )
+            print(colored(f"Sanity Validation done on process index {process_index}", PROCESS_COLOR_MAP[process_index]))
+                
+        while self.latest_step < epochs * train_steps_per_epoch:
+            current_epoch = self.latest_step // train_steps_per_epoch
             print(f"\nEpoch {current_epoch}/{epochs}")
             start_time = time.time()
             epoch_loss = 0
-
-            if process_index == 0:
-                with tqdm.tqdm(total=steps_per_epoch, desc=f'\t\tEpoch {current_epoch}', ncols=100, unit='step') as pbar:
-                    epoch_loss, current_step, train_state, rng_state = train_loop(self.latest_step, pbar, train_state, rng_state)
-            else:
-                epoch_loss, current_step, train_state, rng_state = train_loop(self.latest_step, None, train_state, rng_state)
-                print(colored(f"Epoch done on process index {process_index}", PROCESS_COLOR_MAP[process_index]))
+            
+            epoch_loss, current_step, train_state, rng_state = self.train_loop(
+                train_state,
+                train_step,
+                train_ds,
+                train_steps_per_epoch,
+                self.latest_step,
+                rng_state,
+            )
+            print(colored(f"Epoch done on process index {process_index}", PROCESS_COLOR_MAP[process_index]))
             
             self.latest_step = current_step
             end_time = time.time()
             self.state = train_state
             self.rngstate = rng_state
             total_time = end_time - start_time
-            avg_time_per_step = total_time / steps_per_epoch
-            avg_loss = epoch_loss / steps_per_epoch
+            avg_time_per_step = total_time / train_steps_per_epoch
+            avg_loss = epoch_loss / train_steps_per_epoch
             if avg_loss < self.best_loss:
                 self.best_loss = avg_loss
                 self.best_state = train_state
@@ -437,6 +521,18 @@ class SimpleTrainer:
                         "train/epoch": current_epoch,
                     }, step=current_step)
                 print(colored(f"\n\tEpoch {current_epoch} completed. Avg Loss: {avg_loss}, Time: {total_time:.2f}s, Best Loss: {self.best_loss}", 'green'))
-        print("Training done")
+                    
+            if val_steps_per_epoch > 0:
+                print(f"Validation started for process index {process_index}")
+                # Validation step
+                self.validation_loop(
+                    train_state,
+                    val_step,
+                    data.get('test', None),
+                    val_steps_per_epoch,
+                    current_step,
+                )
+                print(colored(f"Validation done on process index {process_index}", PROCESS_COLOR_MAP[process_index]))
+                
         self.save(epochs)
         return self.state

flaxdiff/trainer/video_diffusion_trainer.py

@@ -0,0 +1,62 @@
+import flax
+from flax import linen as nn
+import jax
+from typing import Callable
+from dataclasses import field
+import jax.numpy as jnp
+import optax
+import functools
+from jax.sharding import Mesh, PartitionSpec as P
+from jax.experimental.shard_map import shard_map
+from typing import Dict, Callable, Sequence, Any, Union, Tuple
+
+from ..schedulers import NoiseScheduler
+from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform
+
+from flaxdiff.utils import RandomMarkovState
+
+from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics
+
+from flaxdiff.models.autoencoder.autoencoder import AutoEncoder
+from flax.training import dynamic_scale as dynamic_scale_lib
+
+class TrainState(SimpleTrainState):
+    rngs: jax.random.PRNGKey
+    ema_params: dict
+
+    def apply_ema(self, decay: float = 0.999):
+        new_ema_params = jax.tree_util.tree_map(
+            lambda ema, param: decay * ema + (1 - decay) * param,
+            self.ema_params,
+            self.params,
+        )
+        return self.replace(ema_params=new_ema_params)
+
+from flaxdiff.models.autoencoder.autoencoder import AutoEncoder
+from flaxdiff.trainer.diffusion_trainer import DiffusionTrainer
+
+class SimpleVideoDiffusionTrainer(DiffusionTrainer):
+    def __init__(self,
+                 model: nn.Module,
+                 input_shapes: Dict[str, Tuple[int]],
+                 optimizer: optax.GradientTransformation,
+                 noise_schedule: NoiseScheduler,
+                 rngs: jax.random.PRNGKey,
+                 unconditional_prob: float = 0.12,
+                 name: str = "SimpleVideoDiffusion",
+                 model_output_transform: DiffusionPredictionTransform = EpsilonPredictionTransform(),
+                 autoencoder: AutoEncoder = None,
+                 **kwargs
+                 ):
+        super().__init__(
+            model=model,
+            input_shapes=input_shapes,
+            optimizer=optimizer,
+            noise_schedule=noise_schedule,
+            unconditional_prob=unconditional_prob,
+            autoencoder=autoencoder,
+            model_output_transform=model_output_transform,
+            rngs=rngs,
+            name=name,
+            **kwargs
+        )

flaxdiff/utils.py

@@ -2,7 +2,12 @@ import jax
 import jax.numpy as jnp
 import flax.struct as struct
 import flax.linen as nn
-from typing import Any
+from typing import Any, Callable
+from dataclasses import dataclass
+from functools import partial
+import numpy as np
+from jax.sharding import Mesh, PartitionSpec as P
+from abc import ABC, abstractmethod
 
 class MarkovState(struct.PyTreeNode):
     pass
@@ -17,6 +22,30 @@ class RandomMarkovState(MarkovState):
 def clip_images(images, clip_min=-1, clip_max=1):
     return jnp.clip(images, clip_min, clip_max)
 
+def _build_global_shape_and_sharding(
+    local_shape: tuple[int, ...], global_mesh: Mesh
+) -> tuple[tuple[int, ...], jax.sharding.NamedSharding]:
+  sharding = jax.sharding.NamedSharding(global_mesh, P(global_mesh.axis_names))
+  global_shape = (jax.process_count() * local_shape[0],) + local_shape[1:]
+  return global_shape, sharding
+
+
+def form_global_array(path, array: np.ndarray, global_mesh: Mesh) -> jax.Array:
+  """Put local sharded array into local devices"""
+  global_shape, sharding = _build_global_shape_and_sharding(np.shape(array), global_mesh)
+  try:
+    local_device_arrays = np.split(array, len(global_mesh.local_devices), axis=0)
+  except ValueError as array_split_error:
+    raise ValueError(
+        f"Unable to put to devices shape {array.shape} with "
+        f"local device count {len(global_mesh.local_devices)} "
+    ) from array_split_error
+  local_device_buffers = jax.device_put(local_device_arrays, global_mesh.local_devices)
+  return jax.make_array_from_single_device_arrays(global_shape, sharding, local_device_buffers)
+
+def convert_to_global_tree(global_mesh, pytree):
+    return jax.tree_util.tree_map_with_path(partial(form_global_array, global_mesh=global_mesh), pytree)
+
 class RMSNorm(nn.Module):
     """
     From "Root Mean Square Layer Normalization" by https://arxiv.org/abs/1910.07467
@@ -86,4 +115,78 @@ class RMSNorm(nn.Module):
             ).reshape(feature_shape)
             mul *= scale
         y = mul * x
-        return jnp.asarray(y, dtype)
+        return jnp.asarray(y, dtype)
+
+@dataclass
+class ConditioningEncoder(ABC):
+    model: nn.Module
+    tokenizer: Callable
+
+    def __call__(self, data):
+        tokens = self.tokenize(data)
+        outputs = self.encode_from_tokens(tokens)
+        return outputs
+        
+    def encode_from_tokens(self, tokens):
+        outputs = self.model(input_ids=tokens['input_ids'],
+                        attention_mask=tokens['attention_mask'])
+        last_hidden_state = outputs.last_hidden_state
+        return last_hidden_state
+    
+    def tokenize(self, data):
+        tokens = self.tokenizer(data, padding="max_length",
+                        max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="np")
+        return tokens
+    
+@dataclass
+class TextEncoder(ConditioningEncoder):
+    # def __call__(self, data):
+    #     tokens = self.tokenize(data)
+    #     outputs = self.encode_from_tokens(tokens)
+    #     return outputs
+        
+    # def encode_from_tokens(self, tokens):
+    #     outputs = self.model(input_ids=tokens['input_ids'],
+    #                     attention_mask=tokens['attention_mask'])
+    #     last_hidden_state = outputs.last_hidden_state
+    #     # pooler_output = outputs.pooler_output  # pooled (EOS token) states
+    #     # embed_pooled = pooler_output  # .astype(jnp.float16)
+    #     embed_labels_full = last_hidden_state  # .astype(jnp.float16)
+
+    #     return embed_labels_full
+    pass
+
+class AutoTextTokenizer:
+    def __init__(self, tensor_type="pt", modelname="openai/clip-vit-large-patch14"):
+        from transformers import AutoTokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(modelname)
+        self.tensor_type = tensor_type
+
+    def __call__(self, inputs):
+        # print(caption)
+        tokens = self.tokenizer(inputs, padding="max_length", max_length=self.tokenizer.model_max_length,
+                                truncation=True, return_tensors=self.tensor_type)
+        # print(tokens.keys())
+        return {
+            "input_ids": tokens["input_ids"],
+            "attention_mask": tokens["attention_mask"],
+            "caption": inputs,
+        }
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
+
+def defaultTextEncodeModel(backend="jax"):
+    from transformers import (
+        CLIPTextModel,
+        FlaxCLIPTextModel,
+        AutoTokenizer,
+    )
+    modelname = "openai/clip-vit-large-patch14"
+    if backend == "jax":
+        model = FlaxCLIPTextModel.from_pretrained(
+            modelname, dtype=jnp.bfloat16)
+    else:
+        model = CLIPTextModel.from_pretrained(modelname)
+    tokenizer = AutoTokenizer.from_pretrained(modelname, dtype=jnp.float16)
+    return TextEncoder(model, tokenizer)

{flaxdiff-0.1.35.5.dist-info → flaxdiff-0.1.36.dist-info}/METADATA

@@ -1,15 +1,21 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.1.35.5
+Version: 0.1.36
 Summary: A versatile and easy to understand Diffusion library
 Author: Ashish Kumar Singh
 Author-email: ashishkmr472@gmail.com
 Description-Content-Type: text/markdown
-Requires-Dist: flax >=0.8.4
-Requires-Dist: optax >=0.2.2
-Requires-Dist: jax >=0.4.28
+Requires-Dist: flax>=0.8.4
+Requires-Dist: optax>=0.2.2
+Requires-Dist: jax>=0.4.28
 Requires-Dist: orbax
 Requires-Dist: clu
+Dynamic: author
+Dynamic: author-email
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: requires-dist
+Dynamic: summary
 
 # ![](images/logo.jpeg "FlaxDiff")
 

flaxdiff-0.1.36.dist-info/RECORD

@@ -0,0 +1,43 @@
+flaxdiff/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flaxdiff/utils.py,sha256=b_hFXsam2NICQYCFk0EOcqtBjM-RUqnN0NKTn0lQ070,6532
+flaxdiff/data/__init__.py,sha256=PM3PkHihyohT5SHVYKc8vQ4IeVfGPpCktkSVwvqMjQ4,52
+flaxdiff/data/dataset_map.py,sha256=hcHaoR2IbNQmfyPUhYd6_8xinurxxCqawQijAsDI0Ek,3093
+flaxdiff/data/datasets.py,sha256=YUMoSvF2yAyikRvRofZVlHwfEOU3zXSSG4KkLnVfpoA,5626
+flaxdiff/data/online_loader.py,sha256=1Fi_QRixxRzbt602nORINcDeHEccvCrBpagrz4PURYg,12499
+flaxdiff/models/__init__.py,sha256=FAivVYXxM2JrCFIXf-C3374RB2Hth25dBrzOeNFhH1U,26
+flaxdiff/models/attention.py,sha256=JvrP7-09MV6IfRLRBhqjPmNUU-lkEMk9TOnJSBKcar8,13289
+flaxdiff/models/common.py,sha256=hWsSs2BP2J-JN1s4qLRr-h-KYkcVyl2hOp1Wsm_L-h8,10994
+flaxdiff/models/favor_fastattn.py,sha256=79Ew1nqarsNLPzZaBSd1ILORzJr74CupYeqGiCQK5E4,27689
+flaxdiff/models/simple_unet.py,sha256=L5m2j5580QP7pJ5VIme7U5xYA22PZiGP7qdvcKUnB38,11463
+flaxdiff/models/simple_vit.py,sha256=UCDDr0XVnpf6tbJWKFtEt3_nAqMqOoakXf5amyVWZNo,7929
+flaxdiff/models/autoencoder/__init__.py,sha256=qY-7MldZpsfkF-_T2LqlRK7VHbqfmosz0NmvzDlBkOk,78
+flaxdiff/models/autoencoder/autoencoder.py,sha256=27_hYl0yXAdH9Mx4Xu9J79mSNo-FEKr9SxhVaS3ffn4,591
+flaxdiff/models/autoencoder/diffusers.py,sha256=JHeFLCxiHhu-QHwhKiCuKsQJn4AZumquiuxgZkiYGQ0,3643
+flaxdiff/models/autoencoder/simple_autoenc.py,sha256=UXHPgDmwGTnv3Uts6Zj3p9R9nJXnEiEXbllgarwDfXM,805
+flaxdiff/predictors/__init__.py,sha256=SKkYYRF9Wfgk2zhtZw4vCXOdOeRlrm2Mk6cvuaEvAzc,4403
+flaxdiff/samplers/__init__.py,sha256=_S-9TwDeshrI0VmapV-J2hqjTByOa0-oOeUs_IdovjU,285
+flaxdiff/samplers/common.py,sha256=ZA08VyovxegpRx4wOQq9LSwZi0gSCz2lrbS5oVYOEYg,8488
+flaxdiff/samplers/ddim.py,sha256=pB8Kod8ZLJ3GXev4uM3cOj1Uy6ibR0jsaZa-VE0fyJM,552
+flaxdiff/samplers/ddpm.py,sha256=u1OchQu0XPhc_6w9JXoaFp2wo4y-zXyQNtGAIJwxNLg,2209
+flaxdiff/samplers/euler.py,sha256=Htb-IJeu7jSgY6mvgYr9yl9pUnos49vijlVk5IQsRps,2740
+flaxdiff/samplers/heun_sampler.py,sha256=UyI-hSlyWvt-7VEUJj27zjgyzKkGVl8fDUHV-YpSOCc,1421
+flaxdiff/samplers/multistep_dpm.py,sha256=3Wu3MrMLYaBb1ObraTbWrJmtEtU0adl1dDbz5fPJ4Gs,2735
+flaxdiff/samplers/rk4_sampler.py,sha256=1j1pES_Q2QiaURvEWeedbbT1LHmkc3jsu0GgH83qBL0,1926
+flaxdiff/schedulers/__init__.py,sha256=3id390WEfdf-MN-oLSPAhlRFIXrFWr6ioAHPAwURJyE,375
+flaxdiff/schedulers/common.py,sha256=b-W4iI-aqScpVE8VZbBpiYvAVI6rqDkUP-C_hEVBwCI,4151
+flaxdiff/schedulers/continuous.py,sha256=5c_niOA20fxJ5oJDi09FfayIRogBGwtfG0XThW2IUZk,334
+flaxdiff/schedulers/cosine.py,sha256=9ban0dFHLMm35wQvaBT4nCQwPGmzNsXwQ1xI0oppmJI,2005
+flaxdiff/schedulers/discrete.py,sha256=O54wH2HVu3olJA71NxgAXFW9cr6B6Gl-DR_uZeytpds,3319
+flaxdiff/schedulers/exp.py,sha256=cPTnUJpYdzJRRZqMLYQz0rRUCpEmaP2tXhRumLx94jA,605
+flaxdiff/schedulers/karras.py,sha256=4GN120kGwdxxU-h2mVdhBVy9IORkUMm_vvz3XjthBcI,3355
+flaxdiff/schedulers/linear.py,sha256=6003F5ISq1Wc0h6UAzY95MJgsDIKGMhBzbiVALpea0k,581
+flaxdiff/schedulers/sqrt.py,sha256=1F84ZgQPuoNMhe6yxGTR2G0h7dPOZtm4UDQOakbSsEU,445
+flaxdiff/trainer/__init__.py,sha256=T-vUVq4zHcMK6kpCsG4Gu8vn71q6lZD-lg-Ul7yKfEk,128
+flaxdiff/trainer/autoencoder_trainer.py,sha256=hxihkRL9WCIQVGOP-pc1jjjIUaRXDLcNo3_erTKsuWM,7049
+flaxdiff/trainer/diffusion_trainer.py,sha256=ajOWBgFFwXP_VQScUjcuPoaB4Gk02aF0Ls5LNlA8wqA,12691
+flaxdiff/trainer/simple_trainer.py,sha256=jCD9-qCwX0SC0rN3GrXUBfRrndWNqUI0HmbOAbmYBMM,21906
+flaxdiff/trainer/video_diffusion_trainer.py,sha256=gMkKpnKNTo8QhTx5ptEEkc7W5-7rzXIr9queU53hXyQ,2197
+flaxdiff-0.1.36.dist-info/METADATA,sha256=7fO1e_icIEK6dmSopv538Hm2fQnhnkOAE2Ab9inpcNE,22213
+flaxdiff-0.1.36.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+flaxdiff-0.1.36.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.1.36.dist-info/RECORD,,

{flaxdiff-0.1.35.5.dist-info → flaxdiff-0.1.36.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.1.0)
+Generator: setuptools (78.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any