Repository: huggingface/diffusers

Issue:
[SD-XL] Instruct Pix2Pix fine-tuning script unconditionally upcasts VAE to float32 when built-in VAE is used
```py
    # Move vae, unet and text_encoder to device and cast to weight_dtype
    # The VAE is in float32 to avoid NaN losses.
    if args.pretrained_vae_model_name_or_path is not None:
        vae.to(accelerator.device, dtype=weight_dtype)
    else:
        vae.to(accelerator.device, dtype=torch.float32)
```

I believe this causes the training to fail when operating in mixed precision, eg. `fp16` or `bf16`:

```
RuntimeError: Input type (c10::BFloat16) and bias type (float) should be the same
```


Discussion:
It's not unconditional. 

When you supply a VAE that has better numerical behavior (such as: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix), we don't do the upcast. Otherwise, we upcast. 

Could you help me try to reproduce this issue?
yes, use ptx0/s1 model that has the better vae included
That is model specific then. In my experiments, the upcasting was needed for the worse VAE (with fp16). 

Feel free to reopen if the issue is persistent. 
yes, but the model's scheduler config has it disabled. the training script needs to take this into consideration.

in fact, we can just pass the value through. the problem is that the VAE is being initialized without its config.
Possible to maybe fix this in a PR?
reopen/assign this one on me and i'll get to it :D 
Done :-)
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](https://github.com/huggingface/diffusers/blob/main/CONTRIBUTING.md) are likely to be ignored.
@github-actions bad bot

Resolving patch:
diff --git a/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py b/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py
--- a/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py
+++ b/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py
@@ -63,6 +63,7 @@
     "fusing/instructpix2pix-1000-samples": ("file_name", "edited_image", "edit_prompt"),
 }
 WANDB_TABLE_COL_NAMES = ["file_name", "edited_image", "edit_prompt"]
+TORCH_DTYPE_MAPPING = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}
 
 
 def import_model_class_from_model_name_or_path(
@@ -100,6 +101,16 @@ def parse_args():
         default=None,
         help="Path to an improved VAE to stabilize training. For more details check out: https://github.com/huggingface/diffusers/pull/4038.",
     )
+    parser.add_argument(
+        "--vae_precision",
+        type="choice",
+        choices=["fp32", "fp16", "bf16"],
+        default="fp32",
+        help=(
+            "The vanilla SDXL 1.0 VAE can cause NaNs due to large activation values. Some custom models might already have a solution"
+            " to this problem, and this flag allows you to use mixed precision to stabilize training."
+        ),
+    )
     parser.add_argument(
         "--revision",
         type=str,
@@ -878,7 +889,7 @@ def collate_fn(examples):
     if args.pretrained_vae_model_name_or_path is not None:
         vae.to(accelerator.device, dtype=weight_dtype)
     else:
-        vae.to(accelerator.device, dtype=torch.float32)
+        vae.to(accelerator.device, dtype=TORCH_DTYPE_MAPPING[args.vae_precision])
 
     # We need to recalculate our total training steps as the size of the training dataloader may have changed.
     num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)


Test patch: