local variable 'train_dataset' referenced before assignment

[RissyRan]

Describe the bug

Hi team,

Recently, we are running Diffusion model on JAX, and noticed an UnboundLocalError issue for local variable train_dataset. It could be a simple fix as this forked repo.

Reproduction

Codes to reproduce:

git clone https://github.com/huggingface/diffusers.git && cd diffusers && pip install . && pip install tensorflow clu && pip install -U -r examples/text_to_image/requirements_flax.txt

cd diffusers/examples/text_to_image && TPU_LIBRARY_PATH=/lib/libtpu.so TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --fake_data=1 --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model

Logs

UnboundLocalError: local variable 'train_dataset' referenced before assignment

System Info

Cloud TPU (Ref)

[sayakpaul]

@yiyixuxu could you look into this?

[yiyixuxu]

@RissyRan are you running on more than one node?

[RissyRan]

@yiyixuxu Yes, that's right. We are running test on Cloud TPU pod, and it includes multiple nodes/workers.

[RissyRan]

Hi team, any updates? Thanks!

[yiyixuxu]

@RissyRan

I think removing the jax.process_index() == 0 will probably resolve the unbound variable error, but I don't think it's sufficient to make our script work correctly with multi-tpu training, i.e.data won't be batched correctly (cc @patrickvonplaten for his opinion here since I haven't had any experience working with pods)

It's on my to-do to make it work with TPU pods, but it's not a top priority right now since most of our users only have access to a single node. I will work on it soon and let you know once it's ready for testing.

[patrickvonplaten]

Maybe also cc @andsteing

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.