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

flaxdiff/trainer/simple_trainer.py

@@ -1,538 +0,0 @@
-import orbax.checkpoint
-import tqdm
-from flax import linen as nn
-import jax
-from typing import Callable
-from dataclasses import field
-import jax.numpy as jnp
-import numpy as np
-from functools import partial
-from clu import metrics
-from flax.training import train_state  # Useful dataclass to keep train state
-import optax
-from flax import struct                # Flax dataclasses
-import flax
-import time
-import os
-import orbax
-from flax.training import orbax_utils
-from jax.sharding import Mesh, PartitionSpec as P
-from jax.experimental import mesh_utils
-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
-from flax.training import dynamic_scale as dynamic_scale_lib
-
-PROCESS_COLOR_MAP = {
-    0: "green",
-    1: "yellow",
-    2: "magenta",
-    3: "cyan", 
-    4: "white",
-    5: "light_blue",
-    6: "light_red",
-    7: "light_cyan"
-}
-
-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)
-
-@struct.dataclass
-class Metrics(metrics.Collection):
-    accuracy: metrics.Accuracy
-    loss: metrics.Average#.from_output('loss')
-
-# Define the TrainState
-class SimpleTrainState(train_state.TrainState):
-    metrics: Metrics
-    dynamic_scale: dynamic_scale_lib.DynamicScale
-
-class SimpleTrainer:
-    state: SimpleTrainState
-    best_state: SimpleTrainState
-    best_loss: float
-    model: nn.Module
-    ema_decay: float = 0.999
-
-    def __init__(self,
-                 model: nn.Module,
-                 input_shapes: Dict[str, Tuple[int]],
-                 optimizer: optax.GradientTransformation,
-                 rngs: jax.random.PRNGKey,
-                 train_state: SimpleTrainState = None,
-                 name: str = "Simple",
-                 load_from_checkpoint: str = None,
-                 checkpoint_suffix: str = "",
-                 loss_fn=optax.l2_loss,
-                 param_transforms: Callable = None,
-                 wandb_config: Dict[str, Any] = None,
-                 distributed_training: bool = None,
-                 checkpoint_base_path: str = "./checkpoints",
-                 checkpoint_step: int = None,
-                 use_dynamic_scale: bool = False,
-                 ):
-        if distributed_training is None or distributed_training is True:
-            # Auto-detect if we are running on multiple devices
-            distributed_training = jax.device_count() > 1
-            self.mesh = jax.sharding.Mesh(jax.devices(), 'data')
-        else:
-            self.mesh = None
-
-        self.distributed_training = distributed_training
-        self.model = model
-        self.name = name
-        self.loss_fn = loss_fn
-        self.input_shapes = input_shapes
-        self.checkpoint_base_path = checkpoint_base_path
-        
-        
-        if wandb_config is not None and jax.process_index() == 0:
-            import wandb
-            run = wandb.init(**wandb_config)
-            self.wandb = run
-            
-            # define our custom x axis metric
-            self.wandb.define_metric("train/step")
-            self.wandb.define_metric("train/epoch")
-            
-            self.wandb.define_metric("train/loss", step_metric="train/step")
-            
-            self.wandb.define_metric("train/epoch_time", step_metric="train/epoch")
-            self.wandb.define_metric("train/avg_time_per_step", step_metric="train/epoch")
-            self.wandb.define_metric("train/avg_loss", step_metric="train/epoch")
-            self.wandb.define_metric("train/best_loss", step_metric="train/epoch")
-            
-        # checkpointer = orbax.checkpoint.PyTreeCheckpointer()
-        async_checkpointer = orbax.checkpoint.AsyncCheckpointer(orbax.checkpoint.PyTreeCheckpointHandler(), timeout_secs=60)
-
-        options = orbax.checkpoint.CheckpointManagerOptions(
-            max_to_keep=4, create=True)
-        self.checkpointer = orbax.checkpoint.CheckpointManager(
-            self.checkpoint_path() + checkpoint_suffix, async_checkpointer, options)
-
-        if load_from_checkpoint is not None:
-            latest_epoch, latest_step, old_state, old_best_state, rngstate = self.load(load_from_checkpoint, checkpoint_step)
-        else:
-            latest_epoch, latest_step, old_state, old_best_state, rngstate = 0, 0, None, None, None
-
-        self.latest_step = latest_step
-        
-        if rngstate:
-            self.rngstate = RandomMarkovState(**rngstate)
-        else:
-            self.rngstate = RandomMarkovState(rngs)
-            
-        self.rngstate, subkey = self.rngstate.get_random_key()
-
-        if train_state == None:
-            state, best_state = self.generate_states(
-                optimizer, subkey, old_state, old_best_state, model, param_transforms, use_dynamic_scale
-            )
-            self.init_state(state, best_state)
-        else:
-            self.state = train_state
-            self.best_state = train_state
-            self.best_loss = 1e9
-
-    def get_input_ones(self):
-        return {k: jnp.ones((1, *v)) for k, v in self.input_shapes.items()}
-
-    def generate_states(
-        self,
-        optimizer: optax.GradientTransformation,
-        rngs: jax.random.PRNGKey,
-        existing_state: dict = None,
-        existing_best_state: dict = None,
-        model: nn.Module = None,
-        param_transforms: Callable = None,
-        use_dynamic_scale: bool = False
-    ) -> Tuple[SimpleTrainState, SimpleTrainState]:
-        print("Generating states for SimpleTrainer")
-        rngs, subkey = jax.random.split(rngs)
-
-        if existing_state == None:
-            input_vars = self.get_input_ones()
-            params = model.init(subkey, **input_vars)
-        else:
-            params = existing_state['params']
-
-        state = SimpleTrainState.create(
-            apply_fn=model.apply,
-            params=params,
-            tx=optimizer,
-            metrics=Metrics.empty(),
-            dynamic_scale = dynamic_scale_lib.DynamicScale() if use_dynamic_scale else None
-        )
-        if existing_best_state is not None:
-            best_state = state.replace(
-                params=existing_best_state['params'])
-        else:
-            best_state = state
-            
-        return state, best_state
-
-    def init_state(
-        self,
-        state: SimpleTrainState,
-        best_state: SimpleTrainState,
-    ):
-        self.best_loss = 1e9
-
-        self.state = state
-        self.best_state = best_state
-
-    def get_state(self):
-        return self.get_np_tree(self.state)
-
-    def get_best_state(self):
-        return self.get_np_tree(self.best_state)
-        
-    def get_rngstate(self):
-        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())
-        if not os.path.exists(path):
-            os.makedirs(path)
-        return path
-
-    def tensorboard_path(self):
-        experiment_name = self.name
-        path = os.path.join(os.path.abspath('./tensorboard'), experiment_name)
-        if not os.path.exists(path):
-            os.makedirs(path)
-        return path
-
-    def load(self, checkpoint_path=None, checkpoint_step=None):
-        if checkpoint_path is None:
-            checkpointer = self.checkpointer
-        else:
-            checkpointer = orbax.checkpoint.PyTreeCheckpointer()
-            options = orbax.checkpoint.CheckpointManagerOptions(
-                max_to_keep=4, create=False)
-            checkpointer = orbax.checkpoint.CheckpointManager(
-                checkpoint_path, checkpointer, options)    
-        
-        if checkpoint_step is None:
-            step = checkpointer.latest_step()
-        else:
-            step = checkpoint_step
-        
-        print("Loading model from checkpoint at step ", step)
-        ckpt = checkpointer.restore(step)
-        state = ckpt['state']
-        best_state = ckpt['best_state']
-        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, state=None, rngstate=None):
-        print(f"Saving model at epoch {epoch} step {step}")
-        try:
-            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 outer", e)
-
-    def _define_train_step(self, **kwargs):
-        model = self.model
-        loss_fn = self.loss_fn
-        distributed_training = self.distributed_training
-
-        def train_step(train_state: SimpleTrainState, rng_state: RandomMarkovState, batch, local_device_indexes):
-            """Train for a single step."""
-            images = batch['image']
-            labels = batch['label']
-
-            def model_loss(params):
-                preds = model.apply(params, images)
-                expected_output = labels
-                nloss = loss_fn(preds, expected_output)
-                loss = jnp.mean(nloss)
-                return loss
-            loss, grads = jax.value_and_grad(model_loss)(train_state.params)
-            if distributed_training:
-                grads = jax.lax.pmean(grads, "data")
-            train_state = train_state.apply_gradients(grads=grads)
-            return train_state, loss, rng_state
-        
-        if distributed_training:
-            train_step = shard_map(train_step, mesh=self.mesh, in_specs=(P(), P(), P('data'), P('data')), out_specs=(P(), P('data'), P()))
-            train_step = jax.pmap(train_step)
-        return train_step
-
-    def _define_vaidation_step(self):
-        model = self.model
-        loss_fn = self.loss_fn
-        distributed_training = self.distributed_training
-
-        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
-        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()
-        print(self.model.tabulate(jax.random.key(0), **input_vars,
-              console_kwargs={"width": 200, "force_jupyter": True, }))
-
-    def config(self):
-        return {
-            "model": self.model,
-            "state": self.state,
-            "name": self.name,
-            "input_shapes": self.input_shapes
-        }
-
-    def init_tensorboard(self, batch_size, steps_per_epoch, epochs):
-        from flax.metrics import tensorboard
-        summary_writer = tensorboard.SummaryWriter(self.tensorboard_path())
-        summary_writer.hparams({
-            **self.config(),
-            "steps_per_epoch": steps_per_epoch,
-            "epochs": epochs,
-            "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 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()
-        if self.distributed_training:
-            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)
-
-            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'))
-                # 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 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
-            
-            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 / 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
-                self.save(current_epoch, current_step)
-                
-            if process_index == 0:
-                if self.wandb is not None:
-                    self.wandb.log({
-                        "train/epoch_time": total_time,
-                        "train/avg_time_per_step": avg_time_per_step,
-                        "train/avg_loss": avg_loss,
-                        "train/best_loss": self.best_loss,
-                        "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'))
-                    
-            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

@@ -1,62 +0,0 @@
-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-0.1.36.dist-info/RECORD

@@ -1,43 +0,0 @@
-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,,