flaxdiff 0.1.38__py3-none-any.whl → 0.2.0__py3-none-any.whl

flaxdiff/trainer/simple_trainer.py

@@ -25,6 +25,8 @@ 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
+from dataclasses import dataclass
+import gc
 
 PROCESS_COLOR_MAP = {
     0: "green",
@@ -71,6 +73,7 @@ class SimpleTrainState(train_state.TrainState):
     metrics: Metrics
     dynamic_scale: dynamic_scale_lib.DynamicScale
 
+@dataclass
 class SimpleTrainer:
     state: SimpleTrainState
     best_state: SimpleTrainState
@@ -86,7 +89,6 @@ class SimpleTrainer:
                  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,
@@ -94,6 +96,7 @@ class SimpleTrainer:
                  checkpoint_base_path: str = "./checkpoints",
                  checkpoint_step: int = None,
                  use_dynamic_scale: bool = False,
+                 max_checkpoints_to_keep: int = 2,
                  ):
         if distributed_training is None or distributed_training is True:
             # Auto-detect if we are running on multiple devices
@@ -109,10 +112,9 @@ class SimpleTrainer:
         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)
+            run = wandb.init(resume='allow', **wandb_config)
             self.wandb = run
             
             # define our custom x axis metric
@@ -126,13 +128,18 @@ class SimpleTrainer:
             self.wandb.define_metric("train/avg_loss", step_metric="train/epoch")
             self.wandb.define_metric("train/best_loss", step_metric="train/epoch")
             
+            if self.wandb.sweep_id:
+                api = wandb.Api()
+                self.wandb_sweep = api.sweep(f"{self.wandb.entity}/{self.wandb.project}/{self.wandb.sweep_id}")
+                print(f"Running sweep {self.wandb_sweep.id} with id {self.wandb.sweep_id}")
+            
         # 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)
+            max_to_keep=max_checkpoints_to_keep, create=True)
         self.checkpointer = orbax.checkpoint.CheckpointManager(
-            self.checkpoint_path() + checkpoint_suffix, async_checkpointer, options)
+            self.checkpoint_path(), 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)
@@ -248,6 +255,10 @@ class SimpleTrainer:
             step = checkpoint_step
         
         print("Loading model from checkpoint at step ", step)
+        loaded_checkpoint_path = os.path.join(
+            checkpoint_path if checkpoint_path else self.checkpoint_path(),
+            f"{step}")
+        self.loaded_checkpoint_path = loaded_checkpoint_path
         ckpt = checkpointer.restore(step)
         state = ckpt['state']
         best_state = ckpt['best_state']
@@ -311,7 +322,7 @@ class SimpleTrainer:
             train_step = jax.pmap(train_step)
         return train_step
 
-    def _define_vaidation_step(self):
+    def _define_validation_step(self):
         model = self.model
         loss_fn = self.loss_fn
         distributed_training = self.distributed_training
@@ -418,8 +429,8 @@ class SimpleTrainer:
             
         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 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 
@@ -433,16 +444,40 @@ class SimpleTrainer:
                 # loss = jax.experimental.multihost_utils.process_allgather(loss)
                 loss = jnp.mean(loss) # Just to make sure its a scaler value
                     
