fixes failing smoke tests

README.md

@@ -188,7 +188,7 @@ MaxDiffusion started as a fork of [Diffusers](https://github.com/huggingface/dif
 
 Whether you are forking MaxDiffusion for your own needs or intending to contribute back to the community, a full suite of tests can be found in `tests` and `src/maxdiffusion/tests`.
 
-To run unit tests, simply run:
+To run unit tests simply run:
 ```
 python -m pytest
 ```

docs/getting_started/first_run.md

@@ -12,9 +12,7 @@ multiple hosts.
 1. Clone MaxDiffusion in your TPU VM.
 1. Within the root directory of the MaxDiffusion `git` repo, install dependencies by running:
 ```bash
-pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-pip3 install -r requirements.txt
-pip3 install .
+bash setup.sh MODE=stable
 ```
 
 ## Getting Starting: Multihost development

setup.sh

@@ -108,3 +108,6 @@ fi
 
 # Install dependencies from requirements.txt
 pip3 install -U -r requirements.txt
+
+# Install maxdiffusion
+pip3 install -U .

src/maxdiffusion/configs/base_xl_lightning.yml

@@ -110,6 +110,7 @@ ici_tensor_parallelism: 1
 # Dataset
 # Replace with dataset path or train_data_dir. One has to be set.
 dataset_name: ''
+dataset_type: 'tf'
 train_data_dir: ''
 dataset_config_name: ''
 jax_cache_dir: ''

src/maxdiffusion/max_utils.py

@@ -121,7 +121,7 @@ def write_metrics_for_gcs(metrics, step, config, running_metrics):
   """Writes metrics to gcs"""
   metrics_dict_step = _prepare_metrics_for_json(metrics, step, config.run_name)
   running_metrics.append(metrics_dict_step)
-  if (step + 1) % config.log_period == 0 or step == config.steps - 1:
+  if (step + 1) % config.log_period == 0 or step == config.max_train_steps - 1:
     start_step = (step // config.log_period) * config.log_period
     metrics_filename = f"metrics_step_{start_step:06}_to_step_{step:06}.txt"
     with open(metrics_filename, "w", encoding="utf8") as metrics_for_gcs: