Convert a numpy image or a batch of images to a PIL image.

Convert a torch image to a PIL image.

Loads image to a PIL Image.

Prepares a single grid of images. Useful for visualization purposes.

A helper function to create random tensors on the desired device with the desired dtype. When passing a list of generators, you can seed each batch size individually. If CPU generators are passed, the tensor is always created on the CPU.

Applies layerwise casting to a given module. The module expected here is a Diffusers ModelMixin but it can be any nn.Module using diffusers layers or pytorch primitives.

>>> import torch
>>> from diffusers import CogVideoXTransformer3DModel

>>> transformer = CogVideoXTransformer3DModel.from_pretrained(
...     model_id, subfolder="transformer", torch_dtype=torch.bfloat16
... )

>>> apply_layerwise_casting(
...     transformer,
...     storage_dtype=torch.float8_e4m3fn,
...     compute_dtype=torch.bfloat16,
...     skip_modules_pattern=["patch_embed", "norm", "proj_out"],
...     non_blocking=True,
... )

Applies group offloading to the internal layers of a torch.nn.Module. To understand what group offloading is, and where it is beneficial, we need to first provide some context on how other supported offloading methods work.

Typically, offloading is done at two levels:

• Module-level: In Diffusers, this can be enabled using the ModelMixin::enable_model_cpu_offload() method. It works by offloading each component of a pipeline to the CPU for storage, and onloading to the accelerator device when needed for computation. This method is more memory-efficient than keeping all components on the accelerator, but the memory requirements are still quite high. For this method to work, one needs memory equivalent to size of the model in runtime dtype + size of largest intermediate activation tensors to be able to complete the forward pass.
• Leaf-level: In Diffusers, this can be enabled using the ModelMixin::enable_sequential_cpu_offload() method. It works by offloading the lowest leaf-level parameters of the computation graph to the CPU for storage, and onloading only the leafs to the accelerator device for computation. This uses the lowest amount of accelerator memory, but can be slower due to the excessive number of device synchronizations.

Group offloading is a middle ground between the two methods. It works by offloading groups of internal layers, (either torch.nn.ModuleList or torch.nn.Sequential). This method uses lower memory than module-level offloading. It is also faster than leaf-level/sequential offloading, as the number of device synchronizations is reduced.

Another supported feature (for CUDA devices with support for asynchronous data transfer streams) is the ability to overlap data transfer and computation to reduce the overall execution time compared to sequential offloading. This is enabled using layer prefetching with streams, i.e., the layer that is to be executed next starts onloading to the accelerator device while the current layer is being executed - this increases the memory requirements slightly. Note that this implementation also supports leaf-level offloading but can be made much faster when using streams.

>>> from diffusers import CogVideoXTransformer3DModel
>>> from diffusers.hooks import apply_group_offloading

>>> transformer = CogVideoXTransformer3DModel.from_pretrained(
...     "THUDM/CogVideoX-5b", subfolder="transformer", torch_dtype=torch.bfloat16
... )

>>> apply_group_offloading(
...     transformer,
...     onload_device=torch.device("cuda"),
...     offload_device=torch.device("cpu"),
...     offload_type="block_level",
...     num_blocks_per_group=2,
...     use_stream=True,
... )