Fallback to regular tfrecord iterator for datasets without all the processed features

Signed-off-by: Kunjan <kunjanp@google.com>

Author: coolkp

src/maxdiffusion/input_pipeline/_tfds_data_processing.py

@@ -19,7 +19,7 @@
 import tensorflow.experimental.numpy as tnp
 from datasets import load_dataset, load_from_disk
 
-from maxdiffusion import multihost_dataloading
+from maxdiffusion import multihost_dataloading, max_logging
 
 AUTOTUNE = tf.data.AUTOTUNE
 
@@ -134,7 +134,11 @@ def make_tfrecord_iterator(
   maxdiffusion/pedagogical_examples/to_tfrecords.py
   """
   if config.cache_latents_text_encoder_outputs and os.path.isdir(config.dataset_save_location):
-    return make_cached_tfrecord_iterator(config, dataloading_host_index, dataloading_host_count, mesh, global_batch_size)
+    try:
+      return make_cached_tfrecord_iterator(config, dataloading_host_index, dataloading_host_count, mesh, global_batch_size)
+    except Exception as e:
+      max_logging.log(f"Cached tfrecord dataset doesn't contain required features, making regular tfrecord_iterator {e}")
+      
   feature_description = {
       "moments": tf.io.FixedLenFeature([], tf.string),
       "clip_embeddings": tf.io.FixedLenFeature([], tf.string),