lint

Author: Juan Acevedo

src/maxdiffusion/data_preprocessing/__init__.py

@@ -12,4 +12,4 @@
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
- """
+ """

src/maxdiffusion/data_preprocessing/wan_txt2vid_data_preprocessing.py

@@ -34,6 +34,7 @@
 
 import tensorflow as tf
 
+
 def image_feature(value):
   """Returns a bytes_list from a string / byte."""
   return tf.train.Feature(bytes_list=tf.train.BytesList(value=[tf.io.encode_jpeg(value).numpy()]))
@@ -58,6 +59,7 @@ def float_feature_list(value):
   """Returns a list of float_list from a float / double."""
   return tf.train.Feature(float_list=tf.train.FloatList(value=value))
 
+
 def create_example(latent, hidden_states):
   latent = tf.io.serialize_tensor(latent)
   hidden_states = tf.io.serialize_tensor(hidden_states)
@@ -74,11 +76,13 @@ def text_encode(pipeline, prompt: Union[str, List[str]]):
   encoder_hidden_states = encoder_hidden_states.detach().numpy()
   return encoder_hidden_states
 
+
 def vae_encode(video, rng, vae, vae_cache):
   latent = vae.encode(video, feat_cache=vae_cache)
   latent = latent.latent_dist.sample(rng)
   return latent
-  
+
+
 def generate_dataset(config, pipeline):
 
   tfrecords_dir = config.tfrecords_dir
@@ -99,21 +103,21 @@ def generate_dataset(config, pipeline):
   rng = jax.random.key(config.seed)
 
   vae_scale_factor_spatial = 2 ** len(pipeline.vae.temperal_downsample)
-  video_processor = VideoProcessor(vae_scale_factor=vae_scale_factor_spatial) 
-  
+  video_processor = VideoProcessor(vae_scale_factor=vae_scale_factor_spatial)
+
   # jit vae fun.
   p_vae_encode = jax.jit(functools.partial(vae_encode, vae=pipeline.vae, vae_cache=pipeline.vae_cache))
-  
+
   # Load dataset
-  ds = load_dataset(config.dataset_name, split='train')
+  ds = load_dataset(config.dataset_name, split="train")
   ds = ds.shuffle(seed=config.seed)
   ds = ds.select_columns([config.caption_column, config.image_column])
   batch_size = 10
   for i in range(0, len(ds), batch_size):
     rng, new_rng = jax.random.split(rng)
-    text = ds[i:i+batch_size]['text']
-    videos = ds[i:i+batch_size]['image']
-    
+    text = ds[i : i + batch_size]["text"]
+    videos = ds[i : i + batch_size]["image"]
+
     videos = [video_processor.preprocess_video([video], height=config.height, width=config.width) for video in videos]
     video = jnp.array(np.squeeze(np.array(videos), axis=1), dtype=config.weights_dtype)
     with mesh:
@@ -127,24 +131,23 @@ def generate_dataset(config, pipeline):
 
     if shard_record_count >= no_records_per_shard:
       writer.close()
-      tf_rec_num +=1
+      tf_rec_num += 1
       writer = tf.io.TFRecordWriter(
           tfrecords_dir + "/file_%.2i-%i.tfrec" % (tf_rec_num, (global_record_count + no_records_per_shard))
       )
       shard_record_count = 0
 
 
-
 def run(config):
   pipeline = WanPipeline.from_pretrained(config, load_transformer=False)
   # Don't need the transformer for preprocessing.
   generate_dataset(config, pipeline)
 
 
-
 def main(argv: Sequence[str]) -> None:
   pyconfig.initialize(argv)
   run(pyconfig.config)
 
+
 if __name__ == "__main__":
-  app.run(main)
+  app.run(main)

src/maxdiffusion/generate_wan.py

@@ -21,7 +21,7 @@
 from maxdiffusion.utils import export_to_video
 
 
-def run(config, pipeline=None, filename_prefix=''):
+def run(config, pipeline=None, filename_prefix=""):
   print("seed: ", config.seed)
   if pipeline is None:
     pipeline = WanPipeline.from_pretrained(config)

src/maxdiffusion/input_pipeline/_tfds_data_processing.py

@@ -73,14 +73,9 @@ def make_tf_iterator(
   train_iter = multihost_dataloading.MultiHostDataLoadIterator(train_ds, mesh)
   return train_iter
 
+
 def make_cached_tfrecord_iterator(
-    config,
-    dataloading_host_index,
-    dataloading_host_count,
-    mesh,
-    global_batch_size,
-    feature_description,
-    prepare_sample_fn
+    config, dataloading_host_index, dataloading_host_count, mesh, global_batch_size, feature_description, prepare_sample_fn
 ):
   """
   New iterator for TFRecords that contain the full 4 pre-computed latents and embeddings:
@@ -111,13 +106,7 @@ def _parse_tfrecord_fn(example):
 
 # TODO - https://github.com/google/array_record/blob/main/beam/examples/example_gcs_conversion.py
 def make_tfrecord_iterator(
-    config,
-    dataloading_host_index,
-    dataloading_host_count,
-    mesh,
-    global_batch_size,
-    feature_description,
-    prepare_sample_fn
+    config, dataloading_host_index, dataloading_host_count, mesh, global_batch_size, feature_description, prepare_sample_fn
 ):
   """Iterator for TFRecord format. For Laion dataset,
   check out preparation script
@@ -127,18 +116,20 @@ def make_tfrecord_iterator(
   # set load_tfrecord_cached to True in config to use pre-processed tfrecord dataset.
   # pedagogical_examples/dataset_tf_cache_to_tfrecord.py to convert tf preprocessed dataset to tfrecord.
   # Dataset cache in github runner test doesn't contain all the features since its shared, Use the default tfrecord iterator.
-  if (config.cache_latents_text_encoder_outputs
+  if (
+      config.cache_latents_text_encoder_outputs
       and os.path.isdir(config.dataset_save_location)
-      and 'load_tfrecord_cached'in config.get_keys()
-      and config.load_tfrecord_cached):
+      and "load_tfrecord_cached" in config.get_keys()
+      and config.load_tfrecord_cached
+  ):
     return make_cached_tfrecord_iterator(
-      config,
-      dataloading_host_index,
-      dataloading_host_count,
-      mesh,
-      global_batch_size,
-      feature_description,
-      prepare_sample_fn
+        config,
+        dataloading_host_index,
+        dataloading_host_count,
+        mesh,
+        global_batch_size,
+        feature_description,
+        prepare_sample_fn,
     )
 
   feature_description = {

src/maxdiffusion/input_pipeline/input_pipeline_interface.py

@@ -51,16 +51,23 @@ def make_data_iterator(
     tokenize_fn=None,
     image_transforms_fn=None,
     feature_description=None,
-    prepare_sample_fn=None
+    prepare_sample_fn=None,
 ):
   """Make data iterator for SD1, 2, XL, dataset_types in (hf, tf, tfrecord)"""
-  
+
   if config.dataset_type == "hf" or config.dataset_type == "tf":
     if tokenize_fn is None or image_transforms_fn is None:
       raise ValueError(f"dataset type {config.dataset_type} needs to pass a tokenize_fn and image_transforms_fn")
-  
-  if config.dataset_type == "tfrecord" and config.cache_latents_text_encoder_outputs and feature_description is None or prepare_sample_fn is None:
-    raise ValueError(f"dataset type {config.dataset_type} needs to pass a feature_description dictionary and prepare_sample_fn function when cache_latents_text_encoder_outputs is True.")
+
+  if (
+      config.dataset_type == "tfrecord"
+      and config.cache_latents_text_encoder_outputs
+      and feature_description is None
+      or prepare_sample_fn is None
+  ):
+    raise ValueError(
+        f"dataset type {config.dataset_type} needs to pass a feature_description dictionary and prepare_sample_fn function when cache_latents_text_encoder_outputs is True."
+    )
 
   if config.dataset_type == "hf":
     return _hf_data_processing.make_hf_streaming_iterator(
@@ -98,7 +105,7 @@ def make_data_iterator(
         mesh,
         global_batch_size,
         feature_description,
-        prepare_sample_fn
+        prepare_sample_fn,
     )
   else:
     assert False, f"Unknown dataset_type {config.dataset_type}, dataset_type must be in (tf, tfrecord, hf, grain)"

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -41,6 +41,7 @@
 def basic_clean(text):
   if is_ftfy_available():
     import ftfy
+
     text = ftfy.fix_text(text)
   text = html.unescape(html.unescape(text))
   return text.strip()
@@ -398,7 +399,7 @@ def __call__(
             num_channels_latents=num_channel_latents,
         )
 
-      data_sharding = NamedSharding(self.devices_array, P())
+      data_sharding = NamedSharding(self.mesh, P())
       if len(prompt) % jax.device_count() == 0:
         data_sharding = jax.sharding.NamedSharding(self.mesh, P(*self.config.data_sharding))