-            if loss <= 1e-8:
-                # 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
-                if self.best_state is not None:
-                    print(colored(f"Resetting model to best state", 'red'))
+            if loss <= 1e-8 or jnp.isnan(loss) or jnp.isinf(loss):
+                # If the loss is too low or NaN/Inf, log the issue and attempt recovery
+                print(colored(f"Abnormal loss at step {current_step}: {loss}", 'red'))
+                
+                # Check model parameters for NaN/Inf values
+                params = train_state.params
+                has_nan_or_inf = False
+                
+                if isinstance(params, dict):
+                    for key, value in params.items():
+                        if isinstance(value, jnp.ndarray):
+                            if jnp.isnan(value).any() or jnp.isinf(value).any():
+                                print(colored(f"NaN/inf values found in params[{key}] at step {current_step}", 'red'))
+                                has_nan_or_inf = True
+                                break
+                    
+                    if not has_nan_or_inf:
+                        print(colored(f"Model parameters seem valid despite abnormal loss", 'yellow'))
+                
+                # Try to recover - clear JAX caches and collect garbage
+                gc.collect()
+                if hasattr(jax, "clear_caches"):
+                    jax.clear_caches()
+                
+                # If we have a best state and the loss is truly invalid, consider restoring
+                if (loss <= 1e-8 or jnp.isnan(loss) or jnp.isinf(loss)) and self.best_state is not None:
+                    print(colored(f"Attempting recovery by resetting model to last best state", 'yellow'))
                     train_state = self.best_state
                     loss = self.best_loss
                 else:
-                    exit(1)
+                    # If we can't recover, skip this step but continue training
+                    print(colored(f"Unable to recover - continuing with current state", 'yellow'))
+                    if loss <= 1e-8:
+                        loss = 1.0  # Set to a reasonable default to continue training
                             
             epoch_loss += loss
             current_step += 1
@@ -471,7 +506,7 @@ class SimpleTrainer:
     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)
+        val_step = self._define_validation_step(**validation_step_args)
         train_state = self.state
         rng_state = self.rngstate
         process_index = jax.process_index()

flaxdiff/utils.py

@@ -2,26 +2,145 @@ import jax
 import jax.numpy as jnp
 import flax.struct as struct
 import flax.linen as nn
-from typing import Any, Callable
-from dataclasses import dataclass
+from typing import Any
 from functools import partial
 import numpy as np
+import os
 from jax.sharding import Mesh, PartitionSpec as P
-from abc import ABC, abstractmethod
+from flaxdiff.inputs import TextEncoder, CLIPTextEncoder
+
+# Setup mappings for dtype, precision, and activation
+DTYPE_MAP = {
+    'bfloat16': jnp.bfloat16,
+    'float32': jnp.float32,
+    'jax.numpy.float32': jnp.float32,
+    'jax.numpy.bfloat16': jnp.bfloat16,
+    'None': None,
+    None: None,
+}
+
+PRECISION_MAP = {
+    'high': jax.lax.Precision.HIGH,
+    'HIGH': jax.lax.Precision.HIGH,
+    'default': jax.lax.Precision.DEFAULT,
+    'DEFAULT': jax.lax.Precision.DEFAULT,
+    'highest': jax.lax.Precision.HIGHEST,
+    'HIGHEST': jax.lax.Precision.HIGHEST,
+    'None': None,
+    None: None,
+}
+
+ACTIVATION_MAP = {
+    'swish': jax.nn.swish,
+    'silu': jax.nn.silu,
+    'jax._src.nn.functions.silu': jax.nn.silu,
+    'mish': jax.nn.mish,
+}
+
+def map_nested_config(config):
+    new_config = {}
+    for key, value in config.items():
+        if isinstance(value, dict):
+            new_config[key] = map_nested_config(value)
+        elif isinstance(value, str):
+            if value in DTYPE_MAP:
+                new_config[key] = DTYPE_MAP[value]
+            elif value in PRECISION_MAP:
+                new_config[key] = PRECISION_MAP[value]
+            elif value in ACTIVATION_MAP:
+                new_config[key] = ACTIVATION_MAP[value]
+            elif value == 'None':
+                new_config[key] = None
+            elif '.' in value:
+                # Ignore any other string that contains a dot
+                print(
+                    f"Ignoring key {key} with value {value} as it contains a dot.")
+    return new_config
+
+def serialize_model(model: nn.Module):
+    """
+    Serializes the model to a dictionary format.
+    """
+    model_dict = model.__dict__
+    model_dict = {k: v for k, v in model_dict.items() if not k.startswith('_')}
+    # Convert all callable attributes to their string representation
+    def map(model_dict):
+        for k, v in model_dict.items():
+            if isinstance(v, dict):
+                # Recursively serialize nested dictionaries
+                model_dict[k] = map(v)
+            elif isinstance(v, list):
+                # Recursively serialize lists
+                [map(item) if isinstance(item, dict) else item for item in v]
+            elif callable(v):
+                # If the attribute has __name__, use that as the key
+                if hasattr(v, '__name__'):
+                    model_dict[k] = v.__name__
+                else:
+                    model_dict[k] = str(v).split('.')[-1]
+    map(model_dict)
+    return model_dict
+
+def get_latest_checkpoint(checkpoint_path):
+    checkpoint_files = os.listdir(checkpoint_path)
+    # Sort files by step number
+    checkpoint_files = sorted([int(i) for i in checkpoint_files])
+    latest_step = checkpoint_files[-1]
+    latest_checkpoint = os.path.join(checkpoint_path, str(latest_step))
+    return latest_checkpoint
 
 class MarkovState(struct.PyTreeNode):
     pass
 
 class RandomMarkovState(MarkovState):
     rng: jax.random.PRNGKey
