[training] add an offload utility that can be used as a context manager.

[sayakpaul]

What does this PR do?

Useful for handling offloading in a simple and easy manner. Currently, in some of our training scripts, we have code patterns like so:

if args.offload:
     vae.to(accelerator.device)

# perform computation
....

if args.offload:
    vae.cpu()

Instead of having those conditionals, this PR introduces an offload_models() method that turns this into a simple block of code:

with offload_models(vae, device=accelerator.device, offload=args.offload):
    pixel_values = batch["pixel_values"].to(dtype=vae.dtype)

@DN6 I am including changes to examples/**/*.py to also trigger the fast PR tests on GPU. Hope, that's okay?

Cc: @a-r-r-o-w for visibility.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[linoytsaban]

nice @sayakpaul, thanks! should we also add a test to [test_dreambooth_lora_hidream.py]?(https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/test_dreambooth_lora_hidream.py)

[sayakpaul]

should we also add a test to [test_dreambooth_lora_hidream.py]?

I think it should be added as a utility test under tests/others. Currently, it's a bit non-trivial to test in terms of how the testing suite is designed for the scripts.

[sayakpaul]

@DN6 okay regarding?

I am including changes to examples/**/*.py to also trigger the fast PR tests on GPU. Hope, that's okay?

[sayakpaul]

@linoytsaban two things:

• Would you be able to propagate / introduce the offloading related changes to the other important training scripts?
• In https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora_flux_kontext.py, we introduced the ability to precompute text embeddings even when each image has a custom prompt. We could propagate that to the rest of the Flux training scripts?