实用工具

用于处理 🤗 Diffusers 的实用函数和辅助函数。

numpy_to_pil

diffusers.utils.numpy_to_pil

< source >

( images )

将 numpy 图像或一批图像转换为 PIL 图像。

pt_to_pil

diffusers.utils.pt_to_pil

< source >

( images )

将 torch 图像转换为 PIL 图像。

load_image

diffusers.utils.load_image

< source >

( image: str | PIL.Image.Imageconvert_method: typing.Optional[typing.Callable[[PIL.Image.Image], PIL.Image.Image]] = None ) → PIL.Image.Image

参数

image (str 或 PIL.Image.Image) — 要转换为 PIL Image 格式的图像。
convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], 可选) — 加载图像后应用于该图像的转换方法。当设置为 None 时，图像将被转换为 “RGB”。

返回值

PIL.Image.Image

PIL 图像。

将 image 加载为 PIL 图像。

load_video

diffusers.utils.load_video

< source >

( video: strconvert_method: typing.Optional[typing.Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]]] = Nonereturn_fps: bool = False ) → list[PIL.Image.Image] or tuple[list[PIL.Image.Image], float]

参数

video (str) — 视频 URL 或路径，用于将其转换为 PIL Image 格式的图像列表。
convert_method (Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]], optional) — 加载视频后应用于该视频的转换方法。当设置为 None 时，图像将被转换为“RGB”。
return_fps (bool，可选，默认为 False) — 是否同时返回视频编码时的帧率。需要按模型工作帧率对输入进行重采样的流水线需要此信息，因为帧列表本身不包含该信息。

返回值

list[PIL.Image.Image] 或 tuple[list[PIL.Image.Image], float]

作为 PIL 图像列表的视频，以及当设置了 return_fps 时的帧率。

将 video 加载为 PIL Image 列表。

export_to_gif

diffusers.utils.export_to_gif

< source >

( image: list[PIL.Image.Image]output_gif_path: str = Nonefps: int = 10 )

export_to_video

diffusers.utils.export_to_video

< source >

( video_frames: list[np.ndarray] | list[PIL.Image.Image]output_video_path: str = Nonefps: int = 10quality: float = 5.0bitrate: int | None = Nonemacro_block_size: int | None = 16 )

quality: 视频输出质量。默认值为 5。使用可变比特率。最高质量为 10，最低为 0。设置为 None 可防止向 FFMPEG 传递可变比特率标志，以便您改用 output_params 手动指定。使用 bitrate 指定固定比特率会禁用此参数。

bitrate: 为视频编码设置恒定码率。默认值为 None，此时将改用 quality 参数。使用 quality 可变码率参数，而非通过此参数指定固定码率，可以获得质量更好且文件更小的视频。

macro_block_size: 视频的块大小约束。宽度和高度必须能被该数值整除。如果不能整除，imageio 会指示 ffmpeg 将图像放大到最接近的、能被该数值整除的尺寸。大多数编解码器兼容 16（默认值）的宏块大小，部分编解码器支持更小的值（4、8）。若要禁用此自动功能，请将其设置为 None 或 1，但请注意，许多播放器无法解码尺寸为奇数的视频，且某些编解码器可能会产生较差的结果或失败。参见 https://en.wikipedia.org/wiki/Macroblock。

encode_video

diffusers.utils.encode_video

< source >

参数

video (List[PIL.Image.Image] 或 np.ndarray 或 torch.Tensor) — 形状为 [frames, height, width, channels] 的视频张量，像素值为 [0, 255] 范围内的整数。如果输入是 np.ndarray，则期望它是一个值在 [0, 1] 范围内的浮点数组（这通常是管道在使用 output_type="np" 时返回的结果）。
fps (int) — 编码视频的每秒帧数（FPS）。
output_path (str) — 用于保存编码后视频的路径。
audio (torch.Tensor, 可选) — 形状为 [audio_channels, samples] 的音频波形。
audio_sample_rate — (int，可选)：音频波形的采样率。
video_chunks_number (int，可选，默认为 1) — 用于编码时将视频拆分的块数。每个块将单独进行编码。使用的块数通常取决于视频 VAE 的平铺配置。

使用 PyAV 库对视频进行编码，可选包含音频。基于原始 LTX-2 仓库中的代码：https://github.com/Lightricks/LTX-2/blob/4f410820b198e05074a1e92de793e3b59e9ab5a0/packages/ltx-pipelines/src/ltx_pipelines/utils/media_io.py#L182

make_image_grid

diffusers.utils.make_image_grid

< source >

( images: list[PIL.Image.Image]rows: intcols: intresize: int = None )

准备单个图像网格。适用于可视化目的。

randn_tensor

diffusers.utils.torch_utils.randn_tensor

< source >

一个辅助函数，用于在指定的 device 上创建具有指定 dtype 的随机张量。当传入生成器列表时，可以为每个批次大小单独设置种子。如果传入的是 CPU 生成器，则张量始终在 CPU 上创建。

TorchDeviceBackend

class diffusers.utils.torch_utils.TorchDeviceBackend

< source >

( device: str | torch.device | None = None )

torch.<backend> 命名空间（torch.cuda、torch.xpu、torch.mps 等）中某个设备的代理。该类未定义的属性即为模块自身的属性（synchronize、device_count、Stream、current_stream 等）；此处定义的方法会覆盖那些在不同后端之间可用性存在差异、需要回退机制的操作：缓存清除、随机数种子设置以及内存查询。若未指定 device，则通过 torch.accelerator 检测主机加速器。如果 torch 没有对应后端的模块，则会抛出异常，而不是静默回退到 torch.cuda。

max_memory_allocated

< source >

( )

自上次重置以来，设备上分配的峰值内存（以字节为单位）；如果后端不维护统计信息，则值为 0。

mem_get_info

< source >

( )

以字节为单位的空闲和总设备内存。

apply_layerwise_casting

diffusers.hooks.apply_layerwise_casting

< source >

( module: Modulestorage_dtype: dtypecompute_dtype: dtypeskip_modules_pattern: str | tuple[str, ...] = 'auto'skip_modules_classes: tuple[typing.Type[torch.nn.Module], ...] | None = Nonenon_blocking: bool = False )

参数

module (torch.nn.Module) — 该模块的叶子模块将在计算时转换为高精度 dtype，在存储时转换为低精度 dtype。
storage_dtype (torch.dtype) — 用于在 forward pass 前后将模块转换以进行存储的 dtype。
compute_dtype (torch.dtype) — 在前向传播计算过程中，将模块转换到的数据类型。
skip_modules_pattern (tuple[str, ...]，默认为 "auto") — 用于匹配在逐层转换过程中需要跳过的模块名称的模式列表。如果设置为 "auto"，则使用默认模式。如果设置为 None，则不跳过任何模块。如果设置为 None 且 skip_modules_classes 也为 None，则逐层转换将直接应用于该模块，而不是其内部子模块。
skip_modules_classes (tuple[Type[torch.nn.Module], ...]，默认为 None) — 在逐层转换过程中需要跳过的模块类列表。
non_blocking (bool，默认为 False) — 如果为 True，则权重转换操作为非阻塞的。

对给定模块应用逐层类型转换。此处期望的模块是 Diffusers 的 ModelMixin，但也可以是任何使用 diffusers 层或 PyTorch 原语的 nn.Module。

示例：

>>> 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,
... )

apply_group_offloading

diffusers.hooks.apply_group_offloading

< source >

( module: Moduleonload_device: typing.Union[str, torch.device]offload_device: typing.Union[str, torch.device] = torch.device()offload_type: str | diffusers.hooks.group_offloading.GroupOffloadingType = 'block_level'num_blocks_per_group: int | None = Nonenon_blocking: bool = Falseuse_stream: bool = Falserecord_stream: bool = Falselow_cpu_mem_usage: bool = Falseoffload_to_disk_path: str | None = Noneblock_modules: list[str] | None = Noneexclude_kwargs: list[str] | None = None )

参数

module (torch.nn.Module) — 应用组卸载的模块。
onload_device (torch.device) — 模块组被加载到的设备。
offload_device (torch.device，默认为 torch.device("cpu")) — 模块组被卸载到的设备。通常应为 CPU。默认值为 CPU。
offload_type (str 或 GroupOffloadingType，默认为 “block_level”) — 要应用的卸载类型。可以是 “block_level” 或 “leaf_level” 之一。默认为 “block_level”。
offload_to_disk_path (str，可选，默认为 None) — 参数将被卸载到的目录路径。在内存受限的环境中，若希望获得合理的速度与内存之间的权衡，设置此选项会很有用。
num_blocks_per_group (int, 可选) — 当使用 offload_type=“block_level” 时，每个组中的块数。使用 offload_type=“block_level” 时，此参数为必填项。
non_blocking (bool，默认为 False) — 如果为 True，则使用非阻塞数据传输进行卸载和加载。
use_stream (bool，默认为 False) — 如果为 True，则使用 CUDA 流异步执行卸载和加载操作。这有助于重叠计算和数据传输。
record_stream (bool，默认为 False) — 当与 use_stream 一起启用时，它会将当前张量标记为已被此流使用。这种方式速度更快，但会略微增加内存占用。有关更多详细信息，请参阅 PyTorch 官方文档。
low_cpu_mem_usage (bool，默认为 False) — 如果为 True，将通过即时固定张量（而非预先固定）来最小化 CPU 内存使用。此选项仅在使用流式 CPU 卸载（即 use_stream=True）时才有意义。当 CPU 内存成为瓶颈时，此选项可能有用，但可能会抵消使用流所带来的好处。
block_modules (list[str], 可选) — 应作为 offloading 块处理的模块名称列表。如果提供，则仅这些模块会被考虑用于块级 offloading。如果未提供，将使用默认的块检测逻辑。
exclude_kwargs (list[str], 可选) — 不应由 send_to_device 处理的 kwarg 键列表。这对于需要在前向传播之间保持对象身份的可变状态（例如缓存列表）非常有用。如果未提供，则会在模块存在 _skip_keys 属性时从该属性推断。

将组卸载（group offloading）应用于 torch.nn.Module 的内部层。要了解什么是组卸载以及它在哪些场景下具有优势，我们首先需要介绍其他受支持的卸载方法的工作原理。

通常，卸载操作在两个层面进行：

模块级：在 Diffusers 中，可以使用 ModelMixin::enable_model_cpu_offload() 方法启用此功能。其工作原理是将流水线中的每个组件卸载到 CPU 上进行存储，并在需要计算时加载到加速器设备上。与将所有组件都保留在加速器上相比，此方法在内存效率上更优，但内存需求仍然相当高。要使此方法正常工作，所需的内存量相当于运行时数据类型下模型的大小加上最大的中间激活张量的大小，以便能够完成前向传播。
叶节点级别：在 Diffusers 中，可以通过 ModelMixin::enable_sequential_cpu_offload() 方法启用此功能。其工作原理是将计算图中最底层的叶节点参数卸载到 CPU 上进行存储，并在计算时仅将这些叶节点加载到加速器设备上。这种方式占用的加速器内存最少，但由于设备同步次数过多，速度可能会较慢。

Group offloading 是上述两种方法之间的折中方案。它通过卸载内部层的分组（即 torch.nn.ModuleList 或 torch.nn.Sequential）来实现。与模块级卸载相比，该方法占用的内存更少。同时，由于减少了设备同步的次数，其速度也快于叶子级/顺序卸载。

另一种受支持的功能（适用于支持异步数据传输流的 CUDA 设备）是能够重叠数据传输与计算，从而与顺序卸载相比缩短整体执行时间。这通过带流的层预取来实现，即在执行当前层的同时，下一层开始向加速器设备加载——这会略微增加内存需求。请注意，此实现也支持叶级卸载，但使用流时速度会快得多。

示例：

>>> 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,
... )

Update on GitHub