-
     def get_random_key(self):
         rng, subkey = jax.random.split(self.rng)
         return RandomMarkovState(rng), subkey
 
 def clip_images(images, clip_min=-1, clip_max=1):
+    """Clip image values to a specified range.
+    
+    Args:
+        images: Images to clip
+        clip_min: Minimum value
+        clip_max: Maximum value
+    
+    Returns:
+        Clipped images
+    """
     return jnp.clip(images, clip_min, clip_max)
 
+def denormalize_images(images, target_type=jnp.uint8, source_range=(-1, 1), target_range=(0, 255)):
+    """Convert images from normalized range (e.g. [-1, 1]) to target range (e.g. [0, 255]).
+    
+    Args:
+        images: Normalized images
+        target_type: Target dtype (e.g. jnp.uint8 for standard images)
+        source_range: Tuple of (min, max) for the source normalization range
+        target_range: Tuple of (min, max) for the target range
+        
+    Returns:
+        Denormalized images in the target dtype
+    """
+    src_min, src_max = source_range
+    tgt_min, tgt_max = target_range
+    
+    # First clip to ensure we're in the expected source range
+    images = clip_images(images, src_min, src_max)
+    
+    # Scale to [0, 1]
+    images = (images - src_min) / (src_max - src_min)
+    
+    # Scale to target range
+    images = images * (tgt_max - tgt_min) + tgt_min
+    
+    # Convert to target dtype if needed
+    if target_type is not None:
+        images = images.astype(target_type)
+    
+    return images
+
 def _build_global_shape_and_sharding(
     local_shape: tuple[int, ...], global_mesh: Mesh
 ) -> tuple[tuple[int, ...], jax.sharding.NamedSharding]:
@@ -117,45 +236,6 @@ class RMSNorm(nn.Module):
         y = mul * x
         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
@@ -175,18 +255,9 @@ class AutoTextTokenizer:
 
     def __repr__(self):
         return self.__class__.__name__ + '()'
+    
+# class AutoAudioTokenizer:
 
-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)
+def defaultTextEncodeModel(modelname = "openai/clip-vit-large-patch14", backend="jax"):
+    """Default text encoder model."""
+    return CLIPTextEncoder.from_modelname(modelname=modelname, backend=backend)

{flaxdiff-0.1.38.dist-info → flaxdiff-0.2.0.dist-info}/METADATA

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

flaxdiff-0.2.0.dist-info/RECORD

