[Flax] Add Textual Inversion (#880)

* add textual inversion flax

* make style

* make style

* replicate vae and unet params

* make style

* minor

* save after end of training

* style

* Temporary fix
Co-authored-by: Suraj Patil <surajp815@gmail.com>

* Add Flax instruction
Co-authored-by: Suraj Patil <surajp815@gmail.com>

examples/textual_inversion/README.md

@@ -68,6 +68,24 @@ accelerate launch textual_inversion.py \
 
 A full training run takes ~1 hour on one V100 GPU.
 
+If you want to speed it up even more, Flax implementation is available:
+
+```bash
+export MODEL_NAME="duongna/stable-diffusion-v1-4-flax"
+export DATA_DIR="path-to-dir-containing-images"
+
+python textual_inversion_flax.py \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATA_DIR \
+  --learnable_property="object" \
+  --placeholder_token="<cat-toy>" --initializer_token="toy" \
+  --resolution=512 \
+  --train_batch_size=1 \
+  --max_train_steps=3000 \
+  --learning_rate=5.0e-04 --scale_lr \
+  --output_dir="textual_inversion_cat"
+```
+It should be at least 70% faster than the PyTorch script with the same configuration.
 
 ### Inference