[Flax] Add DreamBooth #1001

duongna21 · 2022-10-26T17:27:47Z

What does this PR do?

Add Flax example for DreamBooth.

How to run (74% faster than PyTorch example with same args on Tesla A100)

export MODEL_NAME="duongna/stable-diffusion-v1-4-flax"
export INSTANCE_DIR="path-to-instance-images"
export CLASS_DIR="path-to-class-images"
export OUTPUT_DIR="path-to-save-model"

python train_dreambooth_flax.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --class_data_dir=$CLASS_DIR \
  --output_dir=$OUTPUT_DIR \
  --with_prior_preservation --prior_loss_weight=1.0 \
  --instance_prompt="a photo of sks dog" \
  --class_prompt="a photo of dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --learning_rate=5e-6 \
  --num_class_images=200 \
  --max_train_steps=800

Prompt: a photo of sks dog

Who can review?

cc @patil-suraj @patrickvonplaten

patil-suraj · 2022-10-26T17:29:00Z

You are on fire @duongna21 ! 🔥

HuggingFaceDocBuilderDev · 2022-10-26T17:35:04Z

The documentation is not available anymore as the PR was closed or merged.

patil-suraj

Looks very good, amazing work! Just left some comments.

Let's make sure that seed is not None as PRNGKey will break. Also let's update the readme to show hot to run this example. Then this should be good to merge :)

examples/dreambooth/train_dreambooth_flax.py

patil-suraj · 2022-10-27T09:04:16Z

examples/dreambooth/train_dreambooth_flax.py

+            for example in tqdm(
+                sample_dataloader, desc="Generating class images", disable=not jax.process_index() == 0
+            ):
+                prompt_ids = pipeline.prepare_inputs(example["prompt"])
+                prompt_ids = shard(prompt_ids)
+                p_params = jax_utils.replicate(params)
+                rng = jax.random.split(rng)[0]
+                sample_rng = jax.random.split(rng, jax.device_count())
+                images = pipeline(prompt_ids, p_params, sample_rng, jit=True).images
+                images = images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:])
+                images = pipeline.numpy_to_pil(np.array(images))
+
+                for i, image in enumerate(images):
+                    hash_image = hashlib.sha1(image.tobytes()).hexdigest()
+                    image_filename = class_images_dir / f"{example['index'][i] + cur_class_images}-{hash_image}.jpg"
+                    image.save(image_filename)


examples/dreambooth/train_dreambooth_flax.py

patil-suraj · 2022-10-27T09:24:14Z

(For future PR, maybe also enable the option to allow training the text_encoder this has been found to improve results significantly)

duongna21 · 2022-10-27T11:46:43Z

(For future PR, maybe also enable the option to allow training the text_encoder this has been found to improve results significantly)

@patil-suraj Actually train_text_encoder has been allowed in this PR. Please check it out.
Also thank you for other very helpful comments. Addressed them!

patil-suraj

Thanks a lot @duongna21 !

Also, it seems like text_encoder is always trained , no ? I think we should add an option called --train_text_encoder and train it only when its True. Because training text_encoder is not always needed.

patil-suraj · 2022-10-27T12:03:24Z

examples/dreambooth/train_dreambooth_flax.py

+    weight_dtype = jnp.float32
+    if args.mixed_precision == "fp16":
+        weight_dtype = jnp.float16
+    elif args.mixed_precision == "bf16":
+        weight_dtype = jnp.bfloat16


patil-suraj · 2022-10-27T12:03:43Z

examples/dreambooth/README.md

+Or use the Flax implementation if you need a speedup
+
+```bash
+export MODEL_NAME="duongna/stable-diffusion-v1-4-flax"
+export INSTANCE_DIR="path-to-instance-images"
+export OUTPUT_DIR="path-to-save-model"
+
+python train_dreambooth_flax.py \
+  --pretrained_model_name_or_path=$MODEL_NAME  \
+  --instance_data_dir=$INSTANCE_DIR \
+  --output_dir=$OUTPUT_DIR \
+  --instance_prompt="a photo of sks dog" \
+  --resolution=512 \
+  --train_batch_size=1 \
+  --learning_rate=5e-6 \


Thanks for adding this!

duongna21 · 2022-10-27T12:14:46Z

Also, it seems like text_encoder is always trained?

@patil-suraj No. You can ctrl + F for if args.train_text_encoder in the script and check if anything is wrong.

patil-suraj · 2022-10-27T12:23:52Z

Ahh, sorry. I missed that, all looks good now, thanks a lot for working on this!

patrickvonplaten · 2022-10-29T07:44:55Z

Great work here!

skirsten · 2022-11-02T23:10:12Z

examples/dreambooth/train_dreambooth_flax.py

+        weight_dtype = jnp.bfloat16
+
+    # Load models and create wrapper for stable diffusion
+    text_encoder = FlaxCLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", dtype=weight_dtype)


Hi, why is this version pulled from the hub and its not just using the one in the text_encoder subfolder in args.pretrained_model_name_or_path?

Is there something wrong with the other version?

@skirsten Nice question. Look at this please.

@patrickvonplaten is the PR in transformers merged, that allows loading Flax clip with subfolder ?

@douwekiela Awesome! Thanks for explaining it

@patil-suraj yes, it's merged

Awesome, @duongna21 would you like to open a PR and to update this then :) We will also need to update the installation instructions for transformers to include that fix

duongna21 added 2 commits October 26, 2022 22:49

[Flax] Add DreamBooth

9f1a925

fix sample rng

de82e59

patil-suraj self-assigned this Oct 26, 2022

style

ff584de

not reuse rng

ebb3e28

patil-suraj approved these changes Oct 27, 2022

View reviewed changes

duongna21 and others added 2 commits October 27, 2022 18:32

add dtype for mixed precision training

c9e1ff8

Add Flax example

f34bbfd

patil-suraj reviewed Oct 27, 2022

View reviewed changes

patil-suraj approved these changes Oct 27, 2022

View reviewed changes

patil-suraj merged commit 90f91ad into huggingface:main Oct 27, 2022

duongna21 deleted the add-dreambooth-flax branch October 27, 2022 14:19

patrickvonplaten mentioned this pull request Nov 2, 2022

train_dreambooth_flax.py 😭 #1088

Closed

skirsten reviewed Nov 2, 2022

View reviewed changes

duongna21 mentioned this pull request Nov 5, 2022

[Flax examples] Load text encoder from subfolder #1147

Merged

PhaneeshB pushed a commit to nod-ai/diffusers that referenced this pull request Mar 1, 2023

Fix dark theme again for exe builds (huggingface#1001)

f64e1fb

[Flax] Add DreamBooth #1001

[Flax] Add DreamBooth #1001

Uh oh!

Conversation

duongna21 commented Oct 26, 2022

What does this PR do?

Who can review?

Uh oh!

patil-suraj commented Oct 26, 2022

Uh oh!

HuggingFaceDocBuilderDev commented Oct 26, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

patil-suraj left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

patil-suraj commented Oct 27, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

duongna21 commented Oct 27, 2022

Uh oh!

patil-suraj left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

duongna21 commented Oct 27, 2022

Uh oh!

patil-suraj commented Oct 27, 2022

Uh oh!

patrickvonplaten commented Oct 29, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

duongna21 Nov 3, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Oct 26, 2022 •

edited

Loading

patil-suraj commented Oct 27, 2022 •

edited

Loading

duongna21 Nov 3, 2022 •

edited

Loading