@@ -0,0 +1,64 @@
+flaxdiff/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flaxdiff/utils.py,sha256=DmlWUY1FGz4ESxIHaPQJf92CHjsdMjyDd651wFUtyNg,8838
+flaxdiff/data/__init__.py,sha256=8W5y7NyAOWtpLi8WRawk4VYeE3DMDnM3B_jKPD8BoFQ,143
+flaxdiff/data/benchmark_decord.py,sha256=x56Db1VPmziv_9KJvWdfS0O7cffsYkF5tt5WvldOKc0,13720
+flaxdiff/data/dataloaders.py,sha256=V4goNCK0JD_TthggXAEgJJD4LxJi1pUDew1x_fMCuO4,22576
+flaxdiff/data/dataset_map.py,sha256=NrLG1XtIxy8GcCsZ-e6eascjgsP0Xq5lVA1z3HIIYyI,5093
+flaxdiff/data/online_loader.py,sha256=t1jEhdB6gWTlwx68ehj1ol_PrImbwXYiRlrJPCmNgCM,35701
+flaxdiff/data/sources/audio_utils.py,sha256=X27gG1yQt_abVOYgMtruYmZD7-8_uQCRhhTSpn4clkI,4514
+flaxdiff/data/sources/av_example.py,sha256=RIcbVKqckFqbfnV65NQotzIBxjdDuM67kD1nY8fqw5Q,3826
+flaxdiff/data/sources/av_utils.py,sha256=n2qwMBQGouoBH025vdE7gitWC6RduUommUrs-SPdWe4,24041
+flaxdiff/data/sources/base.py,sha256=uhF0odJSYRy0SLw1xnI9Q_q_xiVht2DmEYcX1j9AWT4,4246
+flaxdiff/data/sources/images.py,sha256=WpH4ywZhNol26peX3m6m5NrmDJ1K2s6fRcYHvOFlOk8,11102
+flaxdiff/data/sources/utils.py,sha256=kFzM4_kPoThbAu54ulABmEDAR33tR50NgzXIpC0Dzjk,7316
+flaxdiff/data/sources/videos.py,sha256=CVpOH6A4P2D8iv3gZIhd2GB5ATUD8Vsm_wVYbbugWD4,9359
+flaxdiff/data/sources/voxceleb2.py,sha256=BoKfat_hsw6ObDyyaiQmPbBzuFiqgCGlgAZmf-t5Iz8,18621
+flaxdiff/inference/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flaxdiff/inference/pipeline.py,sha256=oMBRjvTtlC3Yzl1FqiBHcI4V34HXGAecCg8UvQbKoOc,8849
+flaxdiff/inference/utils.py,sha256=SRNYo-YtHzEPRpNv0fD8ZrUvnRIK941Rh4tjlsOGRgM,12278
+flaxdiff/inputs/__init__.py,sha256=ybPjQsFAf5sqRVZG1sRiOl99EnwpI-NQ8HE3y7UbXmU,7197
+flaxdiff/inputs/encoders.py,sha256=pjfbx4Rk7bLoE80MOfThZDm6YtsDncRekmn0Bmg_CwI,2963
+flaxdiff/metrics/inception.py,sha256=a5kjMCPMT9gB88c_HCKiek-2vsAyoE35K7nDt4h4pVI,31843
+flaxdiff/metrics/psnr.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flaxdiff/metrics/ssim.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flaxdiff/metrics/utils.py,sha256=YuuOfqvqgIjsceupwNeJ59vQ2TnGeNMIyKdkIqOmoNg,1702
+flaxdiff/models/__init__.py,sha256=amtDF07DfiAdnZsvWX4eaW79nwNEU1s8Zb4PB3ewtg4,118
+flaxdiff/models/attention.py,sha256=-q3xqWy4vQSLG4vXtiUN3FHVBIo7ZjpQsdLT9CkML6c,13367
+flaxdiff/models/common.py,sha256=7x9o5vY9UZvN4BNZ7LHzyuU3PNpsNym9B3m1Wfdddjo,10320
+flaxdiff/models/favor_fastattn.py,sha256=79Ew1nqarsNLPzZaBSd1ILORzJr74CupYeqGiCQK5E4,27689
+flaxdiff/models/general.py,sha256=7xMME6KVKQY8sScyHYH4f-Kek4j1pRfplKShFXwVZd4,587
+flaxdiff/models/simple_unet.py,sha256=pjeixszG_6gEY5PNFbQ7KbOyg4z5bfn4RUbINCJexOM,10758
+flaxdiff/models/simple_vit.py,sha256=6DNpwTeE0Gn2jSie6n0JVUmQncPoyFT7jSSBreqk458,7497
+flaxdiff/models/unet_3d.py,sha256=LF0PMxBKGU-_lAMtO_Coxy1yRE02yKKdgb7i6YZxI_4,20163
+flaxdiff/models/unet_3d_blocks.py,sha256=lRYDc9X1VEu54Kg7xEEphXYiQ09tabPXKi-hEcKFYug,19687
+flaxdiff/models/autoencoder/__init__.py,sha256=qY-7MldZpsfkF-_T2LqlRK7VHbqfmosz0NmvzDlBkOk,78
+flaxdiff/models/autoencoder/autoencoder.py,sha256=8XWdsWvsPsyWGtzpCT8w0KXi_ZLGpRuQpn4oXo1gHKw,6039
+flaxdiff/models/autoencoder/diffusers.py,sha256=tPz77YuctrT--jF2AOL8G6vr0NiIr3RXANNrZCxe0bg,5921
+flaxdiff/models/autoencoder/simple_autoenc.py,sha256=NnGFjrkq-1z8Ouh_UvlvP0PFpkzm2LYaTQKMmN1BhkM,2109
+flaxdiff/predictors/__init__.py,sha256=A6lMnRSqDLwBdwB0a1x4xPpcXcTYCISWFMkldbiK6Vs,4554
+flaxdiff/samplers/__init__.py,sha256=EY9v1pgwEoR64Kiz9K8fAR-4_ir9c03mYeY3hrpUNhE,308
+flaxdiff/samplers/common.py,sha256=-uU9FLkoQp3n3bga8Kfj_onDhtOS4MBggIKNrq3S8n4,18438
+flaxdiff/samplers/ddim.py,sha256=iFgXz96NBYuNMiWGMovf3gLO2TCxdrhJd-o8tvSmVUI,2054
+flaxdiff/samplers/ddpm.py,sha256=JgkNSo7fp7Jm-8rCy4eu5m4YIzXTWzxv-iHf3EQ0z5w,2243
+flaxdiff/samplers/euler.py,sha256=QSkttB4DYnepDGwhWq3EGXYjMAqj4qLOdh7u98HttzY,2791
+flaxdiff/samplers/heun_sampler.py,sha256=EvR3hy4t_D47ZOH4luzRFqPmv2v4z78P_JhqBGEpHU8,1436
+flaxdiff/samplers/multistep_dpm.py,sha256=2M4Abb93-GUVN1f0_ZHBeA6lF0eF15Hi6QOgOu2K45s,2752
+flaxdiff/samplers/rk4_sampler.py,sha256=vcQefFhOUZdNOQGBdzNkb2NgmTC2KWd_nhUhyLtt3yI,2026
+flaxdiff/schedulers/__init__.py,sha256=EIva9gBz3DKHORuGmv1LQCKTtRqCRavFOXMNqxAR_ks,131
+flaxdiff/schedulers/common.py,sha256=PDeje2NmN7X3J5qKGauE0jYPpxjgEX44f_evJHRIG3E,4382
+flaxdiff/schedulers/continuous.py,sha256=5c_niOA20fxJ5oJDi09FfayIRogBGwtfG0XThW2IUZk,334
+flaxdiff/schedulers/cosine.py,sha256=E5pODAmINfdyC4kSYOJSPAvq3GNlKPpKEn3X82vYMz0,2055
+flaxdiff/schedulers/discrete.py,sha256=m1q3bAgeAxU3gTj5di3XFWDm4yLfMKAFJPlYdozLE2Y,3316
+flaxdiff/schedulers/exp.py,sha256=cPTnUJpYdzJRRZqMLYQz0rRUCpEmaP2tXhRumLx94jA,605
+flaxdiff/schedulers/karras.py,sha256=7PS6mHdnZnTqS2Xl_DacBt5YQ1f_CFyAxShyOo55eG0,3804
+flaxdiff/schedulers/linear.py,sha256=pBDTXSQcOS4Z03JTh6S0f9E2qLcTQzF2E-pGoQnRoy0,572
+flaxdiff/schedulers/sqrt.py,sha256=mCd_szmOqF6vqQKiAiEOqV_3eBIPGYrW3VxK0o4rBuo,474
+flaxdiff/trainer/__init__.py,sha256=xSoierfi26gxfgxlNnwvyyPmuPAJ--5i3mEHxt3S-AE,215
+flaxdiff/trainer/autoencoder_trainer.py,sha256=2FP2P-k9c0n_k3eT0trkq73dQrHRdBj9ObK1idcyhSw,6996
+flaxdiff/trainer/diffusion_trainer.py,sha256=reQEVWKTqKAeyCMQ-curPOfSRmBKxKooK8EVtUuorcM,14599
+flaxdiff/trainer/general_diffusion_trainer.py,sha256=VQ5p2ZaTv2R1LM0Epz4e719_EfK2dh1eoKK3WIysIW0,24040
+flaxdiff/trainer/simple_trainer.py,sha256=CF2mMcc6AtBgcR1XiqKevRL0paGS0S9ZJofCns32nRM,24214
+flaxdiff-0.2.0.dist-info/METADATA,sha256=1WLpd9RQy_mJE2E2uOdXptY5Fm3n_MTNcgZyBD7YmGw,23982
+flaxdiff-0.2.0.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
+flaxdiff-0.2.0.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.2.0.dist-info/RECORD,,

{flaxdiff-0.1.38.dist-info → flaxdiff-0.2.0.dist-info}/WHEEL

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

flaxdiff/data/datasets.py

@@ -1,169 +0,0 @@
-import jax.numpy as jnp
-import grain.python as pygrain
-from typing import Dict
-import numpy as np
-import jax
-from flaxdiff.utils import convert_to_global_tree, AutoTextTokenizer
-from .dataset_map import datasetMap, onlineDatasetMap
-import traceback
-from .online_loader import OnlineStreamingDataLoader
-import queue
-from jax.sharding import Mesh
-import threading
-
-def batch_mesh_map(mesh):
-    class augmenters(pygrain.MapTransform):
-        def __init__(self, *args, **kwargs):
-            super().__init__(*args, **kwargs)
-
-        def map(self, batch) -> Dict[str, jnp.array]:
-            return convert_to_global_tree(mesh, batch)
-    return augmenters
-
-def get_dataset_grain(
-    data_name="cc12m",
-    batch_size=64,
-    image_scale=256,
-    count=None,
-    num_epochs=None,
-    method=jax.image.ResizeMethod.LANCZOS3,
-    worker_count=32,
-    read_thread_count=64,
-    read_buffer_size=50,
-    worker_buffer_size=20,
-    seed=0,
-    dataset_source="/mnt/gcs_mount/arrayrecord2/cc12m/",
-):
-    dataset = datasetMap[data_name]
-    data_source = dataset["source"](dataset_source)
-    augmenter = dataset["augmenter"](image_scale, method)
-
-    local_batch_size = batch_size // jax.process_count()
-
-    sampler = pygrain.IndexSampler(
-        num_records=len(data_source) if count is None else count,
-        shuffle=True,
-        seed=seed,
-        num_epochs=num_epochs,
-        shard_options=pygrain.ShardByJaxProcess(),
-    )
-
-    def get_trainset():
-        transformations = [
-            augmenter(),
-            pygrain.Batch(local_batch_size, drop_remainder=True),
-        ]
-        
-        # if mesh != None:
-        #     transformations += [batch_mesh_map(mesh)]
-
-        loader = pygrain.DataLoader(
-            data_source=data_source,
-            sampler=sampler,
-            operations=transformations,
-            worker_count=worker_count,
-            read_options=pygrain.ReadOptions(
-                read_thread_count, read_buffer_size
-            ),
-            worker_buffer_size=worker_buffer_size,
-        )
-        return loader
-    
-
-    return {
-        "train": get_trainset,
-        "train_len": len(data_source),
-        "local_batch_size": local_batch_size,
-        "global_batch_size": batch_size,
-        # "null_labels": null_labels,
-        # "null_labels_full": null_labels_full,
-        # "model": model,
-        # "tokenizer": tokenizer,
-    }
-
-def generate_collate_fn():
-    auto_tokenize = AutoTextTokenizer(tensor_type="np")
-    def default_collate(batch):
-        try:
-            # urls = [sample["url"] for sample in batch]
-            captions = [sample["caption"] for sample in batch]
-            results = auto_tokenize(captions)
-            images = np.stack([sample["image"] for sample in batch], axis=0)
-            return {
-                "image": images,
-                "input_ids": results['input_ids'],
-                "attention_mask": results['attention_mask'],
-            }
-        except Exception as e:
-            print("Error in collate function", e, [sample["image"].shape for sample in batch])
-            traceback.print_exc()
-            
-    return default_collate
-    
-def get_dataset_online(
-        data_name="combined_online",
-        batch_size=64,
-        image_scale=256,
-        count=None,
-        num_epochs=None,
-        method=jax.image.ResizeMethod.LANCZOS3,
-        worker_count=32,
-        read_thread_count=64,
-        read_buffer_size=50,
-        worker_buffer_size=20,
-        seed=0,
-        dataset_source="/mnt/gcs_mount/arrayrecord2/cc12m/",
-    ):
-    local_batch_size = batch_size // jax.process_count()
-    
-    sources = onlineDatasetMap[data_name]["source"]
-    dataloader = OnlineStreamingDataLoader(
-            sources, 
-            batch_size=local_batch_size,
-            num_workers=worker_count,
-            num_threads=read_thread_count,
-            image_shape=(image_scale, image_scale),
-            global_process_count=jax.process_count(),
-            global_process_index=jax.process_index(),
-            prefetch=worker_buffer_size,
-            collate_fn=generate_collate_fn(),
-            default_split="train",
-        )
-    
-    def get_trainset(mesh: Mesh = None):
-        if mesh != None:
-            class dataLoaderWithMesh:
-                def __init__(self, dataloader, mesh):
-                    self.dataloader = dataloader
-                    self.mesh = mesh
-                    self.tmp_queue = queue.Queue(worker_buffer_size)
-                    def batch_loader():
-                        for batch in self.dataloader:
-                            try:
-                                self.tmp_queue.put(convert_to_global_tree(mesh, batch))
-                            except Exception as e:
-                                print("Error processing batch", e)
-                    self.loader_thread = threading.Thread(target=batch_loader)
-                    self.loader_thread.start()
-                    
-                def __iter__(self):
-                    return self
-                
-                def __next__(self):
-                    return self.tmp_queue.get()
-                
-            dataloader_with_mesh = dataLoaderWithMesh(dataloader, mesh)
-                    
-            return dataloader_with_mesh
-        return dataloader
-    
-    return {
-        "train": get_trainset,
-        "train_len": len(dataloader) * jax.process_count(),
-        "local_batch_size": local_batch_size,
-        "global_batch_size": batch_size,
-        # "null_labels": null_labels,
-        # "null_labels_full": null_labels_full,
-        # "model": model,
-        # "tokenizer": tokenizer,
-    }